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 | |
9 | namespace 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: |
14 | template <typename T> |
15 | struct 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 | |
46 | template <typename T> |
47 | struct MaybeOwnedTraits; |
48 | |
49 | // Explicitly enable MaybeOwned<shared_ptr<T>>, rather than allowing |
50 | // MaybeOwned to be used for any type right away. |
51 | template <typename T> |
52 | struct 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. |
63 | template <typename T> |
64 | class 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 | |