small_vector.h

Go to the documentation of this file.

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */
 
#pragma once
 
#include <cassert>
#include <cstddef>
#include <initializer_list>
#include <new>
#include <utility>
 
namespace dispenso {
 
template <typename T, size_t N = 4>
class SmallVector {
  static_assert(N > 0, "SmallVector requires at least 1 inline element. Use std::vector for N=0.");
  static_assert(N < 65536, "SmallVector inline capacity is too large. Use std::vector instead.");
 
 public:
  using value_type = T;
  using size_type = size_t;
  using difference_type = std::ptrdiff_t;
  using reference = T&;
  using const_reference = const T&;
  using pointer = T*;
  using const_pointer = const T*;
  using iterator = T*;
  using const_iterator = const T*;
 
  SmallVector() noexcept : size_(0) {}
 
  explicit SmallVector(size_type count) : size_(0) {
    resize(count);
  }
 
  SmallVector(size_type count, const T& value) : size_(0) {
    resize(count, value);
  }
 
  SmallVector(std::initializer_list<T> init) : size_(0) {
    ensureCapacity(init.size());
    for (const auto& v : init) {
      emplace_back(v);
    }
  }
 
  SmallVector(const SmallVector& other) : size_(0) {
    ensureCapacity(other.rawSize());
    for (const auto& v : other) {
      emplace_back(v);
    }
  }
 
  SmallVector(SmallVector&& other) noexcept : size_(0) {
    if (other.isInline()) {
      for (size_type i = 0; i < other.rawSize(); ++i) {
        new (inlineData() + i) T(std::move(other.inlineData()[i]));
        other.inlineData()[i].~T();
      }
    } else {
      storage_.heap_.ptr = other.storage_.heap_.ptr;
      storage_.heap_.capacity = other.storage_.heap_.capacity;
      size_ |= kHeapBit;
    }
    size_ = (size_ & kHeapBit) | other.rawSize();
    other.size_ = 0;
  }
 
  ~SmallVector() {
    destroyAll();
  }
 
  SmallVector& operator=(const SmallVector& other) {
    if (this != &other) {
      destroyAll();
      size_ = 0;
      ensureCapacity(other.rawSize());
      for (const auto& v : other) {
        emplace_back(v);
      }
    }
    return *this;
  }
 
  SmallVector& operator=(SmallVector&& other) noexcept {
    if (this != &other) {
      destroyAll();
      size_ = 0;
 
      if (other.isInline()) {
        for (size_type i = 0; i < other.rawSize(); ++i) {
          new (inlineData() + i) T(std::move(other.inlineData()[i]));
          other.inlineData()[i].~T();
        }
      } else {
        storage_.heap_.ptr = other.storage_.heap_.ptr;
        storage_.heap_.capacity = other.storage_.heap_.capacity;
        size_ |= kHeapBit;
      }
      size_ = (size_ & kHeapBit) | other.rawSize();
      other.size_ = 0;
    }
    return *this;
  }
 
  // --- Element Access ---
 
  reference operator[](size_type pos) {
    return data()[pos];
  }
  const_reference operator[](size_type pos) const {
    return data()[pos];
  }
  reference front() {
    return data()[0];
  }
  const_reference front() const {
    return data()[0];
  }
  reference back() {
    return data()[rawSize() - 1];
  }
  const_reference back() const {
    return data()[rawSize() - 1];
  }
 
  pointer data() noexcept {
    return isInline() ? inlineData() : storage_.heap_.ptr;
  }
  const_pointer data() const noexcept {
    return isInline() ? inlineData() : storage_.heap_.ptr;
  }
 
  // --- Iterators ---
 
  iterator begin() noexcept {
    return data();
  }
  const_iterator begin() const noexcept {
    return data();
  }
  const_iterator cbegin() const noexcept {
    return data();
  }
  iterator end() noexcept {
    return data() + rawSize();
  }
  const_iterator end() const noexcept {
    return data() + rawSize();
  }
  const_iterator cend() const noexcept {
    return data() + rawSize();
  }
 
  // --- Capacity ---
 
  bool empty() const noexcept {
    return rawSize() == 0;
  }
  size_type size() const noexcept {
    return rawSize();
  }
 
  size_type capacity() const noexcept {
    return isInline() ? N : storage_.heap_.capacity;
  }
 
  void reserve(size_type newCap) {
    ensureCapacity(newCap);
  }
 
  // --- Modifiers ---
 
  void clear() noexcept {
    destroyAll();
    size_ = 0;
  }
 
  void push_back(const T& value) {
    emplace_back(value);
  }
  void push_back(T&& value) {
    emplace_back(std::move(value));
  }
 
