diff --git a/lib/nw/util/memory.cpp b/lib/nw/util/memory.cpp index 5aa064861..f06c78a72 100644 --- a/lib/nw/util/memory.cpp +++ b/lib/nw/util/memory.cpp @@ -97,4 +97,54 @@ bool MemoryArena::do_is_equal(const std::pmr::memory_resource& other) const noex return this == &other; } +// == MemoryScope ============================================================== +// ============================================================================= + +MemoryScope::MemoryScope(MemoryArena* arena) + : arena_{arena} + , marker_{arena->current()} +{ +} + +MemoryScope::MemoryScope(MemoryScope&& other) +{ + this->arena_ = other.arena_; + this->finalizers_ = other.finalizers_; + this->marker_ = other.marker_; + + other.arena_ = nullptr; + other.finalizers_ = nullptr; + other.marker_ = MemoryMarker{}; +} + +MemoryScope& MemoryScope::operator=(MemoryScope&& other) +{ + if (this != &other) { + this->arena_ = other.arena_; + this->finalizers_ = other.finalizers_; + this->marker_ = other.marker_; + + other.arena_ = nullptr; + other.finalizers_ = nullptr; + other.marker_ = MemoryMarker{}; + } + return *this; +} + +MemoryScope::~MemoryScope() +{ + auto f = finalizers_; + while (f) { + auto obj = reinterpret_cast(f) + sizeof(detail::Finalizer); + f->fn(obj); + f = f->next; + } + arena_->rewind(marker_); +} + +void* MemoryScope::alloc(size_t size, size_t alignment) +{ + return arena_->allocate(size, alignment); +} + } // namespace nw diff --git a/lib/nw/util/memory.hpp b/lib/nw/util/memory.hpp index 5042a48de..56fd8b3c9 100644 --- a/lib/nw/util/memory.hpp +++ b/lib/nw/util/memory.hpp @@ -132,4 +132,64 @@ struct ObjectPool { std::vector chunks_; }; +namespace detail { + +template +void destructor(void* ptr) +{ + static_cast(ptr)->~T(); +} + +struct Finalizer { + void (*fn)(void*) = nullptr; + Finalizer* next = nullptr; +}; + +template +struct FinalizedObject { + Finalizer f; + T obj; +}; + +} + +struct MemoryScope { + MemoryScope(MemoryArena* arena); + MemoryScope(const MemoryScope&) = delete; + MemoryScope(MemoryScope&& other); + ~MemoryScope(); + + MemoryScope& operator=(const MemoryScope&) = delete; + MemoryScope& operator=(MemoryScope&& other); + + /// Allocates ``size`` bytes with ``alignment`` + void* alloc(size_t size, size_t alignment = alignof(max_align_t)); + + /// Allocates a non-trivial object and stores pointer to destructor that is run when scope exits. + template + T* alloc_obj(Args&&... args) + { + static_assert(!(std::is_standard_layout_v && std::is_trivial_v), "Use alloc_pod for POD types"); + void* mem = alloc(sizeof(detail::FinalizedObject), alignof(detail::FinalizedObject)); + auto fo = static_cast*>(mem); + fo->f.fn = &detail::destructor; + fo->f.next = finalizers_; // last in, first destructed. + finalizers_ = &fo->f; + return new (&fo->obj) T(std::forward(args)...); + } + + /// Allocates a trivial object no destructor is run on scope exit. + template + T* alloc_pod() + { + static_assert(std::is_standard_layout_v && std::is_trivial_v, "Use alloc_obj for non-trivial types"); + void* mem = alloc(sizeof(T), alignof(T)); + return new (mem) T(); + } + + MemoryArena* arena_ = nullptr; + MemoryMarker marker_; + detail::Finalizer* finalizers_ = nullptr; +}; + } // namespace nw diff --git a/tests/util_memory.cpp b/tests/util_memory.cpp index 781ff9788..d4f9e2503 100644 --- a/tests/util_memory.cpp +++ b/tests/util_memory.cpp @@ -18,3 +18,53 @@ TEST(Memory, Allocator) arena_string = "Hello, World"; EXPECT_EQ(arena_string, "Hello, World"); } + +struct TestStruct { + TestStruct(bool* am_i_destrcuted) + : am_i_destrcuted_{am_i_destrcuted} + { + } + ~TestStruct() + { + *am_i_destrcuted_ = true; + } + + bool* am_i_destrcuted_; +}; + +struct TestPOD { + bool value; + int test; +}; + +TEST(Memory, Scope) +{ + nw::MemoryArena arena; + auto current = arena.current(); + + { + nw::MemoryScope scope(&arena); + scope.alloc(60); + EXPECT_NE(current, arena.current()); + } + EXPECT_EQ(current, arena.current()); + + bool result = false; + { + nw::MemoryScope scope(&arena); + auto value = scope.alloc_obj(&result); + EXPECT_TRUE(!!scope.finalizers_); + EXPECT_EQ((uint8_t*)scope.finalizers_ + sizeof(nw::detail::Finalizer), (uint8_t*)value); + EXPECT_EQ(&result, value->am_i_destrcuted_); + EXPECT_NE(current, arena.current()); + } + EXPECT_TRUE(result); + EXPECT_EQ(current, arena.current()); + + { + nw::MemoryScope scope(&arena); + scope.alloc_pod(); + EXPECT_NE(current, arena.current()); + } + EXPECT_EQ(current, arena.current()); +}