From 86f3790ce1b6a7ddaee494df90061780114e621a Mon Sep 17 00:00:00 2001 From: Patrick Wuttke Date: Thu, 26 Jun 2025 16:55:33 +0200 Subject: [PATCH] Added stack allocator snapshots. --- source/mijin/memory/stack_allocator.hpp | 221 ++++++++++++++++++++++++ 1 file changed, 221 insertions(+) diff --git a/source/mijin/memory/stack_allocator.hpp b/source/mijin/memory/stack_allocator.hpp index 84e65d8..1a73294 100644 --- a/source/mijin/memory/stack_allocator.hpp +++ b/source/mijin/memory/stack_allocator.hpp @@ -80,6 +80,56 @@ public: friend class StlStackAllocator; }; +namespace impl +{ +struct StackAllocatorSnapshotData +{ + struct ChunkSnapshot + { + std::size_t allocated; + }; + std::size_t numChunks; +#if MIJIN_STACK_ALLOCATOR_DEBUG == 1 + std::size_t numAllocations_ = 0; +#elif MIJIN_STACK_ALLOCATOR_DEBUG > 1 + // just for debugging, so we don't care what memory this uses... + std::unordered_map> activeAllocations_; +#endif + ChunkSnapshot chunks[1]; // may be bigger than that +}; +} + +class StackAllocatorSnapshot +{ +private: + impl::StackAllocatorSnapshotData* data = nullptr; + + impl::StackAllocatorSnapshotData* operator->() const MIJIN_NOEXCEPT { return data; } + + + template typename TBacking> requires (allocator_tmpl) + friend class StackAllocator; +}; + +template +class StackAllocatorScope +{ +private: + TStackAllocator* allocator_; + StackAllocatorSnapshot snapshot_; +public: + explicit StackAllocatorScope(TStackAllocator& allocator) : allocator_(&allocator), snapshot_(allocator_->createSnapshot()) {} + StackAllocatorScope(const StackAllocatorScope&) = delete; + StackAllocatorScope(StackAllocatorScope&&) = delete; + ~StackAllocatorScope() MIJIN_NOEXCEPT + { + allocator_->restoreSnapshot(snapshot_); + } + + StackAllocatorScope& operator=(const StackAllocatorScope&) = delete; + StackAllocatorScope& operator=(StackAllocatorScope&&) = delete; +}; + template typename TBacking = MIJIN_DEFAULT_ALLOCATOR> requires (allocator_tmpl) class StackAllocator { @@ -242,6 +292,154 @@ public: { return stl_allocator_t(*this); } + + [[nodiscard]] + std::size_t getNumChunks() const MIJIN_NOEXCEPT + { + std::size_t num = 0; + for (Chunk* chunk = firstChunk_; chunk != nullptr; chunk = chunk->next) + { + ++num; + } + return num; + } + + [[nodiscard]] + StackAllocatorSnapshot createSnapshot() + { + if (firstChunk_ == nullptr) + { + return {}; + } + + using impl::StackAllocatorSnapshotData; + + std::size_t numChunks = getNumChunks(); + std::size_t snapshotSize = calcSnapshotSize(numChunks); + Chunk* prevFirst = firstChunk_; + StackAllocatorSnapshotData* snapshotData = static_cast(allocate(alignof(StackAllocatorSnapshotData), snapshotSize)); + if (snapshotData == nullptr) + { + // couldn't allocate the snapshot + return {}; + } + StackAllocatorSnapshot snapshot; + snapshot.data = snapshotData; + if (firstChunk_ != prevFirst) + { + // new chunk has been added, adjust the snapshot + // the snapshot must be inside the new chunk (but not necessarily at the very beginning, due to alignment) + MIJIN_ASSERT(static_cast(snapshot.data) == alignUp(firstChunk_->data.data(), alignof(StackAllocatorSnapshot)), "Snapshot not where it was expected."); + + // a chunk might be too small for the snapshot if we grow it (unlikely, but not impossible) + if (ACTUAL_CHUNK_SIZE - firstChunk_->allocated < MIJIN_STRIDEOF(StackAllocatorSnapshotData::ChunkSnapshot)) [[unlikely]] + { + firstChunk_->allocated = 0; + return {}; + } + + // looking good, adjust the numbers + ++numChunks; + snapshotSize += MIJIN_STRIDEOF(StackAllocatorSnapshotData::ChunkSnapshot); + firstChunk_->allocated += MIJIN_STRIDEOF(StackAllocatorSnapshotData::ChunkSnapshot); + } + // now fill out the struct + snapshot->numChunks = numChunks; +#if MIJIN_STACK_ALLOCATOR_DEBUG == 1 + snapshot->numAllocations_ = numAllocations_; +#elif MIJIN_STACK_ALLOCATOR_DEBUG > 1 + // just for debugging, so we don't care what memory this uses... + snapshot->activeAllocations_ = activeAllocations_; +#endif + + std::size_t pos = 0; + for (Chunk* chunk = firstChunk_; chunk != nullptr; chunk = chunk->next, ++pos) + { + snapshot->chunks[pos].allocated = chunk->allocated; + } + return snapshot; + } + + void restoreSnapshot(StackAllocatorSnapshot snapshot) MIJIN_NOEXCEPT + { + if (snapshot.data == nullptr) + { + return; + } + + const std::size_t numChunks = getNumChunks(); + MIJIN_ASSERT_FATAL(snapshot->numChunks <= numChunks, "Snapshot contains more chunks than the allocator!"); + +#if MIJIN_STACK_ALLOCATOR_DEBUG == 1 + MIJIN_ASSERT(snapshot->numAllocations_ >= numAllocations_, "Missing deallocation in StackAllocator!"); +#elif MIJIN_STACK_ALLOCATOR_DEBUG > 1 + // TODO: compare and print changes + unsigned numMismatches = 0; + for (const auto& [ptr, stack] : activeAllocations_) + { + if (snapshot->activeAllocations_.contains(ptr)) + { + continue; + } + ++numMismatches; + + if (stack.isError()) + { + std::println(stderr, "Missing deallocation at {}, no stacktrace ({})", ptr, stack.getError().message); + } + else + { + std::println(stderr, "Missing deallocation at 0x{}:\n{}", ptr, stack.getValue()); + } + } +#if 0 // deallocating more than expected shouldn't be a problem + for (const auto& [ptr, stack] : snapshot->activeAllocations_) + { + if (activeAllocations_.contains(ptr)) + { + continue; + } + ++numMismatches; + + if (stack.isError()) + { + std::println(stderr, "Unexpected deallocation at {}, no stacktrace ({})", ptr, stack.getError().message); + } + else + { + std::println(stderr, "Unexpected deallocation at 0x{}:\n{}", ptr, stack.getValue()); + } + } +#endif + if (numMismatches > 0) + { + std::println(stderr, "{} mismatched deallocations when restoring stack allocator snapshot.", numMismatches); + MIJIN_TRAP(); + } +#endif + // if we allocated new chunks since the snapshot, these are completely empty now + const std::size_t emptyChunks = numChunks - snapshot->numChunks; + + Chunk* chunk = firstChunk_; + for (std::size_t idx = 0; idx < emptyChunks; ++idx) + { + chunk->allocated = 0; + chunk = chunk->next; + } + + // the other values are in the snapshot + for (std::size_t idx = 0; idx < snapshot->numChunks; ++idx) + { + chunk->allocated = snapshot->chunks[idx].allocated; + chunk = chunk->next; + } + MIJIN_ASSERT(chunk == nullptr, "Something didn't add up."); + + // finally free the space for the snapshot itself + Chunk* snapshotChunk = findChunk(snapshot.data); + MIJIN_ASSERT_FATAL(snapshotChunk != nullptr, "Snapshot not in chunks?"); + snapshotChunk->allocated -= calcSnapshotSize(snapshot->numChunks); // note: this might miss the alignment bytes of the snapshot, but that should be fine + } private: void initAndAddChunk(Chunk* newChunk) noexcept { @@ -252,6 +450,29 @@ private: firstChunk_ = newChunk; } + bool isInChunk(const void* address, const Chunk& chunk) const MIJIN_NOEXCEPT + { + const std::byte* asByte = static_cast(address); + return asByte >= chunk.data.data() && asByte < chunk.data.data() + ACTUAL_CHUNK_SIZE; + } + + Chunk* findChunk(const void* address) const MIJIN_NOEXCEPT + { + for (Chunk* chunk = firstChunk_; chunk != nullptr; chunk = chunk->next) + { + if (isInChunk(address, *chunk)) + { + return chunk; + } + } + return nullptr; + } + + static std::size_t calcSnapshotSize(std::size_t numChunks) MIJIN_NOEXCEPT + { + return sizeof(impl::StackAllocatorSnapshotData) + ((numChunks - 1) * MIJIN_STRIDEOF(impl::StackAllocatorSnapshotData::ChunkSnapshot)); + } + template friend class StlStackAllocator; };