diff --git a/include/nd-wfc/wfc_allocator.hpp b/include/nd-wfc/wfc_allocator.hpp index 1ddb561..f28fdc7 100644 --- a/include/nd-wfc/wfc_allocator.hpp +++ b/include/nd-wfc/wfc_allocator.hpp @@ -63,13 +63,23 @@ namespace WFC { */ class WFCStackAllocator { private: + struct MemoryPool { + void* ptr; + size_t size; + size_t used; + + MemoryPool() : ptr(nullptr), size(0), used(0) {} + MemoryPool(void* p, size_t s) : ptr(p), size(s), used(0) {} + }; + struct Block { void* ptr; size_t size; size_t alignment; - - Block() : ptr(nullptr), size(0), alignment(0) {} - Block(void* p, size_t s, size_t a) : ptr(p), size(s), alignment(a) {} + size_t poolIndex; // Which pool this allocation came from + + Block() : ptr(nullptr), size(0), alignment(0), poolIndex(0) {} + Block(void* p, size_t s, size_t a, size_t pi) : ptr(p), size(s), alignment(a), poolIndex(pi) {} }; public: @@ -78,19 +88,15 @@ public: * @param initialCapacity Initial memory pool size in bytes */ explicit WFCStackAllocator(size_t initialCapacity = 1024 * 1024) // 1MB default - : m_capacity(initialCapacity) - , m_used(0) - , m_pool(nullptr) { - m_pool = allocate_aligned_memory(64, m_capacity); // 64-byte alignment for cache efficiency - if (!m_pool) { - throw std::bad_alloc(); - } + addPool(initialCapacity); } ~WFCStackAllocator() { - if (m_pool) { - free_aligned_memory(m_pool); + for (auto& pool : m_pools) { + if (pool.ptr) { + free_aligned_memory(pool.ptr); + } } } @@ -107,21 +113,37 @@ public: * @return Pointer to allocated memory */ void* allocate(size_t size, size_t alignment = 8) { - // Align the current position - size_t alignedUsed = alignUp(m_used, alignment); - - // Check if we have enough space - if (alignedUsed + size > m_capacity) { - // Grow the pool if needed - growPool(alignedUsed + size); + // Try to allocate from existing pools + for (size_t i = 0; i < m_pools.size(); ++i) { + auto& pool = m_pools[i]; + + // Align the current position in this pool + size_t alignedUsed = alignUp(pool.used, alignment); + + // Check if we have enough space in this pool + if (alignedUsed + size <= pool.size) { + void* ptr = static_cast(pool.ptr) + alignedUsed; + pool.used = alignedUsed + size; + + // Track this allocation for proper cleanup + m_allocations.emplace_back(ptr, size, alignment, i); + + return ptr; + } } - - void* ptr = static_cast(m_pool) + alignedUsed; - m_used = alignedUsed + size; - + + // No existing pool has enough space, add a new one + size_t newPoolSize = std::max(m_pools.back().size * 2, size * 2); // Grow exponentially + addPool(newPoolSize); + + // Now allocate from the new pool (which is the last one) + auto& newPool = m_pools.back(); + void* ptr = static_cast(newPool.ptr); + newPool.used = size; + // Track this allocation for proper cleanup - m_allocations.emplace_back(ptr, size, alignment); - + m_allocations.emplace_back(ptr, size, alignment, m_pools.size() - 1); + return ptr; } @@ -141,19 +163,39 @@ public: class StackFrame { private: WFCStackAllocator& m_allocator; - size_t m_savedUsed; + std::vector m_savedUsed; size_t m_savedAllocCount; public: - StackFrame(WFCStackAllocator& allocator) + StackFrame(WFCStackAllocator& allocator) : m_allocator(allocator) - , m_savedUsed(allocator.m_used) , m_savedAllocCount(allocator.m_allocations.size()) - {} + { + // Save the current used state of all pools + m_savedUsed.reserve(allocator.m_pools.size()); + for (const auto& pool : allocator.m_pools) { + m_savedUsed.push_back(pool.used); + } + } ~StackFrame() { - // Restore the stack to the previous state - m_allocator.m_used = m_savedUsed; + // Restore the used state of all pools + for (size_t i = 0; i < m_savedUsed.size() && i < m_allocator.m_pools.size(); ++i) { + m_allocator.m_pools[i].used = m_savedUsed[i]; + } + + // Remove any new pools that were added during this frame + if (m_allocator.m_pools.size() > m_savedUsed.size()) { + // Free the additional pools that were added + for (size_t i = m_savedUsed.size(); i < m_allocator.m_pools.size(); ++i) { + if (m_allocator.m_pools[i].ptr) { + free_aligned_memory(m_allocator.m_pools[i].ptr); + } + } + m_allocator.m_pools.resize(m_savedUsed.size()); + } + + // Remove allocations that were made during this frame m_allocator.m_allocations.resize(m_savedAllocCount); } @@ -174,12 +216,24 @@ public: /** * @brief Get current memory usage */ - size_t getUsed() const { return m_used; } + size_t getUsed() const { + size_t total = 0; + for (const auto& pool : m_pools) { + total += pool.used; + } + return total; + } /** * @brief Get total capacity */ - size_t getCapacity() const { return m_capacity; } + size_t getCapacity() const { + size_t total = 0; + for (const auto& pool : m_pools) { + total += pool.size; + } + return total; + } /** * @brief Get allocation count @@ -190,41 +244,27 @@ public: * @brief Reset the allocator (useful for reusing between WFC runs) */ void reset() { - m_used = 0; + for (auto& pool : m_pools) { + pool.used = 0; + } m_allocations.clear(); } private: + void addPool(size_t size) { + void* ptr = allocate_aligned_memory(64, size); // 64-byte alignment for cache efficiency + if (!ptr) { + throw std::bad_alloc(); + } + m_pools.emplace_back(ptr, size); + } + size_t alignUp(size_t value, size_t alignment) const { return (value + alignment - 1) & ~(alignment - 1); } - void growPool(size_t requiredSize) { - size_t newCapacity = std::max(m_capacity * 2, requiredSize); - - void* newPool = allocate_aligned_memory(64, newCapacity); - if (!newPool) { - throw std::bad_alloc(); - } - - // Copy existing data - std::memcpy(newPool, m_pool, m_used); - - // Update all existing pointers in allocations - ptrdiff_t offset = static_cast(newPool) - static_cast(m_pool); - for (auto& block : m_allocations) { - block.ptr = static_cast(block.ptr) + offset; - } - - free_aligned_memory(m_pool); - m_pool = newPool; - m_capacity = newCapacity; - } - private: - size_t m_capacity; - size_t m_used; - void* m_pool; + std::vector m_pools; std::vector m_allocations; };