1 | /* Copyright 2021 The TensorFlow Authors. All Rights Reserved. |
2 | |
3 | Licensed under the Apache License, Version 2.0 (the "License"); |
4 | you may not use this file except in compliance with the License. |
5 | You may obtain a copy of the License at |
6 | |
7 | http://www.apache.org/licenses/LICENSE-2.0 |
8 | |
9 | Unless required by applicable law or agreed to in writing, software |
10 | distributed under the License is distributed on an "AS IS" BASIS, |
11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | See the License for the specific language governing permissions and |
13 | limitations under the License. |
14 | ==============================================================================*/ |
15 | #ifndef TENSORFLOW_TSL_PLATFORM_INTRUSIVE_PTR_H_ |
16 | #define TENSORFLOW_TSL_PLATFORM_INTRUSIVE_PTR_H_ |
17 | |
18 | #include <algorithm> |
19 | namespace tsl { |
20 | namespace core { |
21 | |
22 | // A utility for managing the lifetime of ref-counted objects. |
23 | // |
24 | // Generally used for objects that derive from `tensorflow::RefCounted`. |
25 | template <class T> |
26 | class IntrusivePtr { |
27 | public: |
28 | // add_ref=false indicates that IntrusivePtr owns the underlying pointer. |
29 | // |
30 | // In most cases, we expect this to be called with add_ref=false, except in |
31 | // special circumstances where the lifetime of the underlying RefCounted |
32 | // object needs to be externally managed. |
33 | IntrusivePtr(T* h, bool add_ref) { reset(h, add_ref); } |
34 | IntrusivePtr(const IntrusivePtr& o) { reset(o.handle_, /*add_ref=*/true); } |
35 | IntrusivePtr(IntrusivePtr&& o) { *this = std::move(o); } |
36 | IntrusivePtr() {} |
37 | void reset(T* h, bool add_ref) { |
38 | if (h != handle_) { |
39 | if (add_ref && h) h->Ref(); |
40 | if (handle_) handle_->Unref(); |
41 | handle_ = h; |
42 | } |
43 | } |
44 | IntrusivePtr& operator=(const IntrusivePtr& o) { |
45 | reset(o.handle_, /*add_ref=*/true); |
46 | return *this; |
47 | } |
48 | IntrusivePtr& operator=(IntrusivePtr&& o) { |
49 | if (handle_ != o.handle_) { |
50 | // Must clear o.handle_ before calling reset to capture the case where |
51 | // handle_->member == o. In this case, calling handle_->Unref first would |
52 | // delete o.handle_ so we clear it out first. |
53 | reset(o.detach(), /*add_ref=*/false); |
54 | } |
55 | return *this; |
56 | } |
57 | bool operator==(const IntrusivePtr& o) const { return handle_ == o.handle_; } |
58 | T* operator->() const { return handle_; } |
59 | T& operator*() const { return *handle_; } |
60 | explicit operator bool() const noexcept { return get(); } |
61 | T* get() const { return handle_; } |
62 | // Releases ownership of the pointer without unreffing. Caller is responsible |
63 | // for calling Unref on the returned pointer. |
64 | T* detach() { |
65 | T* handle = handle_; |
66 | handle_ = nullptr; |
67 | return handle; |
68 | } |
69 | |
70 | ~IntrusivePtr() { |
71 | if (handle_) handle_->Unref(); |
72 | } |
73 | |
74 | private: |
75 | T* handle_ = nullptr; |
76 | }; |
77 | |
78 | } // namespace core |
79 | } // namespace tsl |
80 | |
81 | #endif // TENSORFLOW_TSL_PLATFORM_INTRUSIVE_PTR_H_ |
82 | |