weaseldb/src/arena_allocator.hpp

#pragma once

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <limits>
#include <new>
#include <type_traits>
#include <typeinfo>
#include <utility>
#include <vector>

/**
 * @brief A high-performance arena allocator for bulk allocations.
 *
 * ArenaAllocator provides extremely fast memory allocation (~1ns per
 * allocation) by allocating large blocks and serving allocations from them
 * sequentially. It's designed for scenarios where many small objects need to be
 * allocated and can all be deallocated together.
 *
 * ## Key Features:
 * - **Ultra-fast allocation**: ~1ns per allocation vs ~20-270ns for malloc
 * - **Lazy initialization**: No memory allocated until first use
 * - **Intrusive linked list**: Minimal memory overhead using backward-linked
 * blocks
 * - **Geometric growth**: Block sizes double to minimize allocations
 * - **Memory efficient reset**: Frees unused blocks to prevent memory leaks
 * - **Proper alignment**: Respects alignment requirements for all types
 *
 * ## Performance Characteristics:
 * - Allocation: O(1) amortized
 * - Memory tracking: O(1) using accumulated counters
 * - Reset: O(n) where n is number of blocks (but frees memory)
 * - Destruction: O(n) where n is number of blocks
 *
 * ## Usage Examples:
 * ```cpp
 * ArenaAllocator arena(1024);
 * void* ptr = arena.allocate_raw(100);
 * int* num = arena.construct<int>(42);
 * arena.reset(); // Reuse arena memory
 * ```
 *
 * ## Memory Management:
 * - Individual objects cannot be freed (by design)
 * - All memory is freed when the allocator is destroyed
 * - reset() frees all blocks except the first one
 * - Move semantics transfer ownership of all blocks
 *
 * ## Thread Safety:
 * Not thread-safe. Use separate instances per thread or external
 * synchronization.
 */
class ArenaAllocator {
private:
  /**
   * @brief Internal block structure for the intrusive linked list.
   *
   * Each block contains:
   * - The actual data storage immediately following the Block header
   * - Backward pointer to previous block (intrusive linked list)
   * - Accumulated counters for O(1) tracking operations
   */
  struct Block {
    uint32_t size;     ///< Size of this block's data area
    uint32_t offset;   ///< The offset of the first unused byte in the data area
    size_t total_size; ///< Accumulated size of this block + all previous blocks
    size_t total_used; ///< Accumulated offsets of previous blocks
    Block *prev;       ///< Pointer to previous block (nullptr for first block)

    /**
     * @brief Get pointer to the data area of this block.
     * @return Pointer to the start of the data area (after Block header).
     */
    char *data() { return reinterpret_cast<char *>(this + 1); }

    /**
     * @brief Create a new block with the specified size.
     * @param size Size of the data area for this block
     * @param prev Pointer to the previous block (nullptr for first block)
     * @return Pointer to the newly created block
     * @note Prints error to stderr and calls std::abort() if memory allocation
     * fails
     */
    static Block *create(size_t size, Block *prev) {
      if (size > std::numeric_limits<uint32_t>::max()) {
        std::fprintf(
            stderr,
            "ArenaAllocator: Block size %zu exceeds maximum uint32_t value\n",
            size);
        std::abort();
      }
      void *memory = std::aligned_alloc(
          alignof(Block), align_up(sizeof(Block) + size, alignof(Block)));
      if (!memory) {
        std::fprintf(
            stderr,
            "ArenaAllocator: Failed to allocate memory block of size %zu\n",
            size);
        std::abort();
      }
      size_t total_size = size + (prev ? prev->total_size : 0);
      size_t total_used = prev ? prev->total_used + prev->offset : 0;
      Block *block = new (memory)
          Block{uint32_t(size), /*offset*/ 0, total_size, total_used, prev};
      return block;
    }
  };

public:
  /**
   * @brief Construct an ArenaAllocator with the specified initial block size.
   *
   * No memory is allocated until the first allocation request (lazy
   * initialization). The initial block size is used for the first block and as
   * the baseline for geometric growth.
   *
   * @param initial_size Size in bytes for the first block (default: 1024)
   */
  explicit ArenaAllocator(size_t initial_size = 1024)
      : initial_block_size_(initial_size), current_block_(nullptr) {}

  /**
   * @brief Destructor - frees all allocated blocks.
   *
   * Traverses the intrusive linked list backwards from current_block_,
   * freeing each block. This ensures no memory leaks.
   */
  ~ArenaAllocator();

  /// Copy construction is not allowed (would be expensive and error-prone)
  ArenaAllocator(const ArenaAllocator &) = delete;
  /// Copy assignment is not allowed (would be expensive and error-prone)
  ArenaAllocator &operator=(const ArenaAllocator &) = delete;

