StackAllocator.hpp

#ifndef INTNS_MEMORY_STACKALLOCATOR_HPP
#define INTNS_MEMORY_STACKALLOCATOR_HPP
 
#include <cstddef>
#include <cstdint>
#include <optional>
#include <stdexcept>
 
#include "Alignment.hpp"
 
namespace intns::memory {
 
class StackAllocator {
 public:
  using marker_type = std::uintptr_t;
  using checkpoint_t = marker_type;
  using size_type = std::size_t;
 
  // Validate our assumptions
  static_assert(sizeof(marker_type) >= sizeof(void*),
                "marker_type must be able to hold pointer values");
  static_assert(std::is_same_v<marker_type, std::uintptr_t>,
                "marker_type should be uintptr_t for pointer arithmetic");
  static_assert(sizeof(size_type) >= sizeof(std::size_t),
                "size_type must accommodate all possible allocation sizes");
 
  StackAllocator(size_type size = 1000);
 
  StackAllocator(void* memory, size_type size);
 
  ~StackAllocator();
 
  // Don't allow moving or copying
  StackAllocator(const StackAllocator&) = delete;
  StackAllocator& operator=(const StackAllocator&) = delete;
  StackAllocator(StackAllocator&&) = delete;
  StackAllocator& operator=(StackAllocator&&) = delete;
 
  template <typename T>
  [[nodiscard]] T* alloc_t() noexcept {
    static_assert(!std::is_reference_v<T>,
                  "Cannot allocate storage for reference types");
    static_assert(!std::is_void_v<T>, "Cannot allocate storage for void");
    static_assert(std::is_destructible_v<T>, "Type must be destructible");
    static_assert(sizeof(T) > 0, "Type must have non-zero size");
 
    constexpr size_type type_size = sizeof(T);
    constexpr size_type align_req = alignof(T);  // Always power of 2 per spec
 
    if (type_size > capacity_) [[unlikely]] {
      return nullptr;
    }
 
    const auto aligned_pos = align_address(active_marker_, align_req);
 
    // This calculates 'type_size' BACK from the end, where are we last valid?
    const auto last_valid_start_pos = (start_marker_ + capacity_) - type_size;
    if (aligned_pos < start_marker_ || aligned_pos > last_valid_start_pos) {
      return nullptr;
    }
 
    // Get aligned address and advance the marker
    T* new_addr = reinterpret_cast<T*>(aligned_pos);
    active_marker_ = aligned_pos + type_size;
    return new_addr;
  }
 
  [[nodiscard]] void* alloc(size_type size, size_type alignment = alignof(
                                                std::max_align_t)) noexcept {
    if (size == 0 || size > capacity_) {
      return nullptr;
    }
 
    // If anything is awry with the alignment, bail
    if (!is_power_of_two(alignment) || alignment > capacity_) [[unlikely]] {
      return nullptr;
    }
 
    const auto aligned_pos = align_address(active_marker_, alignment);
 
    // This calculates 'size' BACK from the end, where are we last valid?
    const auto last_valid_start_pos = (start_marker_ + capacity_) - size;
    if (aligned_pos < start_marker_ || aligned_pos > last_valid_start_pos) {
      return nullptr;
    }
 
    // Get aligned address and advance the marker
    void* new_addr = reinterpret_cast<void*>(aligned_pos);
    active_marker_ = aligned_pos + size;
    return new_addr;
  }
 
  [[nodiscard]] checkpoint_t save_checkpoint() const noexcept {
    return active_marker_;
  }
 
  void restore_checkpoint(checkpoint_t checkpoint) {
    // Check the bounds of the checkpoint: [S] checkpoint lives here [E]
    if (checkpoint < start_marker_ || checkpoint > start_marker_ + capacity_)
        [[unlikely]] {
      throw std::runtime_error(
          "StackAllocator::restore_checkpoint: Invalid checkpoint");
    }
 
    active_marker_ = checkpoint;
  }
 
  [[nodiscard]] size_type bytes_used() const noexcept {
    return active_marker_ - start_marker_;
  }
 
  [[nodiscard]] size_type bytes_remaining() const noexcept {
    return capacity_ - bytes_used();
  }
 
  [[nodiscard]] size_type capacity() const noexcept { return capacity_; }
 
  void reset() noexcept { active_marker_ = start_marker_; }
 
 private:
  // Marker to track the beginning of the stack
  marker_type start_marker_ = 0;
 
  // Marker to track the current position of the stack
  marker_type active_marker_ = 0;
 
  // The full size of the stack
  size_type capacity_ = 0;
 
  // Tracks ownership of the start marker
  bool owns_memory_ = true;
};
 
class StackCheckpoint {
  // The parent allocator
  StackAllocator& allocator_;
 
  // The position at object creation, to be restored on destruction
  StackAllocator::checkpoint_t saved_ = 0;
 
 public:
  StackCheckpoint(StackAllocator& a)
      : allocator_(a), saved_(a.save_checkpoint()) {}
 
  ~StackCheckpoint() { allocator_.restore_checkpoint(saved_); }
 
  // Don't allow moving or copying - this is a scoped checkpoint!
  StackCheckpoint(const StackCheckpoint&) = delete;
  StackCheckpoint& operator=(const StackCheckpoint&) = delete;
  StackCheckpoint(StackCheckpoint&&) = delete;
  StackCheckpoint& operator=(StackCheckpoint&&) = delete;
};
 
}  // namespace intns::memory
 
#endif