1#pragma once
2
3#include <c10/macros/Macros.h>
4#include <c10/util/Exception.h>
5#include <c10/util/in_place.h>
6
7#include <type_traits>
8
9namespace c10 {
10
11/// MaybeOwnedTraits<T> describes how to borrow from T. Here is how we
12/// can implement borrowing from an arbitrary type T using a raw
13/// pointer to const:
14template <typename T>
15struct MaybeOwnedTraitsGenericImpl {
16 using owned_type = T;
17 using borrow_type = const T*;
18
19 static borrow_type createBorrow(const owned_type& from) {
20 return &from;
21 }
22
23 static void assignBorrow(borrow_type& lhs, borrow_type rhs) {
24 lhs = rhs;
25 }
26
27 static void destroyBorrow(borrow_type& /*toDestroy*/) {}
28
29 static const owned_type& referenceFromBorrow(const borrow_type& borrow) {
30 return *borrow;
31 }
32
33 static const owned_type* pointerFromBorrow(const borrow_type& borrow) {
34 return borrow;
35 }
36
37 static bool debugBorrowIsValid(const borrow_type& borrow) {
38 return borrow != nullptr;
39 }
40};
41
42/// It is possible to eliminate the extra layer of indirection for
43/// borrows for some types that we control. For examples, see
44/// intrusive_ptr.h and TensorBody.h.
45
46template <typename T>
47struct MaybeOwnedTraits;
48
49// Explicitly enable MaybeOwned<shared_ptr<T>>, rather than allowing
50// MaybeOwned to be used for any type right away.
51template <typename T>
52struct MaybeOwnedTraits<std::shared_ptr<T>>
53 : public MaybeOwnedTraitsGenericImpl<std::shared_ptr<T>> {};
54
55/// A smart pointer around either a borrowed or owned T. When
56/// constructed with borrowed(), the caller MUST ensure that the
57/// borrowed-from argument outlives this MaybeOwned<T>. Compare to
58/// Rust's std::borrow::Cow
59/// (https://doc.rust-lang.org/std/borrow/enum.Cow.html), but note
60/// that it is probably not suitable for general use because C++ has
61/// no borrow checking. Included here to support
62/// Tensor::expect_contiguous.
63template <typename T>
64class MaybeOwned final {
65 using borrow_type = typename MaybeOwnedTraits<T>::borrow_type;
66 using owned_type = typename MaybeOwnedTraits<T>::owned_type;
67
68 bool isBorrowed_;
69 union {
70 borrow_type borrow_;
71 owned_type own_;
72 };
73
74 /// Don't use this; use borrowed() instead.
75 explicit MaybeOwned(const owned_type& t)
76 : isBorrowed_(true), borrow_(MaybeOwnedTraits<T>::createBorrow(t)) {}
77
78 /// Don't use this; use owned() instead.
79 explicit MaybeOwned(T&& t) noexcept(
80 std::is_nothrow_move_constructible<T>::value)
81 : isBorrowed_(false), own_(std::move(t)) {}
82
83 /// Don't use this; use owned() instead.
84 template <class... Args>
85 explicit MaybeOwned(in_place_t, Args&&... args)
86 : isBorrowed_(false), own_(std::forward<Args>(args)...) {}
87
88 public:
89 explicit MaybeOwned() : isBorrowed_(true), borrow_() {}
90
91 // Copying a borrow yields another borrow of the original, as with a
92 // T*. Copying an owned T yields another owned T for safety: no
93 // chains of borrowing by default! (Note you could get that behavior
94 // with MaybeOwned<T>::borrowed(*rhs) if you wanted it.)
95 MaybeOwned(const MaybeOwned& rhs) : isBorrowed_(rhs.isBorrowed_) {
96 if (C10_LIKELY(rhs.isBorrowed_)) {
97 MaybeOwnedTraits<T>::assignBorrow(borrow_, rhs.borrow_);
98 } else {
99 new (&own_) T(rhs.own_);
100 }
101 }
102
103 MaybeOwned& operator=(const MaybeOwned& rhs) {
104 if (this == &rhs) {
105 return *this;
106 }
107 if (C10_UNLIKELY(!isBorrowed_)) {
108 if (rhs.isBorrowed_) {
109 own_.~T();
110 MaybeOwnedTraits<T>::assignBorrow(borrow_, rhs.borrow_);
111 isBorrowed_ = true;
112 } else {
113 own_ = rhs.own_;
114 }
115 } else {
116 if (C10_LIKELY(rhs.isBorrowed_)) {
117 MaybeOwnedTraits<T>::assignBorrow(borrow_, rhs.borrow_);
118 } else {
119 MaybeOwnedTraits<T>::destroyBorrow(borrow_);
120 new (&own_) T(rhs.own_);
121 isBorrowed_ = false;
122 }
123 }
124 TORCH_INTERNAL_ASSERT_DEBUG_ONLY(isBorrowed_ == rhs.isBorrowed_);
125 return *this;
126 }
127
128 MaybeOwned(MaybeOwned&& rhs) noexcept(
129 std::is_nothrow_move_constructible<T>::value)
130 : isBorrowed_(rhs.isBorrowed_) {
131 if (C10_LIKELY(rhs.isBorrowed_)) {
132 MaybeOwnedTraits<T>::assignBorrow(borrow_, rhs.borrow_);
133 } else {
134 new (&own_) T(std::move(rhs.own_));
135 }
136 }
137
138 MaybeOwned& operator=(MaybeOwned&& rhs) noexcept(
139 std::is_nothrow_move_assignable<T>::value) {
140 if (this == &rhs) {
141 return *this;
142 }
143 if (C10_UNLIKELY(!isBorrowed_)) {
144 if (rhs.isBorrowed_) {
145 own_.~T();
146 MaybeOwnedTraits<T>::assignBorrow(borrow_, rhs.borrow_);
147 isBorrowed_ = true;
148 } else {
149 own_ = std::move(rhs.own_);
150 }
151 } else {
152 if (C10_LIKELY(rhs.isBorrowed_)) {
153 MaybeOwnedTraits<T>::assignBorrow(borrow_, rhs.borrow_);
154 } else {
155 MaybeOwnedTraits<T>::destroyBorrow(borrow_);
156 new (&own_) T(std::move(rhs.own_));
157 isBorrowed_ = false;
158 }
159 }
160 TORCH_INTERNAL_ASSERT_DEBUG_ONLY(isBorrowed_ == rhs.isBorrowed_);
161 return *this;
162 }
163
164 static MaybeOwned borrowed(const T& t) {
165 return MaybeOwned(t);
166 }
167
168 static MaybeOwned owned(T&& t) noexcept(
169 std::is_nothrow_move_constructible<T>::value) {
170 return MaybeOwned(std::move(t));
171 }
172
173 template <class... Args>
174 static MaybeOwned owned(in_place_t, Args&&... args) {
175 return MaybeOwned(in_place, std::forward<Args>(args)...);
176 }
177
178 ~MaybeOwned() {
179 if (C10_UNLIKELY(!isBorrowed_)) {
180 own_.~T();
181 } else {
182 MaybeOwnedTraits<T>::destroyBorrow(borrow_);
183 }
184 }
185
186 // This is an implementation detail! You should know what you're doing
187 // if you are testing this. If you just want to guarantee ownership move
188 // this into a T
189 bool unsafeIsBorrowed() const {
190 return isBorrowed_;
191 }
192
193 const T& operator*() const& {
194 if (isBorrowed_) {
195 TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
196 MaybeOwnedTraits<T>::debugBorrowIsValid(borrow_));
197 }
198 return C10_LIKELY(isBorrowed_)
199 ? MaybeOwnedTraits<T>::referenceFromBorrow(borrow_)
200 : own_;
201 }
202
203 const T* operator->() const {
204 if (isBorrowed_) {
205 TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
206 MaybeOwnedTraits<T>::debugBorrowIsValid(borrow_));
207 }
208 return C10_LIKELY(isBorrowed_)
209 ? MaybeOwnedTraits<T>::pointerFromBorrow(borrow_)
210 : &own_;
211 }
212
213 // If borrowed, copy the underlying T. If owned, move from
214 // it. borrowed/owned state remains the same, and either we
215 // reference the same borrow as before or we are an owned moved-from
216 // T.
217 T operator*() && {
218 if (isBorrowed_) {
219 TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
220 MaybeOwnedTraits<T>::debugBorrowIsValid(borrow_));
221 return MaybeOwnedTraits<T>::referenceFromBorrow(borrow_);
222 } else {
223 return std::move(own_);
224 }
225 }
226};
227
228} // namespace c10
229