[util] add: MemoryScope

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<uint8_t*>(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

lib/nw/util/memory.hpp

@@ -132,4 +132,64 @@ struct ObjectPool {
     std::vector<T*> chunks_;
 };
 
+namespace detail {
+
+template <typename T>
+void destructor(void* ptr)
+{
+    static_cast<T*>(ptr)->~T();
+}
+
+struct Finalizer {
+    void (*fn)(void*) = nullptr;
+    Finalizer* next = nullptr;
+};
+
+template <typename T>
+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 <typename T, typename... Args>
+    T* alloc_obj(Args&&... args)
+    {
+        static_assert(!(std::is_standard_layout_v<T> && std::is_trivial_v<T>), "Use alloc_pod for POD types");
+        void* mem = alloc(sizeof(detail::FinalizedObject<T>), alignof(detail::FinalizedObject<T>));
+        auto fo = static_cast<detail::FinalizedObject<T>*>(mem);
+        fo->f.fn = &detail::destructor<T>;
+        fo->f.next = finalizers_; // last in, first destructed.
+        finalizers_ = &fo->f;
+        return new (&fo->obj) T(std::forward<Args>(args)...);
+    }
+
+    /// Allocates a trivial object no destructor is run on scope exit.
+    template <typename T>
+    T* alloc_pod()
+    {
+        static_assert(std::is_standard_layout_v<T> && std::is_trivial_v<T>, "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

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<TestStruct>(&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<TestPOD>();
+        EXPECT_NE(current, arena.current());
+    }
+    EXPECT_EQ(current, arena.current());
+}