From 0e3acf05e6cb4a34d242822360b5cd33e6002d99 Mon Sep 17 00:00:00 2001
From: jmd <joshua.m.dean@gmail.com>
Date: Wed, 25 Sep 2024 21:29:34 -0700
Subject: [PATCH] [util] add: C++ Allocator interface for Memory Arena

add: Some utility functions for calculating mb, kb, and gb
---
 lib/nw/util/memory.hpp | 63 ++++++++++++++++++++++++++++++++++++++++++
 tests/CMakeLists.txt   |  1 +
 tests/util_memory.cpp  | 21 ++++++++++++++
 3 files changed, 85 insertions(+)
 create mode 100644 tests/util_memory.cpp

diff --git a/lib/nw/util/memory.hpp b/lib/nw/util/memory.hpp
index ea6f10471..11a3f19bc 100644
--- a/lib/nw/util/memory.hpp
+++ b/lib/nw/util/memory.hpp
@@ -12,6 +12,25 @@
 
 namespace nw {
 
+/// Convert kilobytes to bytes
+constexpr std::uint64_t KB(std::uint64_t kb)
+{
+    return kb * 1024ULL;
+}
+
+/// Convert megabytes to bytes
+constexpr std::uint64_t MB(std::uint64_t mb)
+{
+    return mb * 1024ULL * 1024ULL;
+}
+
+/// Convert gigabytes to bytes
+constexpr std::uint64_t GB(std::uint64_t gb)
+{
+    return gb * 1024ULL * 1024ULL * 1024ULL;
+}
+
+/// A growable Memory Arena
 struct MemoryArena {
     MemoryArena(size_t blockSize = 1024);
     MemoryArena(const MemoryArena&) = delete;
@@ -37,6 +56,50 @@ struct MemoryArena {
     void alloc_block_(size_t size);
 };
 
+/// C++ Allocator interface for Memory arena
+template <typename T>
+class ArenaAllocator {
+public:
+    using value_type = T;
+    ArenaAllocator(MemoryArena* arena)
+        : arena_(arena)
+    {
+    }
+
+    template <typename U>
+    ArenaAllocator(const ArenaAllocator<U>& other)
+        : arena_(other.arena_)
+    {
+    }
+
+    /// Allocate memory for n objects of type T.
+    T* allocate(size_t n)
+    {
+        if (!arena_) { return nullptr; }
+        size_t size = n * sizeof(T);
+        void* ptr = arena_->allocate(size, alignof(T));
+        return static_cast<T*>(ptr);
+    }
+
+    /// Deallocate memory. a no-op.
+    void deallocate(T*, size_t) { }
+
+    template <typename U>
+    bool operator==(const ArenaAllocator<U>& other) const
+    {
+        return arena_ == other.arena_;
+    }
+
+    template <typename U>
+    bool operator!=(const ArenaAllocator<U>& other) const
+    {
+        return !(*this == other);
+    }
+
+private:
+    MemoryArena* arena_ = nullptr;
+};
+
 // This is very simple and naive.
 template <typename T, size_t chunk_size>
 struct ObjectPool {
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 4d992700d..9a7b30283 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -66,6 +66,7 @@ add_executable(rollnw_test
 
     serial_gff.cpp
 
+    util_memory.cpp
     util_string.cpp
     util_tokenizer.cpp
 )
diff --git a/tests/util_memory.cpp b/tests/util_memory.cpp
new file mode 100644
index 000000000..9cf49b43b
--- /dev/null
+++ b/tests/util_memory.cpp
@@ -0,0 +1,21 @@
+#include <gtest/gtest.h>
+
+#include <nw/util/memory.hpp>
+
+using namespace std::literals;
+
+TEST(Memory, Helpers)
+{
+    EXPECT_EQ(nw::KB(1), 1024ULL);
+    EXPECT_EQ(nw::MB(1), 1024ULL * 1024ULL);
+    EXPECT_EQ(nw::GB(1), 1024ULL * 1024ULL * 1024ULL);
+}
+
+TEST(Memory, Allocator)
+{
+    nw::MemoryArena arena;
+    nw::ArenaAllocator<char> allocator(&arena);
+    std::basic_string<char, std::char_traits<char>, nw::ArenaAllocator<char>> arena_string(allocator);
+    arena_string = "Hello, World";
+    EXPECT_EQ(arena_string, "Hello, World");
+}