From e045dea04cfe0fc421abe967dc615f3cbc97a1ef Mon Sep 17 00:00:00 2001
From: Siarhei Fedartsou <siarhei.fedartsou@gmail.com>
Date: Thu, 11 Jul 2024 23:21:48 +0200
Subject: [PATCH] Use pool in std::unordered_map

---
 include/util/pool_allocator.hpp | 46 +++++++++++++++++----------------
 include/util/query_heap.hpp     |  6 ++++-
 2 files changed, 29 insertions(+), 23 deletions(-)

diff --git a/include/util/pool_allocator.hpp b/include/util/pool_allocator.hpp
index a0203b941a..9015d111e5 100644
--- a/include/util/pool_allocator.hpp
+++ b/include/util/pool_allocator.hpp
@@ -14,14 +14,11 @@
 namespace osrm::util
 {
 
-template <typename T>
-class PoolAllocator;
-
-template <typename T>
 class MemoryManager
 {
 private:
     constexpr static size_t MIN_ITEMS_IN_BLOCK = 1024;
+
 public:
     static std::shared_ptr<MemoryManager> instance()
     {
@@ -33,9 +30,10 @@ class MemoryManager
         return instance;
     }
 
+    template <typename T>
     T *allocate(std::size_t n)
     {
-        size_t free_list_index = get_next_power_of_two_exponent(n);
+        size_t free_list_index = get_next_power_of_two_exponent(n * sizeof(T));
         auto &free_list = free_lists_[free_list_index];
         const auto items_in_block = 1u << free_list_index;
         if (free_list.empty())
@@ -43,32 +41,34 @@ class MemoryManager
             // Check if there is space in current block
             if (current_block_left_items_ < items_in_block)
             {
-                allocate_block(items_in_block);
+                allocate_block<T>(items_in_block);
             }
 
             free_list.push_back(current_block_ptr_);
             current_block_left_items_ -= items_in_block;
-            current_block_ptr_ += items_in_block;
+            current_block_ptr_ += items_in_block * sizeof(T);
         }
-        auto ptr = free_list.back();
+        auto ptr = static_cast<T*>(free_list.back());
         free_list.pop_back();
         return ptr;
     }
 
+    template <typename T>
     void deallocate(T *p, std::size_t n) noexcept
     {
-        size_t free_list_index = get_next_power_of_two_exponent(n);
+        size_t free_list_index = get_next_power_of_two_exponent(n * sizeof(T));
         free_lists_[free_list_index].push_back(p);
     }
 
     ~MemoryManager()
     {
-       // std::cerr << "~MemoryManager()" << std::endl;
+        std::cerr << "~MemoryManager()" << std::endl;
         for (auto block : blocks_)
         {
             std::free(block);
         }
     }
+
 private:
     MemoryManager() = default;
     MemoryManager(const MemoryManager &) = delete;
@@ -80,25 +80,26 @@ class MemoryManager
         return (sizeof(size_t) * 8) - std::countl_zero(n - 1);
     }
 
+    template <typename T>
     void allocate_block(size_t items_in_block)
     {
         items_in_block = std::max(items_in_block, MIN_ITEMS_IN_BLOCK);
 
         size_t block_size = items_in_block * sizeof(T);
-        T *block = static_cast<T *>(std::malloc(block_size));
+        void *block = std::malloc(block_size);
         if (!block)
         {
             throw std::bad_alloc();
         }
         total_allocated_ += block_size;
         blocks_.push_back(block);
-        current_block_ptr_ = block;
+        current_block_ptr_ = static_cast<uint8_t*>(block);
         current_block_left_items_ = items_in_block;
     }
 
-    std::array<std::vector<T *>, 32> free_lists_;
-    std::vector<T *> blocks_;
-    T *current_block_ptr_ = nullptr;
+    std::array<std::vector<void *>, 32> free_lists_;
+    std::vector<void *> blocks_;
+    uint8_t *current_block_ptr_ = nullptr;
     size_t current_block_left_items_ = 0;
 
     size_t total_allocated_ = 0;
@@ -110,10 +111,10 @@ class PoolAllocator
 public:
     using value_type = T;
 
-    PoolAllocator() noexcept : pool(MemoryManager<T>::instance()) {};
+    PoolAllocator() noexcept : pool(MemoryManager::instance()) {};
 
     template <typename U>
-    PoolAllocator(const PoolAllocator<U> &) noexcept : pool(MemoryManager<T>::instance()) {}
+    PoolAllocator(const PoolAllocator<U> &) noexcept : pool(MemoryManager::instance()) {}
 
     template <typename U>
     struct rebind
@@ -123,15 +124,17 @@ class PoolAllocator
 
     T *allocate(std::size_t n)
     {
-        return pool->allocate(n);
+        return pool->allocate<T>(n);
     }
 
     void deallocate(T *p, std::size_t n) noexcept
     {
-        pool->deallocate(p, n);
+        pool->deallocate<T>(p, n);
     }
 
-    ~PoolAllocator() = default;
+    ~PoolAllocator() {
+        std::cerr << "~PoolAllocator()" << std::endl;
+    }
 
     PoolAllocator(const PoolAllocator &) = default;
     PoolAllocator &operator=(const PoolAllocator &) = default;
@@ -139,9 +142,8 @@ class PoolAllocator
     PoolAllocator &operator=(PoolAllocator &&) noexcept = default;
 
 private:
-    std::shared_ptr<MemoryManager<T>> pool;
+    std::shared_ptr<MemoryManager> pool;
 };
-
 template <typename T, typename U>
 bool operator==(const PoolAllocator<T> &, const PoolAllocator<U> &)
 {
diff --git a/include/util/query_heap.hpp b/include/util/query_heap.hpp
index 44f1835533..367796ce41 100644
--- a/include/util/query_heap.hpp
+++ b/include/util/query_heap.hpp
@@ -101,6 +101,10 @@ template <typename NodeID, typename Key> class UnorderedMapStorage
   public:
     explicit UnorderedMapStorage(std::size_t) { nodes.rehash(1000); }
 
+    ~UnorderedMapStorage() {
+        std::cerr << "~UnorderedMapStorage()" << std::endl;
+    }
+
     Key &operator[](const NodeID node) { return nodes[node]; }
 
     Key peek_index(const NodeID node) const
@@ -124,7 +128,7 @@ template <typename NodeID, typename Key> class UnorderedMapStorage
   private:
     template <typename K, typename V>
     using UnorderedMap = std::
-        unordered_map<K, V/*, std::hash<K>, std::equal_to<K>, PoolAllocator<std::pair<const K, V>>*/>;
+        unordered_map<K, V, std::hash<K>, std::equal_to<K>, PoolAllocator<std::pair<const K, V>>>;
 
     UnorderedMap<NodeID, Key> nodes;
 };