Don't traverse to count; fix memory leak in reset

src/arena_allocator.hpp

@@ -1,41 +1,45 @@
 #pragma once
 
+#include <algorithm>
 #include <cstddef>
+#include <cstdint>
 #include <cstdlib>
-#include <memory>
 #include <new>
+#include <utility>
 
 class ArenaAllocator {
 private:
   struct Block {
     size_t size;
-    Block *next;
+    Block *prev;
+    size_t total_size;  // Accumulated size of this block + all previous blocks
+    size_t block_count; // Number of blocks including this one + all previous
+                        // blocks
 
     char *data() { return reinterpret_cast<char *>(this + 1); }
 
-    static Block *create(size_t size) {
+    static Block *create(size_t size, Block *prev) {
       void *memory = std::aligned_alloc(alignof(Block), sizeof(Block) + size);
       if (!memory) {
         throw std::bad_alloc();
       }
-      Block *block = new (memory) Block{size, nullptr};
+      size_t total_size = size + (prev ? prev->total_size : 0);
+      size_t block_count = 1 + (prev ? prev->block_count : 0);
+      Block *block = new (memory) Block{size, prev, total_size, block_count};
       return block;
     }
   };
 
 public:
   explicit ArenaAllocator(size_t initial_size = 1024)
-      : initial_block_size_(initial_size), first_block_(nullptr),
-        current_block_(nullptr), current_offset_(0) {
-    add_block(initial_size);
-  }
+      : initial_block_size_(initial_size), current_block_(nullptr),
+        current_offset_(0) {}
 
   ~ArenaAllocator() {
-    Block *current = first_block_;
-    while (current) {
-      Block *next = current->next;
-      std::free(current);
-      current = next;
+    while (current_block_) {
+      Block *prev = current_block_->prev;
+      std::free(current_block_);
+      current_block_ = prev;
     }
   }
 
@@ -44,28 +48,24 @@ public:
 
   ArenaAllocator(ArenaAllocator &&other) noexcept
       : initial_block_size_(other.initial_block_size_),
-        first_block_(other.first_block_), current_block_(other.current_block_),
+        current_block_(other.current_block_),
         current_offset_(other.current_offset_) {
-    other.first_block_ = nullptr;
     other.current_block_ = nullptr;
     other.current_offset_ = 0;
   }
 
   ArenaAllocator &operator=(ArenaAllocator &&other) noexcept {
     if (this != &other) {
-      Block *current = first_block_;
-      while (current) {
-        Block *next = current->next;
-        std::free(current);
-        current = next;
+      while (current_block_) {
+        Block *prev = current_block_->prev;
+        std::free(current_block_);
+        current_block_ = prev;
       }
 
       initial_block_size_ = other.initial_block_size_;
-      first_block_ = other.first_block_;
       current_block_ = other.current_block_;
       current_offset_ = other.current_offset_;
 
-      other.first_block_ = nullptr;
       other.current_block_ = nullptr;
       other.current_offset_ = 0;
     }
@@ -78,7 +78,8 @@ public:
     }
 
     if (!current_block_) {
-      return nullptr;
+      size_t block_size = std::max(size, initial_block_size_);
+      add_block(block_size);
     }
 
     char *block_start = current_block_->data();
@@ -106,28 +107,45 @@ public:
   }
 
   void reset() {
-    current_block_ = first_block_;
+    if (!current_block_) {
+      return;
+    }
+
+    // Find the first block by traversing backwards
+    Block *first_block = current_block_;
+    while (first_block && first_block->prev) {
+      first_block = first_block->prev;
+    }
+
+    // Free all blocks after the first one
+    Block *current = current_block_;
+    while (current && current != first_block) {
+      Block *prev = current->prev;
+      std::free(current);
+      current = prev;
+    }
+
+    // Update first block's counters to reflect only itself
+    if (first_block) {
+      first_block->total_size = first_block->size;
+      first_block->block_count = 1;
+    }
+
+    current_block_ = first_block;
     current_offset_ = 0;
   }
 
   size_t total_allocated() const {
-    size_t total = 0;
-    Block *current = first_block_;
-    while (current) {
-      total += current->size;
-      current = current->next;
-    }
-    return total;
+    return current_block_ ? current_block_->total_size : 0;
   }
 
   size_t used_bytes() const {
-    size_t total = current_offset_;
-    Block *current = first_block_;
-    while (current && current != current_block_) {
-      total += current->size;
-      current = current->next;
+    if (!current_block_) {
+      return 0;
     }
-    return total;
+    size_t prev_total =
+        current_block_->prev ? current_block_->prev->total_size : 0;
+    return prev_total + current_offset_;
   }
 
   size_t available_in_current_block() const {
@@ -135,29 +153,12 @@ public:
   }
 
   size_t num_blocks() const {
-    size_t count = 0;
-    Block *current = first_block_;
-    while (current) {
-      count++;
-      current = current->next;
-    }
-    return count;
+    return current_block_ ? current_block_->block_count : 0;
   }
 
 private:
   void add_block(size_t size) {
-    Block *new_block = Block::create(size);
-
-    if (!first_block_) {
-      first_block_ = new_block;
-    } else {
-      Block *current = first_block_;
-      while (current->next) {
-        current = current->next;
-      }
-      current->next = new_block;
-    }
-
+    Block *new_block = Block::create(size, current_block_);
     current_block_ = new_block;
     current_offset_ = 0;
   }
@@ -178,7 +179,6 @@ private:
   }
 
   size_t initial_block_size_;
-  Block *first_block_;
   Block *current_block_;
   size_t current_offset_;
 };