Add ArenaAllocator::realloc

src/arena_allocator.cpp

@@ -1,4 +1,5 @@
 #include "arena_allocator.hpp"
+#include <cassert>
 
 ArenaAllocator::~ArenaAllocator() {
   while (current_block_) {
@@ -63,6 +64,70 @@ void ArenaAllocator::reset() {
   current_offset_ = 0;
 }
 
+void *ArenaAllocator::realloc(void *ptr, size_t old_size, size_t new_size,
+                              size_t alignment) {
+  if (ptr == nullptr) {
+    return allocate(new_size, alignment);
+  }
+
+  if (new_size == old_size) {
+    return ptr;
+  }
+
+  // Assert that we have a current block if ptr is not null
+  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
+  // valid callers)
+  assert(current_offset_ >= old_size &&
+         "current_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;
+
+  if (ptr == expected_last_alloc_start) {
+    // This is indeed the last allocation
+    if (new_size > old_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) {
+        // We can extend in place
+        current_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;
+      return new_size == 0 ? nullptr : ptr;
+    }
+  }
+
+  // Can't extend in place
+  if (new_size == 0) {
+    // For non-last allocations, we can't reclaim memory but still return
+    // nullptr
+    return nullptr;
+  }
+
+  if (new_size <= old_size) {
+    // Shrinking - no need to allocate, just return the same pointer
+    return ptr;
+  }
+
+  // Growing but can't extend in place - need to allocate new space and copy
+  void *new_ptr = allocate(new_size, alignment);
+  if (new_ptr && ptr) {
+    // Copy all the old data since we're growing
+    std::memcpy(new_ptr, ptr, old_size);
+  }
+
+  return new_ptr;
+}
+
 std::vector<ArenaAllocator::PointerInfo>
 ArenaAllocator::find_intra_arena_pointers() const {
   std::vector<PointerInfo> pointers;
@@ -380,4 +445,4 @@ void ArenaAllocator::dump_memory_contents(std::ostream &out, const char *data,
     }
     out << "|" << std::dec << std::endl;
   }
-}
+}