1 | #pragma once |
2 | #include <array> |
3 | #include <type_traits> |
4 | |
5 | /** Helper class for allocating temporary fixed size arrays with SBO. |
6 | * |
7 | * This is intentionally much simpler than SmallVector, to improve performace at |
8 | * the expense of many features: |
9 | * - No zero-initialization for numeric types |
10 | * - No resizing after construction |
11 | * - No copy/move |
12 | * - No non-trivial types |
13 | */ |
14 | |
15 | namespace c10 { |
16 | |
17 | template <typename T, size_t N> |
18 | class SmallBuffer { |
19 | static_assert( |
20 | std::is_trivial<T>::value, |
21 | "SmallBuffer is intended for POD types" ); |
22 | |
23 | std::array<T, N> storage_; |
24 | size_t size_{}; |
25 | T* data_{}; |
26 | |
27 | public: |
28 | SmallBuffer(size_t size) : size_(size) { |
29 | if (size > N) { |
30 | data_ = new T[size]; |
31 | } else { |
32 | data_ = &storage_[0]; |
33 | } |
34 | } |
35 | |
36 | SmallBuffer(const SmallBuffer&) = delete; |
37 | SmallBuffer& operator=(const SmallBuffer&) = delete; |
38 | |
39 | // move constructor is needed in function return |
40 | SmallBuffer(SmallBuffer&& rhs) noexcept : size_{rhs.size_} { |
41 | rhs.size_ = 0; |
42 | if (size_ > N) { |
43 | data_ = rhs.data_; |
44 | rhs.data_ = nullptr; |
45 | } else { |
46 | storage_ = std::move(rhs.storage_); |
47 | data_ = &storage_[0]; |
48 | } |
49 | } |
50 | |
51 | SmallBuffer& operator=(SmallBuffer&&) = delete; |
52 | |
53 | ~SmallBuffer() { |
54 | if (size_ > N) { |
55 | delete[] data_; |
56 | } |
57 | } |
58 | |
59 | T& operator[](int64_t idx) { |
60 | return data()[idx]; |
61 | } |
62 | const T& operator[](int64_t idx) const { |
63 | return data()[idx]; |
64 | } |
65 | T* data() { |
66 | return data_; |
67 | } |
68 | const T* data() const { |
69 | return data_; |
70 | } |
71 | size_t size() const { |
72 | return size_; |
73 | } |
74 | T* begin() { |
75 | return data_; |
76 | } |
77 | const T* begin() const { |
78 | return data_; |
79 | } |
80 | T* end() { |
81 | return data_ + size_; |
82 | } |
83 | const T* end() const { |
84 | return data_ + size_; |
85 | } |
86 | }; |
87 | |
88 | } // namespace c10 |
89 | |