1 #pragma once 2 #include <array> 3 #include <cstddef> 4 #include <cstdint> 5 #include <type_traits> 6 7 /** Helper class for allocating temporary fixed size arrays with SBO. 8 * 9 * This is intentionally much simpler than SmallVector, to improve performance 10 * at the expense of many features: 11 * - No zero-initialization for numeric types 12 * - No resizing after construction 13 * - No copy/move 14 * - No non-trivial types 15 */ 16 17 namespace c10 { 18 19 template <typename T, size_t N> 20 class SmallBuffer { 21 static_assert(std::is_trivial_v<T>, "SmallBuffer is intended for POD types"); 22 23 std::array<T, N> storage_; 24 size_t size_{}; 25 T* data_{}; 26 27 public: SmallBuffer(size_t size)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 SmallBuffer(SmallBuffer && rhs)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 ~SmallBuffer()53 ~SmallBuffer() { 54 if (size_ > N) { 55 delete[] data_; 56 } 57 } 58 T& operator[](size_t idx) { 59 return data()[idx]; 60 } 61 const T& operator[](size_t idx) const { 62 return data()[idx]; 63 } data()64 T* data() { 65 return data_; 66 } data()67 const T* data() const { 68 return data_; 69 } size()70 size_t size() const { 71 return size_; 72 } begin()73 T* begin() { 74 return data_; 75 } begin()76 const T* begin() const { 77 return data_; 78 } end()79 T* end() { 80 return data_ + size_; 81 } end()82 const T* end() const { 83 return data_ + size_; 84 } 85 }; 86 87 } // namespace c10 88