  /**
   * @brief Move constructor - transfers ownership of all blocks.
   * @param other The ArenaAllocator to move from (will be left empty)
   */
  ArenaAllocator(ArenaAllocator &&other) noexcept;

  /**
   * @brief Move assignment operator - transfers ownership of all blocks.
   *
   * Frees any existing blocks in this allocator before taking ownership
   * of blocks from the other allocator.
   *
   * @param other The ArenaAllocator to move from (will be left empty)
   * @return Reference to this allocator
   */
  ArenaAllocator &operator=(ArenaAllocator &&other) noexcept;

  /**
   * @brief Allocate raw memory with the specified size and alignment.
   *
   * This is the core allocation method providing ~1ns allocation performance.
   * It performs lazy initialization on first use and automatically grows
   * the arena when needed using geometric growth (doubling block sizes).
   *
   * For type-safe allocation, prefer the allocate<T>() template method.
   *
   * @param size Number of bytes to allocate (0 returns nullptr)
   * @param alignment Required alignment (default: alignof(std::max_align_t))
   * @return Pointer to allocated memory, or nullptr if size is 0
   * @note Prints error to stderr and calls std::abort() if memory allocation
   * fails
   *
   * ## Performance:
   * - O(1) amortized allocation time
   * - Respects alignment requirements with minimal padding
   * - Automatically creates new blocks when current block is exhausted
   *
   * @note This method is kept inline for maximum performance (~1ns allocation).
   */
  void *allocate_raw(uint32_t size,
                     size_t alignment = alignof(std::max_align_t)) {
    if (size == 0) {
      return nullptr;
    }

    if (!current_block_) {
      size_t block_size = std::max(size, initial_block_size_);
      add_block(block_size);
    }

    char *block_start = current_block_->data();
    uintptr_t block_addr = reinterpret_cast<uintptr_t>(block_start);
    size_t aligned_offset =
        align_up(block_addr + current_block_->offset, alignment) - block_addr;

    if (aligned_offset + size > current_block_->size) {
      size_t next_block_size = calculate_next_block_size(size);
      add_block(next_block_size);
      block_start = current_block_->data();
      block_addr = reinterpret_cast<uintptr_t>(block_start);
      aligned_offset = align_up(block_addr, alignment) - block_addr;
    }

    void *ptr = block_start + aligned_offset;
    current_block_->offset = aligned_offset + size;

    return ptr;
  }

  /**
   * @brief Reallocate memory, extending in place if possible or copying to a
   * new location.
   *
   * This method provides realloc-like functionality for the arena allocator.
   * If the given pointer was the last allocation and there's sufficient space
   * in the current block to extend it, the allocation is grown in place.
   * Otherwise, a new allocation is made and the old data is copied.
   *
   * @param ptr Pointer to the existing allocation (must be from this allocator)
   * @param old_size Size of the existing allocation in bytes
   * @param new_size Desired new size in bytes
   * @param alignment Required alignment. Defaults to
   * `alignof(std::max_align_t)`
   * @return Pointer to the reallocated memory (may be the same as ptr or
   * different)
   * @note Prints error to stderr and calls std::abort() if memory allocation
   * fails
   *
   * ## Behavior:
   * - If new_size == old_size, returns ptr unchanged
   * - If new_size == 0, returns nullptr (no deallocation occurs)
   * - If ptr is null, behaves like allocate(new_size, alignment)
   * - If ptr was the last allocation and space exists, extends in place
   *
   *
   * ## Safety Notes:
   * - The caller must provide the correct old_size - this is not tracked
   * - The old pointer becomes invalid if a copy occurs
   * - Like malloc/realloc, the contents beyond old_size are uninitialized
   * - When copying to new location, uses the specified alignment
   */
  void *realloc_raw(void *ptr, uint32_t old_size, uint32_t new_size,
                    uint32_t alignment = alignof(std::max_align_t));

  /**
   * @brief Type-safe version of realloc_raw for arrays of type T.
   *
   * @param ptr Pointer to the existing allocation (must be from this allocator)
   * @param old_size Size of the existing allocation in number of T objects
   * @param new_size Desired new size in number of T objects
   * @return Pointer to the reallocated memory (may be the same as ptr or
   * different)
   * @note Prints error to stderr and calls std::abort() if memory allocation
   * fails or size overflow occurs
   */
  template <typename T>
  T *realloc(T *ptr, uint32_t old_size, uint32_t new_size) {
    if (size_t(new_size) * sizeof(T) > std::numeric_limits<uint32_t>::max()) {
      std::fprintf(stderr,
                   "ArenaAllocator: Reallocation size overflow for type %s "
                   "(new_size=%u, sizeof(T)=%zu)\n",
                   typeid(T).name(), new_size, sizeof(T));
      std::abort();
    }
    return static_cast<T *>(realloc_raw(ptr, old_size * sizeof(T),
                                        new_size * sizeof(T), alignof(T)));
  }

