1 /*
2  * Copyright (c) 2022 NVIDIA Corporation
3  *
4  * Licensed under the Apache License Version 2.0 with LLVM Exceptions
5  * (the "License"); you may not use this file except in compliance with
6  * the License. You may obtain a copy of the License at
7  *
8  *   https://llvm.org/LICENSE.txt
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 #pragma once
17 
18 #include "../concepts.hpp"
19 #include "__config.hpp"
20 #include "__meta.hpp"
21 
22 #include <atomic>
23 #include <cstddef>
24 #include <memory>
25 #include <new>
26 #include <type_traits>
27 
28 #if STDEXEC_TSAN()
29 #include <sanitizer/tsan_interface.h>
30 #endif
31 
32 namespace stdexec
33 {
34 namespace __ptr
35 {
36 template <class _Ty>
37 struct __make_intrusive_t;
38 
39 template <class _Ty>
40 class __intrusive_ptr;
41 
42 template <class _Ty>
43 struct __enable_intrusive_from_this
44 {
45     __intrusive_ptr<_Ty> __intrusive_from_this() noexcept;
46     __intrusive_ptr<const _Ty> __intrusive_from_this() const noexcept;
47 
48   private:
49     friend _Ty;
50     void __inc_ref() noexcept;
51     void __dec_ref() noexcept;
52 };
53 
54 STDEXEC_PRAGMA_PUSH()
55 STDEXEC_PRAGMA_IGNORE_GNU("-Wtsan")
56 
57 template <class _Ty>
58 struct __control_block
59 {
60     alignas(_Ty) unsigned char __value_[sizeof(_Ty)];
61     std::atomic<unsigned long> __refcount_;
62 
63     template <class... _Us>
64     explicit __control_block(_Us&&... __us) noexcept(noexcept(_Ty{
65         __declval<_Us>()...})) :
66         __refcount_(1u)
67     {
68         // Construct the value *after* the initialization of the
69         // atomic in case the constructor of _Ty calls
70         // __intrusive_from_this() (which increments the atomic):
71         ::new ((void*)__value_) _Ty{(_Us&&)__us...};
72     }
73 
74     ~__control_block()
75     {
76         __value().~_Ty();
77     }
78 
79     _Ty& __value() const noexcept
80     {
81         return *(_Ty*)__value_;
82     }
83 
84     void __inc_ref_() noexcept
85     {
86         __refcount_.fetch_add(1, std::memory_order_relaxed);
87     }
88 
89     void __dec_ref_() noexcept
90     {
91         if (1u == __refcount_.fetch_sub(1, std::memory_order_release))
92         {
93             std::atomic_thread_fence(std::memory_order_acquire);
94             // TSan does not support std::atomic_thread_fence, so we
95             // need to use the TSan-specific __tsan_acquire instead:
96             STDEXEC_TSAN(__tsan_acquire(&__refcount_));
97             delete this;
98         }
99     }
100 };
101 
102 STDEXEC_PRAGMA_POP()
103 
104 template <class _Ty>
105 class __intrusive_ptr
106 {
107     using _UncvTy = std::remove_cv_t<_Ty>;
108     friend struct __make_intrusive_t<_Ty>;
109     friend struct __enable_intrusive_from_this<_UncvTy>;
110 
111     __control_block<_UncvTy>* __data_{nullptr};
112 
113     explicit __intrusive_ptr(__control_block<_UncvTy>* __data) noexcept :
114         __data_(__data)
115     {}
116 
117     void __inc_ref_() noexcept
118     {
119         if (__data_)
120         {
121             __data_->__inc_ref_();
122         }
123     }
124 
125     void __dec_ref_() noexcept
126     {
127         if (__data_)
128         {
129             __data_->__dec_ref_();
130         }
131     }
132 
133   public:
134     using element_type = _Ty;
135 
136     __intrusive_ptr() = default;
137 
138     __intrusive_ptr(__intrusive_ptr&& __that) noexcept :
139         __data_(std::exchange(__that.__data_, nullptr))
140     {}
141 
142     __intrusive_ptr(const __intrusive_ptr& __that) noexcept :
143         __data_(__that.__data_)
144     {
145         __inc_ref_();
146     }
147 
148     __intrusive_ptr(__enable_intrusive_from_this<_Ty>* __that) noexcept :
149         __intrusive_ptr(__that ? __that->__intrusive_from_this()
150                                : __intrusive_ptr())
151     {}
152 
153     __intrusive_ptr& operator=(__intrusive_ptr&& __that) noexcept
154     {
155         [[maybe_unused]] __intrusive_ptr __old{
156             std::exchange(__data_, std::exchange(__that.__data_, nullptr))};
157         return *this;
158     }
159 
160     __intrusive_ptr& operator=(const __intrusive_ptr& __that) noexcept
161     {
162         return operator=(__intrusive_ptr(__that));
163     }
164 
165     __intrusive_ptr&
166         operator=(__enable_intrusive_from_this<_Ty>* __that) noexcept
167     {
168         return operator=(__that ? __that->__intrusive_from_this()
169                                 : __intrusive_ptr());
170     }
171 
172     ~__intrusive_ptr()
173     {
174         __dec_ref_();
175     }
176 
177     void reset() noexcept
178     {
179         operator=({});
180     }
181 
182     void swap(__intrusive_ptr& __that) noexcept
183     {
184         std::swap(__data_, __that.__data_);
185     }
186 
187     _Ty* get() const noexcept
188     {
189         return &__data_->__value();
190     }
191 
192     _Ty* operator->() const noexcept
193     {
194         return &__data_->__value();
195     }
196 
197     _Ty& operator*() const noexcept
198     {
199         return __data_->__value();
200     }
201 
202     explicit operator bool() const noexcept
203     {
204         return __data_ != nullptr;
205     }
206 
207     bool operator!() const noexcept
208     {
209         return __data_ == nullptr;
210     }
211 
212     bool operator==(const __intrusive_ptr&) const = default;
213 
214     bool operator==(std::nullptr_t) const noexcept
215     {
216         return __data_ == nullptr;
217     }
218 };
219 
220 template <class _Ty>
221 __intrusive_ptr<_Ty>
222     __enable_intrusive_from_this<_Ty>::__intrusive_from_this() noexcept
223 {
224     auto* __data = (__control_block<_Ty>*)static_cast<_Ty*>(this);
225     __data->__inc_ref_();
226     return __intrusive_ptr<_Ty>{__data};
227 }
228 
229 template <class _Ty>
230 __intrusive_ptr<const _Ty>
231     __enable_intrusive_from_this<_Ty>::__intrusive_from_this() const noexcept
232 {
233     auto* __data = (__control_block<_Ty>*)static_cast<const _Ty*>(this);
234     __data->__inc_ref_();
235     return __intrusive_ptr<const _Ty>{__data};
236 }
237 
238 template <class _Ty>
239 void __enable_intrusive_from_this<_Ty>::__inc_ref() noexcept
240 {
241     auto* __data = (__control_block<_Ty>*)static_cast<_Ty*>(this);
242     __data->__inc_ref_();
243 }
244 
245 template <class _Ty>
246 void __enable_intrusive_from_this<_Ty>::__dec_ref() noexcept
247 {
248     auto* __data = (__control_block<_Ty>*)static_cast<_Ty*>(this);
249     __data->__dec_ref_();
250 }
251 
252 template <class _Ty>
253 struct __make_intrusive_t
254 {
255     template <class... _Us>
256         requires constructible_from<_Ty, _Us...>
257     __intrusive_ptr<_Ty> operator()(_Us&&... __us) const
258     {
259         using _UncvTy = std::remove_cv_t<_Ty>;
260         return __intrusive_ptr<_Ty>{
261             ::new __control_block<_UncvTy>{(_Us&&)__us...}};
262     }
263 };
264 } // namespace __ptr
265 
266 using __ptr::__enable_intrusive_from_this;
267 using __ptr::__intrusive_ptr;
268 template <class _Ty>
269 inline constexpr __ptr::__make_intrusive_t<_Ty> __make_intrusive{};
270 
271 } // namespace stdexec
272