diff --git a/src/arena_allocator.hpp b/src/arena_allocator.hpp index 83451e4..4b55649 100644 --- a/src/arena_allocator.hpp +++ b/src/arena_allocator.hpp @@ -1,41 +1,45 @@ #pragma once +#include #include +#include #include -#include #include +#include 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(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_; }; diff --git a/tests/test_arena_allocator.cpp b/tests/test_arena_allocator.cpp index e62622c..73a1c3a 100644 --- a/tests/test_arena_allocator.cpp +++ b/tests/test_arena_allocator.cpp @@ -6,17 +6,17 @@ TEST_CASE("ArenaAllocator basic construction") { ArenaAllocator arena; - CHECK(arena.num_blocks() == 1); + CHECK(arena.num_blocks() == 0); CHECK(arena.used_bytes() == 0); - CHECK(arena.total_allocated() == 1024); - CHECK(arena.available_in_current_block() == 1024); + CHECK(arena.total_allocated() == 0); + CHECK(arena.available_in_current_block() == 0); } TEST_CASE("ArenaAllocator custom initial size") { ArenaAllocator arena(2048); - CHECK(arena.num_blocks() == 1); - CHECK(arena.total_allocated() == 2048); - CHECK(arena.available_in_current_block() == 2048); + CHECK(arena.num_blocks() == 0); + CHECK(arena.total_allocated() == 0); + CHECK(arena.available_in_current_block() == 0); } TEST_CASE("ArenaAllocator basic allocation") { @@ -93,7 +93,7 @@ TEST_CASE("ArenaAllocator block management") { SUBCASE("allocation larger than block size grows arena") { void *ptr = arena.allocate(200); CHECK(ptr != nullptr); - CHECK(arena.num_blocks() == 2); + CHECK(arena.num_blocks() == 1); } } @@ -146,12 +146,43 @@ TEST_CASE("ArenaAllocator reset functionality") { CHECK(arena.used_bytes() == 50); } +TEST_CASE("ArenaAllocator reset memory leak test") { + ArenaAllocator arena(32); // Smaller initial size + + // Force multiple blocks + arena.allocate(30); // First block (32 bytes) + CHECK(arena.num_blocks() == 1); + + arena.allocate(30); // Should create second block (64 bytes due to doubling) + CHECK(arena.num_blocks() == 2); + + arena.allocate(100); // Should create third block (128 bytes due to doubling, + // or larger for 100) + CHECK(arena.num_blocks() == 3); + + size_t total_before = arena.total_allocated(); + CHECK(total_before > 32); + + arena.reset(); + + // After reset, only first block should remain (others freed to prevent memory + // leak) + CHECK(arena.num_blocks() == 1); + CHECK(arena.total_allocated() == 32); // Only first block size + CHECK(arena.used_bytes() == 0); + + // Should be able to use the first block again + void *ptr = arena.allocate(20); + CHECK(ptr != nullptr); + CHECK(arena.used_bytes() == 20); +} + TEST_CASE("ArenaAllocator memory tracking") { ArenaAllocator arena(512); - CHECK(arena.total_allocated() == 512); + CHECK(arena.total_allocated() == 0); CHECK(arena.used_bytes() == 0); - CHECK(arena.available_in_current_block() == 512); + CHECK(arena.available_in_current_block() == 0); arena.allocate(100); CHECK(arena.used_bytes() >= 100);