From 73b9e9c3beed18996d2852db150a8568c59dbcd4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=CE=B6eh=20Matt?= <5415177+ZehMatt@users.noreply.github.com>
Date: Thu, 5 Sep 2024 07:56:01 +0300
Subject: [PATCH] Unify the object pools, specialization for pool types

---
 .../src/benchmarks/benchmark.assembler.cpp    |  9 ++--
 zasm/src/zasm/src/program/program.cpp         | 20 +++++----
 zasm/src/zasm/src/program/program.state.hpp   | 44 ++++++++++++-------
 3 files changed, 45 insertions(+), 28 deletions(-)

diff --git a/benchmark/src/benchmarks/benchmark.assembler.cpp b/benchmark/src/benchmarks/benchmark.assembler.cpp
index 9d9fd15..e9ba5a9 100644
--- a/benchmark/src/benchmarks/benchmark.assembler.cpp
+++ b/benchmark/src/benchmarks/benchmark.assembler.cpp
@@ -120,10 +120,13 @@ namespace zasm::benchmarks
             assembler.setCursor(nullptr);
             state.ResumeTiming();
 
-            for (const auto& instr : zasm::tests::data::Instructions)
+            for (int i = 0; i < 100; i++)
             {
-                instr.emitter(assembler);
-                numInstructions++;
+                for (const auto& instr : zasm::tests::data::Instructions)
+                {
+                    instr.emitter(assembler);
+                    numInstructions++;
+                }
             }
         }
 
diff --git a/zasm/src/zasm/src/program/program.cpp b/zasm/src/zasm/src/program/program.cpp
index 6574e81..031cec9 100644
--- a/zasm/src/zasm/src/program/program.cpp
+++ b/zasm/src/zasm/src/program/program.cpp
@@ -376,10 +376,9 @@ namespace zasm
         auto* nodeToDestroy = detail::toInternal(node);
 
         node->visit([&](auto& ptr) {
-
             using T = std::decay_t<decltype(ptr)>;
 
-            auto& objectPool = state.nodePools.getPool<T>();
+            auto& objectPool = state.objectPools.get<T>();
             objectPool.destroy(&ptr);
 
             if (!quickDestroy)
@@ -387,13 +386,14 @@ namespace zasm
                 objectPool.deallocate(&ptr, 1);
             }
         });
-        
-        state.nodePool.destroy(nodeToDestroy);
+
+        auto& nodePool = state.objectPools.get<Node>();
+        nodePool.destroy(nodeToDestroy);
 
         if (!quickDestroy)
         {
             // Release memory, when quickDestroy is true the entire pool will be cleared at once.
-            state.nodePool.deallocate(nodeToDestroy, 1);
+            nodePool.deallocate(nodeToDestroy, 1);
 
             // Remove mapping.
             auto& nodeMap = state.nodeMap;
@@ -438,7 +438,7 @@ namespace zasm
         _state->sections.clear();
         _state->labels.clear();
         _state->symbolNames.clear();
-        _state->nodePool.reset();
+        _state->objectPools.reset();
     }
 
     void Program::setEntryPoint(const Label& label)
@@ -456,19 +456,21 @@ namespace zasm
         const auto nextId = state.nextNodeId;
         state.nextNodeId = static_cast<Node::Id>(static_cast<std::underlying_type_t<Node::Id>>(nextId) + 1U);
 
-        auto& pool = state.nodePool;
-        auto* node = detail::toInternal(pool.allocate(1));
+        auto& nodePool = state.objectPools.get<Node>();
+        auto* node = detail::toInternal(nodePool.allocate(1));
         if (node == nullptr)
         {
             return nullptr;
         }
 
+        // Construct object.
         using ObjectType = std::decay_t<T>;
-        auto& objectPool = state.nodePools.getPool<ObjectType>();
+        auto& objectPool = state.objectPools.get<ObjectType>();
 
         auto* obj = objectPool.allocate(1);
         ::new ((void*)obj) ObjectType(std::move(object));
 
+        // Construct node.
         ::new ((void*)node) detail::Node(nextId, obj);
 
         notifyObservers<true>(&Observer::onNodeCreated, state.observer, node);
diff --git a/zasm/src/zasm/src/program/program.state.hpp b/zasm/src/zasm/src/program/program.state.hpp
index b80f2fd..4c2e2a1 100644
--- a/zasm/src/zasm/src/program/program.state.hpp
+++ b/zasm/src/zasm/src/program/program.state.hpp
@@ -23,8 +23,6 @@ namespace zasm
 
 namespace zasm::detail
 {
-    constexpr std::size_t PoolSize = 1U << 10;
-
     struct LabelData
     {
         Label::Id id{ Label::Id::Invalid };
@@ -47,25 +45,36 @@ namespace zasm::detail
         zasm::Node* node{};
     };
 
-    template<typename... TTypes> struct ObjectPools
+    namespace detail
     {
-        std::tuple<ObjectPool<TTypes, PoolSize>...> pools;
+        template<typename T>
+        struct PoolSize
+        {
+            static constexpr std::size_t kSize = 50'000;
+        };
 
-        template<typename T> auto& getPool()
+        template<> struct PoolSize<Instruction>
         {
-            return std::get<ObjectPool<T, PoolSize>>(pools);
-        }
-    };
+            static constexpr std::size_t kSize = 30'000;
+        };
 
-    using NodePools = ObjectPools<Sentinel, Instruction, Label, EmbeddedLabel, Data, Section, Align>;
+        template<typename... TTypes> struct ObjectPools
+        {
+            std::tuple<ObjectPool<TTypes, PoolSize<TTypes>::kSize>...> pools;
 
-    struct NodeStorage
-    {
-        ObjectPool<Node, PoolSize> nodePool;
-        Node::Id nextNodeId{};
+            template<typename T> auto& get()
+            {
+                return std::get<ObjectPool<T, PoolSize<T>::kSize>>(pools);
+            }
 
-        NodePools nodePools;
-    };
+            void reset()
+            {
+                std::apply([](auto&... pool) { (pool.reset(), ...); }, pools);
+            }
+        };
+    } // namespace detail
+
+    using ObjectPools = detail::ObjectPools<zasm::Node, Sentinel, Instruction, Label, EmbeddedLabel, Data, Section, Align>;
 
     struct NodeList
     {
@@ -79,7 +88,7 @@ namespace zasm::detail
         StringPool symbolNames;
     };
 
-    struct ProgramState : NodeStorage, NodeList, Symbols
+    struct ProgramState : NodeList, Symbols
     {
         MachineMode mode{};
 
@@ -94,6 +103,9 @@ namespace zasm::detail
 
         Label entryPoint{ Label::Id::Invalid };
 
+        ObjectPools objectPools;
+        Node::Id nextNodeId{};
+
         ProgramState(MachineMode m)
             : mode(m)
         {