mijin2/source/mijin/memory/stack_allocator.hpp

485 lines
16 KiB
C++

#pragma once
#if !defined(MIJIN_MEMORY_STACK_ALLOCATOR_HPP_INCLUDED)
#define MIJIN_MEMORY_STACK_ALLOCATOR_HPP_INCLUDED 1
#include <memory>
#include <utility>
#include "../debug/assert.hpp"
#include "../internal/common.hpp"
#include "../util/align.hpp"
#include "../util/concepts.hpp"
#include "../util/traits.hpp"
#if !defined(MIJIN_STACK_ALLOCATOR_DEBUG)
#if defined(MIJIN_DEBUG)
#define MIJIN_STACK_ALLOCATOR_DEBUG 1
#else
#define MIJIN_STACK_ALLOCATOR_DEBUG 0
#endif
#endif // !defined(MIJIN_STACK_ALLOCATOR_DEBUG)
#if MIJIN_STACK_ALLOCATOR_DEBUG > 1
#include <print>
#include <unordered_map>
#include "../debug/stacktrace.hpp"
#endif
namespace mijin
{
template<typename TValue, typename TStackAllocator>
class StlStackAllocator
{
public:
using value_type = TValue;
private:
TStackAllocator* base_;
public:
explicit StlStackAllocator(TStackAllocator& base) MIJIN_NOEXCEPT : base_(&base) {}
template<typename TOtherValue>
StlStackAllocator(const StlStackAllocator<TOtherValue, TStackAllocator>& other) MIJIN_NOEXCEPT : base_(other.base_) {}
template<typename TOtherValue>
StlStackAllocator& operator=(const StlStackAllocator<TOtherValue, TStackAllocator>& other) MIJIN_NOEXCEPT
{
base_ = other.base_;
return *this;
}
auto operator<=>(const StlStackAllocator&) const noexcept = default;
[[nodiscard]]
TValue* allocate(std::size_t count) const
{
void* result = base_->allocate(alignof(TValue), count * sizeof(TValue));
#if MIJIN_STACK_ALLOCATOR_DEBUG == 1
++base_->numAllocations_;
#elif MIJIN_STACK_ALLOCATOR_DEBUG > 1
base_->activeAllocations_.emplace(result, captureStacktrace(1));
#endif
return static_cast<TValue*>(result);
}
void deallocate([[maybe_unused]] TValue* ptr, std::size_t /* count */) const MIJIN_NOEXCEPT
{
#if MIJIN_STACK_ALLOCATOR_DEBUG == 1
MIJIN_ASSERT(base_->numAllocations_ > 0, "Unbalanced allocations in stack allocators!");
--base_->numAllocations_;
#elif MIJIN_STACK_ALLOCATOR_DEBUG > 1
auto it = base_->activeAllocations_.find(ptr);
if (it != base_->activeAllocations_.end())
{
base_->activeAllocations_.erase(it);
}
else
{
MIJIN_ERROR("Deallocating invalid pointer from StackAllocator.");
}
#endif
}
template<typename TOtherValue, typename TOtherAllocator>
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
{
public:
using backing_t = TBacking<void>;
static constexpr std::size_t ACTUAL_CHUNK_SIZE = chunkSize - sizeof(void*) - sizeof(std::size_t);
template<typename T>
using stl_allocator_t = StlStackAllocator<T, StackAllocator<chunkSize, TBacking>>;
private:
struct Chunk
{
std::array<std::byte, ACTUAL_CHUNK_SIZE> data;
Chunk* next;
std::size_t allocated;
};
[[no_unique_address]] TBacking<Chunk> backing_;
Chunk* firstChunk_ = nullptr;
#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
public:
StackAllocator() MIJIN_NOEXCEPT_IF(std::is_nothrow_default_constructible_v<backing_t>) = default;
explicit StackAllocator(backing_t backing) MIJIN_NOEXCEPT_IF((std::is_nothrow_constructible_v<TBacking<Chunk>, backing_t&&>))
: backing_(std::move(backing)) {}
StackAllocator(const StackAllocator&) = delete;
StackAllocator(StackAllocator&& other) MIJIN_NOEXCEPT_IF(std::is_nothrow_move_constructible_v<TBacking<Chunk>>)
: firstChunk_(std::exchange(other.firstChunk_, nullptr)) {}
~StackAllocator() noexcept
{
Chunk* chunk = firstChunk_;
while (chunk != nullptr)
{
Chunk* nextChunk = firstChunk_->next;
backing_.deallocate(chunk, 1);
chunk = nextChunk;
}
}
StackAllocator& operator=(const StackAllocator&) = delete;
StackAllocator& operator=(StackAllocator&& other) MIJIN_NOEXCEPT_IF(std::is_nothrow_move_assignable_v<TBacking<Chunk>>)
{
if (this != &other)
{
backing_ = std::move(other.backing_);
firstChunk_ = std::exchange(other.firstChunk_, nullptr);
}
return *this;
}
void* allocate(std::size_t alignment, std::size_t size)
{
// first check if this can ever fit
if (size > ACTUAL_CHUNK_SIZE)
{
return nullptr;
}
// then try to find space in the current chunks
for (Chunk* chunk = firstChunk_; chunk != nullptr; chunk = chunk->next)
{
const std::size_t remaining = ACTUAL_CHUNK_SIZE - chunk->allocated;
if (remaining < size)
{
continue;
}
std::byte* start = &chunk->data[chunk->allocated];
std::byte* pos = mijin::alignUp(start, alignment);
const std::ptrdiff_t alignmentBytes = pos - start;
const std::size_t combinedSize = size + alignmentBytes;
if (remaining < combinedSize)
{
continue;
}
chunk->allocated += combinedSize;
return pos;
}
// no free space in any chunk? allocate a new one
Chunk* newChunk = backing_.allocate(1);
if (newChunk == nullptr)
{
return nullptr;
}
initAndAddChunk(newChunk);
// now try with the new chunk
std::byte* start = newChunk->data.data();
std::byte* pos = mijin::alignUp(start, alignment);
const std::ptrdiff_t alignmentBytes = pos - start;
const std::size_t combinedSize = size + alignmentBytes;
// doesn't fit (due to alignment), time to give up
if (ACTUAL_CHUNK_SIZE < combinedSize)
{
return nullptr;
}
newChunk->allocated = combinedSize;
return pos;
}
void reset() noexcept
{
#if MIJIN_STACK_ALLOCATOR_DEBUG == 1
MIJIN_ASSERT(numAllocations_ == 0, "Missing deallocation in StackAllocator!");
#elif MIJIN_STACK_ALLOCATOR_DEBUG > 1
if (!activeAllocations_.empty())
{
std::println(stderr, "{} active allocations in StackAllocator when resetting!", activeAllocations_.size());
for (const auto& [ptr, stack] : activeAllocations_)
{
if (stack.isError())
{
std::println(stderr, "at {}, no stacktrace ({})", ptr, stack.getError().message);
}
else
{
std::println(stderr, "at 0x{}:\n{}", ptr, stack.getValue());
}
}
MIJIN_TRAP();
}
#endif
for (Chunk* chunk = firstChunk_; chunk != nullptr; chunk = chunk->next)
{
chunk->allocated = 0;
}
}
bool createChunks(std::size_t count)
{
if (count == 0)
{
return true;
}
Chunk* newChunks = backing_.allocate(count);
if (newChunks == nullptr)
{
return false;
}
// reverse so the chunks are chained from 0 to count (new chunks are inserted in front)
for (std::size_t pos = count; pos > 0; --pos)
{
initAndAddChunk(&newChunks[pos-1]);
}
return true;
}
template<typename T>
stl_allocator_t<T> makeStlAllocator() MIJIN_NOEXCEPT
{
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
{
::new (newChunk) Chunk();
// put it in the front
newChunk->next = firstChunk_;
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;
};
} // namespace mijin
#endif // !defined(MIJIN_MEMORY_STACK_ALLOCATOR_HPP_INCLUDED)