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;
   }
-}
+}

src/arena_allocator.hpp

@@ -208,6 +208,50 @@ public:
     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 (default: alignof(std::max_align_t))
+   * @return Pointer to the reallocated memory (may be the same as ptr or
+   * different)
+   * @throws std::bad_alloc 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
+   * - Otherwise allocates new space and copies min(old_size, new_size) bytes
+   *
+   * ## Performance:
+   * - In-place extension: O(1) - just updates offset
+   * - Copy required: O(min(old_size, new_size)) for the memcpy
+   *
+   * ## Example:
+   * ```cpp
+   * void* ptr = arena.allocate(100);
+   * // ... use ptr ...
+   * ptr = arena.realloc(ptr, 100, 200);  // May extend in place or copy
+   * ```
+   *
+   * ## 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(void *ptr, size_t old_size, size_t new_size,
+                size_t alignment = alignof(std::max_align_t));
+
   /**
    * @brief Construct an object of type T in the arena using placement new.
    *

tests/test_arena_allocator.cpp

@@ -356,3 +356,179 @@ TEST_CASE("ArenaAllocator alignment edge cases") {
     CHECK(reinterpret_cast<uintptr_t>(ptr) % 128 == 0);
   }
 }
+
+TEST_CASE("ArenaAllocator realloc functionality") {
+  ArenaAllocator arena;
+
+  SUBCASE("realloc edge cases") {
+    // realloc with new_size == 0 returns nullptr and reclaims memory if it's
+    // the last allocation
+    ArenaAllocator fresh_arena(256);
+    void *ptr = fresh_arena.allocate(100);
+    size_t used_before = fresh_arena.used_bytes();
+    CHECK(used_before == 100);
+
+    void *result = fresh_arena.realloc(ptr, 100, 0);
+    CHECK(result == nullptr);
+    CHECK(fresh_arena.used_bytes() == 0); // Memory should be reclaimed
+
+    // Test case where it's NOT the last allocation - memory cannot be reclaimed
+    ArenaAllocator arena2(256);
+    void *ptr1 = arena2.allocate(50);
+    void *ptr2 = arena2.allocate(50);
+    size_t used_before2 = arena2.used_bytes();
+    CHECK(used_before2 >= 100); // At least 100 bytes due to potential alignment
+
+    void *result2 =
+        arena2.realloc(ptr1, 50, 0); // ptr1 is not the last allocation
+    CHECK(result2 == nullptr);
+    CHECK(arena2.used_bytes() ==
+          used_before2); // Memory should NOT be reclaimed
+
+    // realloc with nullptr behaves like allocate
+    void *new_ptr = arena.realloc(nullptr, 0, 150);
+    CHECK(new_ptr != nullptr);
+    CHECK(arena.used_bytes() >= 150);
+
+    // realloc with same size returns same pointer
+    void *same_ptr = arena.realloc(new_ptr, 150, 150);
+    CHECK(same_ptr == new_ptr);
+  }
+
+  SUBCASE("in-place extension - growing") {
+    ArenaAllocator fresh_arena(1024);
+    void *ptr = fresh_arena.allocate(100);
+    CHECK(ptr != nullptr);
+
+    // Fill the allocation with test data
+    std::memset(ptr, 0xAB, 100);
+
+    // Should extend in place since it's the last allocation
+    void *extended_ptr = fresh_arena.realloc(ptr, 100, 200);
+    CHECK(extended_ptr == ptr); // Should be same pointer (in-place)
+
+    // Check that original data is preserved
+    char *data = static_cast<char *>(extended_ptr);
+    for (size_t i = 0; i < 100; ++i) {
+      CHECK(data[i] == static_cast<char>(0xAB));
+    }
+
+    CHECK(fresh_arena.used_bytes() == 200);
+  }
+
+  SUBCASE("in-place shrinking") {
+    ArenaAllocator fresh_arena(1024);
+    void *ptr = fresh_arena.allocate(200);
+    std::memset(ptr, 0xCD, 200);
+
+    // Should shrink in place
+    void *shrunk_ptr = fresh_arena.realloc(ptr, 200, 100);
+    CHECK(shrunk_ptr == ptr); // Same pointer
+    CHECK(fresh_arena.used_bytes() == 100);
+
+    // Check that remaining data is preserved
+    char *data = static_cast<char *>(shrunk_ptr);
+    for (size_t i = 0; i < 100; ++i) {
+      CHECK(data[i] == static_cast<char>(0xCD));
+    }
+  }
+
+  SUBCASE("copy when can't extend in place") {
+    ArenaAllocator fresh_arena(256); // Larger block to avoid edge cases
+
+    // Allocate first chunk
+    void *ptr1 = fresh_arena.allocate(60);
+    std::memset(ptr1, 0x11, 60);
+
+    // Allocate second chunk (this prevents in-place extension of ptr1)
+    void *ptr2 = fresh_arena.allocate(30);
+    std::memset(ptr2, 0x22, 30);
+
+    // Try to reallocate ptr1 - should copy since ptr2 is now the last
+    // allocation
+    void *realloc_ptr1 = fresh_arena.realloc(ptr1, 60, 120);
+    CHECK(realloc_ptr1 != ptr1); // Should be different pointer (copy occurred)
+
+    // Check that data was copied correctly
+    char *data = static_cast<char *>(realloc_ptr1);
+    for (size_t i = 0; i < 60; ++i) {
+      CHECK(data[i] == static_cast<char>(0x11));
+    }
+
+    // Try to reallocate ptr2 (last allocation) - should extend in place if
+    // space allows
+    void *realloc_ptr2 = fresh_arena.realloc(ptr2, 30, 50);
+    // Note: this might not be in-place if the realloc of ptr1 used up space in
+    // the block Let's just check the result is valid and data is preserved
+    CHECK(realloc_ptr2 != nullptr);
+
+    char *data2 = static_cast<char *>(realloc_ptr2);
+    for (size_t i = 0; i < 30; ++i) {
+      CHECK(data2[i] == static_cast<char>(0x22));
+    }
+  }
+
+  SUBCASE("copy when insufficient space for extension") {
+    ArenaAllocator fresh_arena(100);
+
+    // Allocate almost all space
+    void *ptr = fresh_arena.allocate(90);
+    std::memset(ptr, 0x33, 90);
+
+    // Try to extend beyond block size - should copy to new block
+    void *extended_ptr = fresh_arena.realloc(ptr, 90, 150);
+    CHECK(extended_ptr != ptr);           // Should be different (new block)
+    CHECK(fresh_arena.num_blocks() == 2); // Should have created new block
+
+    // Check data preservation
+    char *data = static_cast<char *>(extended_ptr);
+    for (size_t i = 0; i < 90; ++i) {
+      CHECK(data[i] == static_cast<char>(0x33));
+    }
+  }
+
+  SUBCASE("realloc with custom alignment") {
+    ArenaAllocator fresh_arena(1024);
+
+    // Allocate with specific alignment
+    void *ptr = fresh_arena.allocate(50, 16);
+    CHECK(reinterpret_cast<uintptr_t>(ptr) % 16 == 0);
+    std::memset(ptr, 0x44, 50);
+
+    // Realloc with same alignment - should extend in place
+    void *extended_ptr = fresh_arena.realloc(ptr, 50, 100, 16);
+    CHECK(extended_ptr == ptr); // In place
+    CHECK(reinterpret_cast<uintptr_t>(extended_ptr) % 16 == 0);
+
+    // Check data preservation
+    char *data = static_cast<char *>(extended_ptr);
+    for (size_t i = 0; i < 50; ++i) {
+      CHECK(data[i] == static_cast<char>(0x44));
+    }
+  }
+
+  SUBCASE("realloc stress test") {
+    ArenaAllocator fresh_arena(512);
+    void *ptr = fresh_arena.allocate(50);
+    size_t current_size = 50;
+
+    // Fill with pattern
+    for (size_t i = 0; i < 50; ++i) {
+      static_cast<char *>(ptr)[i] = static_cast<char>(i & 0xFF);
+    }
+
+    // Grow in multiple steps
+    for (size_t new_size = 100; new_size <= 500; new_size += 100) {
+      void *new_ptr = fresh_arena.realloc(ptr, current_size, new_size);
+      CHECK(new_ptr != nullptr);
+
+      // Verify original data is still there
+      for (size_t i = 0; i < 50; ++i) {
+        CHECK(static_cast<char *>(new_ptr)[i] == static_cast<char>(i & 0xFF));
+      }
+
+      ptr = new_ptr;
+      current_size = new_size;
+    }
+  }
+}