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
15namespace c10 {
16
17template <typename T, size_t N>
18class 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