Added stack allocator snapshots.

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<void*, Result<Stacktrace>> 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<std::size_t chunkSize, template<typename> typename TBacking> requires (allocator_tmpl<TBacking>)
+    friend class StackAllocator;
+};
+
+template<typename TStackAllocator>
+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<std::size_t chunkSize = 4096, template<typename> typename TBacking = MIJIN_DEFAULT_ALLOCATOR> requires (allocator_tmpl<TBacking>)
 class StackAllocator
 {
@@ -242,6 +292,154 @@ public:
     {
         return stl_allocator_t<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<StackAllocatorSnapshotData*>(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<const void*>(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<const std::byte*>(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<typename TValue, typename TStackAllocator>
     friend class StlStackAllocator;
 };