1 | /* Copyright 2017 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 | |
16 | #ifndef TENSORFLOW_TSL_UTIL_PTR_UTIL_H_ |
17 | #define TENSORFLOW_TSL_UTIL_PTR_UTIL_H_ |
18 | |
19 | // Utility functions for pointers. |
20 | |
21 | #include <stddef.h> |
22 | |
23 | #include <memory> |
24 | #include <type_traits> |
25 | #include <utility> |
26 | |
27 | namespace tsl { |
28 | |
29 | namespace helper { |
30 | |
31 | // Trait to select overloads and return types for MakeUnique. |
32 | template <typename T> |
33 | struct MakeUniqueResult { |
34 | using scalar = std::unique_ptr<T>; |
35 | }; |
36 | template <typename T> |
37 | struct MakeUniqueResult<T[]> { |
38 | using array = std::unique_ptr<T[]>; |
39 | }; |
40 | template <typename T, size_t N> |
41 | struct MakeUniqueResult<T[N]> { |
42 | using invalid = void; |
43 | }; |
44 | |
45 | } // namespace helper |
46 | |
47 | // Transfers ownership of a raw pointer to a std::unique_ptr of deduced type. |
48 | // Example: |
49 | // X* NewX(int, int); |
50 | // auto x = WrapUnique(NewX(1, 2)); // 'x' is std::unique_ptr<X>. |
51 | // |
52 | // WrapUnique is useful for capturing the output of a raw pointer factory. |
53 | // However, prefer 'MakeUnique<T>(args...) over 'WrapUnique(new T(args...))'. |
54 | // auto x = WrapUnique(new X(1, 2)); // works, but nonideal. |
55 | // auto x = MakeUnique<X>(1, 2); // safer, standard, avoids raw 'new'. |
56 | // |
57 | // Note: Cannot wrap pointers to array of unknown bound (i.e. U(*)[]). |
58 | template <typename T> |
59 | std::unique_ptr<T> WrapUnique(T* ptr) { |
60 | static_assert(!std::is_array<T>::value || std::extent<T>::value != 0, |
61 | "types T[0] or T[] are unsupported" ); |
62 | return std::unique_ptr<T>(ptr); |
63 | } |
64 | |
65 | template <typename T, typename... Args> |
66 | typename helper::MakeUniqueResult<T>::scalar MakeUnique(Args&&... args) { |
67 | return std::unique_ptr<T>(new T(std::forward<Args>(args)...)); |
68 | } |
69 | |
70 | // Overload for array of unknown bound. |
71 | // The allocation of arrays needs to use the array form of new, |
72 | // and cannot take element constructor arguments. |
73 | template <typename T> |
74 | typename helper::MakeUniqueResult<T>::array MakeUnique(size_t n) { |
75 | return std::unique_ptr<T>(new typename std::remove_extent<T>::type[n]()); |
76 | } |
77 | |
78 | } // namespace tsl |
79 | |
80 | #endif // TENSORFLOW_TSL_UTIL_PTR_UTIL_H_ |
81 | |