  /**
   * @brief Construct an object of type T in the arena using placement new.
   *
   * This is a convenience method that combines allocation with in-place
   * construction. It properly handles alignment requirements for type T.
   *
   * @tparam T The type of object to construct (must be trivially destructible)
   * @tparam Args Types of constructor arguments
   * @param args Arguments to forward to T's constructor
   * @return Pointer to the constructed object
   * @note Prints error to stderr and calls std::abort() if memory allocation
   * fails
   *
   * ## Type Requirements:
   * T must be trivially destructible (std::is_trivially_destructible_v<T>).
   * This prevents subtle bugs since destructors are never called for objects
   * constructed in the arena.
   *
   *
   * ## Note:
   * Objects constructed this way cannot be individually destroyed.
   * Their destructors will NOT be called automatically - hence the requirement
   * for trivially destructible types.
   */
  template <typename T, typename... Args> T *construct(Args &&...args) {
    static_assert(
        std::is_trivially_destructible_v<T>,
        "ArenaAllocator::construct requires trivially destructible types. "
        "Objects constructed in the arena will not have their destructors "
        "called.");
    void *ptr = allocate_raw(sizeof(T), alignof(T));
    return new (ptr) T(std::forward<Args>(args)...);
  }

  /**
   * @brief Allocate space for an array of size T objects with proper alignment.
   *
   * This is a type-safe convenience method that combines sizing and alignment
   * calculations for allocating arrays of type T. It's preferred over calling
   * allocate_raw() directly as it prevents common errors with size calculations
   * and alignment requirements.
   *
   * @tparam T The type of objects to allocate space for (must be trivially
   * destructible)
   * @param size Number of T objects to allocate space for
   * @return Pointer to allocated memory suitable for constructing an array of T
   * objects
   * @note Prints error to stderr and calls std::abort() if memory allocation
   * fails
   *
   * ## Type Requirements:
   * T must be trivially destructible (std::is_trivially_destructible_v<T>).
   * This ensures consistency with the arena allocator's design where
   * destructors are never called.
   *
   *
   * ## Note:
   * This method only allocates memory - it does not construct objects.
   * Use placement new or other initialization methods as needed.
   */
  template <typename T> T *allocate(uint32_t size) {
    static_assert(
        std::is_trivially_destructible_v<T>,
        "ArenaAllocator::allocate requires trivially destructible types. "
        "Objects allocated in the arena will not have their destructors "
        "called.");
    if (size == 0) {
      return nullptr;
    }
    if (size_t(size) * sizeof(T) > std::numeric_limits<uint32_t>::max()) {
      std::fprintf(stderr,
                   "ArenaAllocator: Allocation size overflow for type %s "
                   "(size=%u, sizeof(T)=%zu)\n",
                   typeid(T).name(), size, sizeof(T));
      std::abort();
    }
    void *ptr = allocate_raw(sizeof(T) * size, alignof(T));
    return static_cast<T *>(ptr);
  }

  /**
   * @brief Reset the allocator to reuse the first block, freeing all others.
   *
   * This method provides memory-efficient reset behavior by:
   * 1. Keeping the first block for reuse
   * 2. Freeing all subsequent blocks to prevent memory leaks
   * 3. Resetting allocation position to the start of the first block
   *
   * If no blocks have been allocated yet, this is a no-op.
   *
   * ## Performance:
   * - O(n) where n is the number of blocks to free
   * - Prevents memory leaks by freeing unused blocks
   * - Faster than destroying and recreating the allocator
   *
   */
  void reset();

  /**
   * @brief Get the total number of bytes allocated across all blocks.
   *
   * Uses O(1) accumulated counters for fast retrieval.
   *
   * @return Total allocated bytes, or 0 if no blocks exist
   */
  size_t total_allocated() const {
    return current_block_ ? current_block_->total_size : 0;
  }

  /**
   * @brief Get the number of bytes currently used for allocations.
   *
   * This includes all fully used previous blocks plus the used portion
   * of the current block. Uses O(1) accumulated counters.
   *
   * @return Number of bytes in use
   */
  size_t used_bytes() const {
    if (!current_block_) {
      return 0;
    }
    return current_block_->total_used + current_block_->offset;
  }

