Accurately track used bytes in Arena

src/arena_allocator.cpp

@@ -1,6 +1,7 @@
 #include "arena_allocator.hpp"
 #include <cassert>
 #include <iomanip>
+#include <limits>
 
 ArenaAllocator::~ArenaAllocator() {
   while (current_block_) {
@@ -12,10 +13,8 @@ ArenaAllocator::~ArenaAllocator() {
 
 ArenaAllocator::ArenaAllocator(ArenaAllocator &&other) noexcept
     : initial_block_size_(other.initial_block_size_),
-      current_block_(other.current_block_),
+      current_block_(other.current_block_) {
-      current_offset_(other.current_offset_) {
   other.current_block_ = nullptr;
-  other.current_offset_ = 0;
 }
 
 ArenaAllocator &ArenaAllocator::operator=(ArenaAllocator &&other) noexcept {
@@ -28,10 +27,8 @@ ArenaAllocator &ArenaAllocator::operator=(ArenaAllocator &&other) noexcept {
 
     initial_block_size_ = other.initial_block_size_;
     current_block_ = other.current_block_;
-    current_offset_ = other.current_offset_;
 
     other.current_block_ = nullptr;
-    other.current_offset_ = 0;
   }
   return *this;
 }
@@ -58,15 +55,15 @@ void ArenaAllocator::reset() {
   // Update first block's counters to reflect only itself
   if (first_block) {
     first_block->total_size = first_block->size;
-    first_block->block_count = 1;
+    first_block->total_used = 0;
   }
 
   current_block_ = first_block;
-  current_offset_ = 0;
+  current_block_->offset = 0;
 }
 
-void *ArenaAllocator::realloc_raw(void *ptr, size_t old_size, size_t new_size,
+void *ArenaAllocator::realloc_raw(void *ptr, uint32_t old_size,
-                                  size_t alignment) {
+                                  uint32_t new_size, uint32_t alignment) {
   if (ptr == nullptr) {
     return allocate_raw(new_size, alignment);
   }
@@ -79,14 +76,14 @@ void *ArenaAllocator::realloc_raw(void *ptr, size_t old_size, size_t new_size,
   assert(current_block_ &&
          "realloc called with non-null ptr but no current block exists");
 
-  // Assert that current_offset_ is large enough (should always be true for
+  // Assert that offset is large enough (should always be true for
   // valid callers)
-  assert(current_offset_ >= old_size &&
+  assert(current_block_->offset >= old_size &&
-         "current_offset_ must be >= old_size for valid last allocation");
+         "offset must be >= old_size for valid last allocation");
 
   // Check if this was the last allocation by comparing with expected location
   char *expected_last_alloc_start =
-      current_block_->data() + current_offset_ - old_size;
+      current_block_->data() + current_block_->offset - old_size;
 
   if (ptr == expected_last_alloc_start) {
     // This is indeed the last allocation
@@ -94,15 +91,16 @@ void *ArenaAllocator::realloc_raw(void *ptr, size_t old_size, size_t new_size,
       // Growing - check if we have space
       size_t additional_space_needed = new_size - old_size;
 
-      if (current_offset_ + additional_space_needed <= current_block_->size) {
+      if (current_block_->offset + additional_space_needed <=
+          current_block_->size) {
         // We can extend in place
-        current_offset_ += additional_space_needed;
+        current_block_->offset += additional_space_needed;
         return ptr;
       }
     } else {
       // Shrinking - just update the offset
       size_t space_to_free = old_size - new_size;
-      current_offset_ -= space_to_free;
+      current_block_->offset -= space_to_free;
       return new_size == 0 ? nullptr : ptr;
     }
   }
@@ -154,14 +152,7 @@ ArenaAllocator::find_intra_arena_pointers() const {
       uintptr_t block_start = reinterpret_cast<uintptr_t>(b->data());
 
       // Calculate used bytes in this specific block
-      size_t block_used;
+      size_t block_used = b->offset;
-      if (block_idx == 0) {
-        // Current block - use current_offset_
-        block_used = current_offset_;
-      } else {
-        // Previous blocks are fully used
-        block_used = b->size;
-      }
 
       uintptr_t block_used_end = block_start + block_used;
 
@@ -181,8 +172,8 @@ ArenaAllocator::find_intra_arena_pointers() const {
     // Calculate used bytes in this specific block
     size_t block_used;
     if (block_idx == 0) {
-      // Current block - use current_offset_
+      // Current block - use offset
-      block_used = current_offset_;
+      block_used = current_block_->offset;
     } else {
       // Previous blocks are fully used
       block_used = b->size;
@@ -241,8 +232,8 @@ ArenaAllocator::find_address_location(const void *addr) const {
     // Calculate used bytes in this specific block
     size_t block_used;
     if (block_idx == 0) {
-      // Current block - use current_offset_
+      // Current block - use offset
-      block_used = current_offset_;
+      block_used = current_block_->offset;
     } else {
       // Previous blocks are fully used
       block_used = b->size;
@@ -273,9 +264,17 @@ void ArenaAllocator::debug_dump(std::ostream &out, bool show_memory_map,
     return;
   }
 
+  // Build list of blocks from current to first
+  std::vector<Block *> blocks;
+  Block *block = current_block_;
+  while (block) {
+    blocks.push_back(block);
+    block = block->prev;
+  }
+
   // Overall statistics
+  size_t used = this->used_bytes();
   size_t total_alloc = this->total_allocated();
-  size_t used = used_bytes();
   double utilization = total_alloc > 0 ? (100.0 * used / total_alloc) : 0.0;
 
   out << "Total allocated: " << total_alloc << " bytes across " << num_blocks()
@@ -286,14 +285,6 @@ void ArenaAllocator::debug_dump(std::ostream &out, bool show_memory_map,
       << std::endl;
   out << std::endl;
 
-  // Build list of blocks from current to first
-  std::vector<Block *> blocks;
-  Block *block = current_block_;
-  while (block) {
-    blocks.push_back(block);
-    block = block->prev;
-  }
-
   out << "Block Chain (newest to oldest):" << std::endl;
 
   // Display blocks in reverse order (current first)
@@ -301,14 +292,7 @@ void ArenaAllocator::debug_dump(std::ostream &out, bool show_memory_map,
     Block *b = blocks[i];
 
     // Calculate used bytes in this specific block
-    size_t block_used;
+    size_t block_used = b->offset;
-    if (i == 0) {
-      // Current block - use current_offset_
-      block_used = current_offset_;
-    } else {
-      // Previous blocks are fully used
-      block_used = b->size;
-    }
 
     double block_util = b->size > 0 ? (100.0 * block_used / b->size) : 0.0;
 
@@ -365,8 +349,8 @@ void ArenaAllocator::debug_dump(std::ostream &out, bool show_memory_map,
       // Calculate used bytes in this specific block
       size_t block_used;
       if (i == 0) {
-        // Current block - use current_offset_
+        // Current block - use offset
-        block_used = current_offset_;
+        block_used = current_block_->offset;
       } else {
         // Previous blocks are fully used
         block_used = b->size;
@@ -396,13 +380,12 @@ void ArenaAllocator::debug_dump(std::ostream &out, bool show_memory_map,
 void ArenaAllocator::add_block(size_t size) {
   Block *new_block = Block::create(size, current_block_);
   current_block_ = new_block;
-  current_offset_ = 0;
 }
 
 size_t ArenaAllocator::calculate_next_block_size(size_t required_size) const {
-  size_t current_size =
+  size_t doubled_size = total_allocated() * 2;
-      current_block_ ? current_block_->size : initial_block_size_;
+  doubled_size =
-  size_t doubled_size = current_size * 2;
+      std::min<size_t>(doubled_size, std::numeric_limits<uint32_t>::max());
 
   return std::max(required_size, doubled_size);
 }

src/arena_allocator.hpp

@@ -6,6 +6,7 @@
 #include <cstdlib>
 #include <cstring>
 #include <iostream>
+#include <limits>
 #include <new>
 #include <type_traits>
 #include <utility>
@@ -73,11 +74,11 @@ private:
    * - Accumulated counters for O(1) tracking operations
    */
   struct Block {
-    size_t size;       ///< Size of this block's data area
+    uint32_t size;     ///< Size of this block's data area
-    Block *prev;       ///< Pointer to previous block (nullptr for first block)
+    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 block_count; ///< Number of blocks including this one + all previous
+    size_t total_used; ///< Accumulated offsets of previous blocks
-                        ///< blocks
+    Block *prev;       ///< Pointer to previous block (nullptr for first block)
 
     /**
      * @brief Get pointer to the data area of this block.
@@ -93,14 +94,18 @@ private:
      * @throws std::bad_alloc if memory allocation fails
      */
     static Block *create(size_t size, Block *prev) {
+      if (size > std::numeric_limits<uint32_t>::max()) {
+        throw std::bad_alloc();
+      }
       void *memory = std::aligned_alloc(
           alignof(Block), align_up(sizeof(Block) + size, alignof(Block)));
       if (!memory) {
         throw std::bad_alloc();
       }
       size_t total_size = size + (prev ? prev->total_size : 0);
-      size_t block_count = 1 + (prev ? prev->block_count : 0);
+      size_t total_used = prev ? prev->total_used + prev->offset : 0;
-      Block *block = new (memory) Block{size, prev, total_size, block_count};
+      Block *block = new (memory)
+          Block{uint32_t(size), /*offset*/ 0, total_size, total_used, prev};
       return block;
     }
   };
@@ -116,8 +121,7 @@ public:
    * @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),
+      : initial_block_size_(initial_size), current_block_(nullptr) {}
-        current_offset_(0) {}
 
   /**
    * @brief Destructor - frees all allocated blocks.
@@ -181,7 +185,7 @@ public:
    * The allocation path is extremely hot and inlining eliminates function
    * call overhead, allowing the ~1ns allocation performance.
    */
-  void *allocate_raw(size_t size,
+  void *allocate_raw(uint32_t size,
                      size_t alignment = alignof(std::max_align_t)) {
     if (size == 0) {
       return nullptr;
@@ -195,7 +199,7 @@ public:
     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_offset_, alignment) - block_addr;
+        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);
@@ -206,7 +210,7 @@ public:
     }
 
     void *ptr = block_start + aligned_offset;
-    current_offset_ = aligned_offset + size;
+    current_block_->offset = aligned_offset + size;
 
     return ptr;
   }
@@ -247,8 +251,8 @@ public:
    * - Like malloc/realloc, the contents beyond old_size are uninitialized
    * - When copying to new location, uses the specified alignment
    */
-  void *realloc_raw(void *ptr, size_t old_size, size_t new_size,
+  void *realloc_raw(void *ptr, uint32_t old_size, uint32_t new_size,
-                    size_t alignment = alignof(std::max_align_t));
+                    uint32_t alignment = alignof(std::max_align_t));
 
   /**
    * @brief Reallocate memory, extending in place if possible or copying to a
@@ -285,7 +289,11 @@ public:
    * - Like malloc/realloc, the contents beyond old_size are uninitialized
    * - When copying to new location, uses the specified alignment
    */
-  template <typename T> T *realloc(T *ptr, size_t old_size, size_t new_size) {
+  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()) {
+      throw std::bad_alloc();
+    }
     return static_cast<T *>(realloc_raw(ptr, old_size * sizeof(T),
                                         new_size * sizeof(T), alignof(T)));
   }
@@ -367,7 +375,7 @@ public:
    * This method only allocates memory - it does not construct objects.
    * Use placement new or other initialization methods as needed.
    */
-  template <typename T> T *allocate(size_t size) {
+  template <typename T> T *allocate(uint32_t size) {
     static_assert(
         std::is_trivially_destructible_v<T>,
         "ArenaAllocator::allocate requires trivially destructible types. "
@@ -376,6 +384,9 @@ public:
     if (size == 0) {
       return nullptr;
     }
+    if (size_t(size) * sizeof(T) > std::numeric_limits<uint32_t>::max()) {
+      throw std::bad_alloc();
+    }
     void *ptr = allocate_raw(sizeof(T) * size, alignof(T));
     return static_cast<T *>(ptr);
   }
@@ -427,9 +438,7 @@ public:
     if (!current_block_) {
       return 0;
     }
-    size_t prev_total =
+    return current_block_->total_used + current_block_->offset;
-        current_block_->prev ? current_block_->prev->total_size : 0;
-    return prev_total + current_offset_;
   }
 
   /**
@@ -438,18 +447,18 @@ public:
    * @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_offset_ : 0;
+    return current_block_ ? current_block_->size - current_block_->offset : 0;
   }
 
   /**
    * @brief Get the total number of blocks in the allocator.
-   *
-   * Uses O(1) accumulated counters for fast retrieval.
-   *
-   * @return Number of blocks, or 0 if no blocks exist
    */
   size_t num_blocks() const {
-    return current_block_ ? current_block_->block_count : 0;
+    size_t result = 0;
+    for (auto *p = current_block_; p != nullptr; p = p->prev) {
+      ++result;
+    }
+    return result;
   }
 
   /**
@@ -595,11 +604,9 @@ private:
                                    size_t size);
 
   /// Size used for the first block and baseline for geometric growth
-  size_t initial_block_size_;
+  uint32_t initial_block_size_;
   /// Pointer to the current (most recent) block, or nullptr if no blocks exist
   Block *current_block_;
-  /// Current offset within the current block's data area
-  size_t current_offset_;
 };
 
 /**

tests/test_arena_allocator.cpp

@@ -2,7 +2,6 @@
 #include "arena_allocator.hpp"
 #include <cstring>
 #include <doctest/doctest.h>
-#include <string>
 #include <vector>
 
 TEST_CASE("ArenaAllocator basic construction") {