  template <typename... Args>
  reference emplace_back(Args&&... args) {
    T* ptr;
    if (isInline()) {
      size_type sz = rawSize();
      if (sz < N) {
        ptr = inlineData();
      } else {
        growToHeap(N * 2);
        ptr = storage_.heap_.ptr;
      }
    } else {
      size_type sz = rawSize();
      if (sz == storage_.heap_.capacity) {
        growToHeap(storage_.heap_.capacity * 2);
      }
      ptr = storage_.heap_.ptr;
    }
    size_type idx = rawSize();
    new (ptr + idx) T(std::forward<Args>(args)...);
    // Increment preserves heap bit naturally
    ++size_;
    assert(rawSize() > 0 && "Size overflow into heap bit");
    return ptr[idx];
  }
 
  void pop_back() {
    T* ptr = data();
    size_type sz = rawSize();
    ptr[sz - 1].~T();
    // Decrement preserves heap bit naturally
    --size_;
  }
 
  void resize(size_type count) {
    size_type sz = rawSize();
    if (count > sz) {
      ensureCapacity(count);
      T* ptr = data();
      for (size_type i = sz; i < count; ++i) {
        new (ptr + i) T();
      }
      setSize(count);
    } else if (count < sz) {
      T* ptr = data();
      for (size_type i = count; i < sz; ++i) {
        ptr[i].~T();
      }
      setSize(count);
    }
  }
 
  void resize(size_type count, const T& value) {
    size_type sz = rawSize();
    if (count > sz) {
      ensureCapacity(count);
      T* ptr = data();
      for (size_type i = sz; i < count; ++i) {
        new (ptr + i) T(value);
      }
      setSize(count);
    } else if (count < sz) {
      T* ptr = data();
      for (size_type i = count; i < sz; ++i) {
        ptr[i].~T();
      }
      setSize(count);
    }
  }
 
  iterator erase(const_iterator pos) {
    T* ptr = data();
    size_type sz = rawSize();
    size_type index = pos - ptr;
 
    for (size_type i = index; i + 1 < sz; ++i) {
      ptr[i] = std::move(ptr[i + 1]);
    }
    ptr[sz - 1].~T();
    --size_; // Preserves heap bit
    return data() + index;
  }
 
 private:
  // High bit of size_ tracks heap vs inline mode.
  // Since kHeapBit >> N, comparisons like size_ <= N and size_ < N
  // are naturally false when the heap bit is set, so isInline() and
  // many internal checks work without masking.
  static constexpr size_type kHeapBit = size_type(1) << (sizeof(size_type) * 8 - 1);
  static constexpr size_type kSizeMask = ~kHeapBit;
 
  bool isInline() const noexcept {
    return (size_ & kHeapBit) == 0;
  }
 
  size_type rawSize() const noexcept {
    return size_ & kSizeMask;
  }
 
  // Set the size portion, preserving the heap bit.
  void setSize(size_type s) noexcept {
    assert((s & kHeapBit) == 0 && "Size overflow into heap bit");
    size_ = (size_ & kHeapBit) | s;
  }
 
  T* inlineData() noexcept {
    return reinterpret_cast<T*>(&storage_.inline_);
  }
  const T* inlineData() const noexcept {
    return reinterpret_cast<const T*>(&storage_.inline_);
  }
 
  void destroyAll() noexcept {
    T* ptr = data();
    size_type sz = rawSize();
    for (size_type i = 0; i < sz; ++i) {
      ptr[i].~T();
    }
    if (!isInline()) {
      ::operator delete(storage_.heap_.ptr);
    }
  }
 
  // Grow to heap storage with the specified capacity.
  // Moves existing elements, frees old heap if applicable, sets heap bit.
  void growToHeap(size_type newCap) {
    T* newData = static_cast<T*>(::operator new(newCap * sizeof(T)));
    T* oldData = data();
    size_type sz = rawSize();
 
    for (size_type i = 0; i < sz; ++i) {
      new (newData + i) T(std::move(oldData[i]));
      oldData[i].~T();
    }
 
    if (!isInline()) {
      ::operator delete(storage_.heap_.ptr);
    }
 
    storage_.heap_.ptr = newData;
    storage_.heap_.capacity = newCap;
    size_ = kHeapBit | sz;
  }
 
  void ensureCapacity(size_type newCap) {
    if (newCap <= N && isInline()) {
      return;
    }
    if (isInline()) {
      growToHeap(newCap);
    } else if (newCap > storage_.heap_.capacity) {
      growToHeap(newCap);
    }
  }
 
  size_type size_;
 
  struct HeapStorage {
    T* ptr;
    size_type capacity;
  };
 
  union Storage {
    alignas(T) unsigned char inline_[sizeof(T) * N];
    HeapStorage heap_;
    Storage() noexcept {}
    ~Storage() {}
  } storage_;
};
 
} // namespace dispenso