  /**
   * @brief Get the number of bytes available in the current block.
   *
   * @return Available bytes in current block, or 0 if no blocks exist
   */
  size_t available_in_current_block() const {
    return current_block_ ? current_block_->size - current_block_->offset : 0;
  }

  /**
   * @brief Get the total number of blocks in the allocator.
   *
   * @note This function is primarily used for testing and debugging.
   * It has O(n) complexity as it traverses the entire block chain.
   */
  size_t num_blocks() const {
    size_t result = 0;
    for (auto *p = current_block_; p != nullptr; p = p->prev) {
      ++result;
    }
    return result;
  }

  /**
   * @brief Debug function to visualize the arena's layout and contents.
   *
   * @note This function is intended for testing, debugging, and development
   * only. It should not be used in production code due to performance overhead.
   *
   * Prints a detailed breakdown of all blocks, memory usage, and allocation
   * patterns. This is useful for understanding memory fragmentation and
   * allocation behavior during development and debugging.
   *
   * @param out Output stream to write debug information to (default: std::cout)
   * @param show_memory_map If true, shows a visual memory map of used/free
   * space
   * @param show_content If true, shows actual memory contents in hex and ASCII
   * @param content_limit Maximum bytes of content to show per block (default:
   * 256)
   */
  void debug_dump(std::ostream &out = std::cout, bool show_memory_map = false,
                  bool show_content = false, size_t content_limit = 256) const;

private:
  /**
   * @brief Add a new block with the specified size to the allocator.
   *
   * Creates a new block and makes it the current block. Updates all
   * accumulated counters automatically through Block::create().
   *
   * @param size Size of the data area for the new block
   */
  void add_block(size_t size);

  /**
   * @brief Calculate the size for the next block using geometric growth.
   *
   * Uses a doubling strategy to minimize the number of blocks while
   * ensuring large allocations are handled efficiently.
   *
   * @param required_size Minimum size needed for the allocation
   * @return Size for the next block (max of required_size and doubled current
   * size)
   */
  size_t calculate_next_block_size(size_t required_size) const;

  /**
   * @brief Align a value up to the specified alignment boundary.
   *
   * Uses bit manipulation for efficient alignment calculation.
   * Only works with power-of-2 alignments.
   *
   * This method is kept inline in the header for maximum performance
   * as it's called in the hot allocation path and benefits from inlining.
   *
   * @param value The value to align
   * @param alignment The alignment boundary (must be power of 2)
   * @return The aligned value
   */
  static size_t align_up(size_t value, size_t alignment) {
    if (alignment == 0 || (alignment & (alignment - 1)) != 0) {
      return value;
    }
    return (value + alignment - 1) & ~(alignment - 1);
  }

  /**
   * @brief Dump memory contents in hex/ASCII format.
   *
   * Displays memory in the classic hex dump format with 16 bytes per line,
   * showing both hexadecimal values and ASCII representation.
   *
   * @param out Output stream to write to
   * @param data Pointer to the memory to dump
   * @param size Number of bytes to dump
   */
  static void dump_memory_contents(std::ostream &out, const char *data,
                                   size_t size);

  /// Size used for the first block and baseline for geometric growth
  uint32_t initial_block_size_;
  /// Pointer to the current (most recent) block, or nullptr if no blocks exist
  Block *current_block_;
};

/**
 * @brief STL-compatible allocator that uses ArenaAllocator for memory
 * management.
 * @tparam T The type of objects to allocate
 */
template <typename T> class ArenaStlAllocator {
public:
  using value_type = T;
  using pointer = T *;
  using const_pointer = const T *;
  using reference = T &;
  using const_reference = const T &;
  using size_type = std::size_t;
  using difference_type = std::ptrdiff_t;

  template <typename U> struct rebind {
    using other = ArenaStlAllocator<U>;
  };

  explicit ArenaStlAllocator(ArenaAllocator *arena) noexcept : arena_(arena) {}

  template <typename U>
  ArenaStlAllocator(const ArenaStlAllocator<U> &other) noexcept
      : arena_(other.arena_) {}

  T *allocate(size_type n) {
    if (n == 0)
      return nullptr;
    return arena_->allocate<T>(n);
  }

  void deallocate(T *, size_type) noexcept {
    // Arena allocator doesn't support individual deallocation
  }

  template <typename U>
  bool operator==(const ArenaStlAllocator<U> &other) const noexcept {
    return arena_ == other.arena_;
  }

  template <typename U>
  bool operator!=(const ArenaStlAllocator<U> &other) const noexcept {
    return arena_ != other.arena_;
  }

  ArenaAllocator *arena_;

  template <typename U> friend class ArenaStlAllocator;
};