From 2c247fa75eaa32a2a982c64e14c178e87290c50e Mon Sep 17 00:00:00 2001
From: Andrew Noyes <andrew@weaselab.dev>
Date: Thu, 14 Aug 2025 13:05:05 -0400
Subject: [PATCH] Add ArenaStlAllocator, and use arena throughout CommitRequest

---
 src/arena_allocator.hpp | 52 +++++++++++++++++++++++++++++++++++++++++
 src/commit_request.hpp  | 51 +++++++++++++++++++++++++++-------------
 2 files changed, 87 insertions(+), 16 deletions(-)

diff --git a/src/arena_allocator.hpp b/src/arena_allocator.hpp
index 2006ee1..8e65ed8 100644
--- a/src/arena_allocator.hpp
+++ b/src/arena_allocator.hpp
@@ -423,3 +423,55 @@ private:
   /// Current offset within the current block's data area
   size_t current_offset_;
 };
+
+/**
+ * @brief STL-compatible allocator that uses ArenaAllocator for memory
+ * management.
+ * @tparam T The type of objects to allocate
+ */
+template <typename T> class ArenaStlAllocator {
+public:
+  using value_type = T;
+  using pointer = T *;
+  using const_pointer = const T *;
+  using reference = T &;
+  using const_reference = const T &;
+  using size_type = std::size_t;
+  using difference_type = std::ptrdiff_t;
+
+  template <typename U> struct rebind {
+    using other = ArenaStlAllocator<U>;
+  };
+
+  explicit ArenaStlAllocator(ArenaAllocator *arena) noexcept : arena_(arena) {}
+
+  template <typename U>
+  ArenaStlAllocator(const ArenaStlAllocator<U> &other) noexcept
+      : arena_(other.arena_) {}
+
+  T *allocate(size_type n) {
+    if (n == 0)
+      return nullptr;
+    return static_cast<T *>(arena_->allocate(n * sizeof(T), alignof(T)));
+  }
+
+  void deallocate(T *ptr, size_type n) noexcept {
+    // Arena allocator doesn't support individual deallocation
+    (void)ptr;
+    (void)n;
+  }
+
+  template <typename U>
+  bool operator==(const ArenaStlAllocator<U> &other) const noexcept {
+    return arena_ == other.arena_;
+  }
+
+  template <typename U>
+  bool operator!=(const ArenaStlAllocator<U> &other) const noexcept {
+    return arena_ != other.arena_;
+  }
+
+  ArenaAllocator *arena_;
+
+  template <typename U> friend class ArenaStlAllocator;
+};
diff --git a/src/commit_request.hpp b/src/commit_request.hpp
index 511c7e5..5787356 100644
--- a/src/commit_request.hpp
+++ b/src/commit_request.hpp
@@ -2,6 +2,7 @@
 
 #include "arena_allocator.hpp"
 #include <optional>
+#include <span>
 #include <stack>
 #include <string>
 #include <string_view>
@@ -61,10 +62,16 @@ public:
   };
 
   struct ParserContext {
-    std::stack<ParseState> state_stack;
-    std::string current_key;
-    std::string current_string;
-    std::string current_number;
+    using ArenaString = std::basic_string<char, std::char_traits<char>,
+                                          ArenaStlAllocator<char>>;
+    using ArenaStateStack =
+        std::stack<ParseState,
+                   std::vector<ParseState, ArenaStlAllocator<ParseState>>>;
+
+    ArenaStateStack state_stack;
+    ArenaString current_key;
+    ArenaString current_string;
+    ArenaString current_number;
     bool in_key = false;
     bool parse_error = false;
     bool parse_complete = false;
@@ -74,8 +81,18 @@ public:
     Operation current_operation{};
 
     // Parsing state for nested structures
-    std::string precondition_type;
-    std::string operation_type;
+    ArenaString precondition_type;
+    ArenaString operation_type;
+
+    // Constructor to initialize arena-allocated containers
+    explicit ParserContext(ArenaAllocator *arena)
+        : state_stack(ArenaStateStack::container_type(
+              ArenaStlAllocator<ParseState>(arena))),
+          current_key(ArenaStlAllocator<char>(arena)),
+          current_string(ArenaStlAllocator<char>(arena)),
+          current_number(ArenaStlAllocator<char>(arena)),
+          precondition_type(ArenaStlAllocator<char>(arena)),
+          operation_type(ArenaStlAllocator<char>(arena)) {}
   };
 
 private:
@@ -83,8 +100,8 @@ private:
   std::optional<std::string_view> request_id_;
   std::string_view leader_id_;
   uint64_t read_version_ = 0;
-  std::vector<Precondition> preconditions_;
-  std::vector<Operation> operations_;
+  std::vector<Precondition, ArenaStlAllocator<Precondition>> preconditions_;
+  std::vector<Operation, ArenaStlAllocator<Operation>> operations_;
   ParserContext parser_context_;
   WeaselJsonParser *json_parser_ = nullptr;
 
@@ -93,7 +110,11 @@ public:
    * @brief Construct a new CommitRequest with the given initial arena size.
    * @param arena_size Initial size for the arena allocator
    */
-  explicit CommitRequest(size_t arena_size = 4096) : arena_(arena_size) {}
+  explicit CommitRequest(size_t arena_size = 4096)
+      : arena_(arena_size),
+        preconditions_(ArenaStlAllocator<Precondition>(&arena_)),
+        operations_(ArenaStlAllocator<Operation>(&arena_)),
+        parser_context_(&arena_) {}
 
   /**
    * @brief Destructor - cleans up any active parser.
@@ -203,17 +224,15 @@ public:
 
   /**
    * @brief Get the preconditions.
-   * @return Vector of preconditions
+   * @return span of preconditions
    */
-  const std::vector<Precondition> &preconditions() const {
-    return preconditions_;
-  }
+  std::span<const Precondition> preconditions() const { return preconditions_; }
 
   /**
    * @brief Get the operations.
-   * @return Vector of operations
+   * @return span of operations
    */
-  const std::vector<Operation> &operations() const { return operations_; }
+  std::span<const Operation> operations() const { return operations_; }
 
   /**
    * @brief Get the total allocated bytes in the arena.
@@ -243,7 +262,7 @@ public:
       WeaselJsonParser_destroy(json_parser_);
       json_parser_ = nullptr;
     }
-    parser_context_ = ParserContext{};
+    parser_context_ = ParserContext(&arena_);
     parser_context_.state_stack.push(ParseState::Root);
   }