From b34317e7eefe0aa38824dfa3bf939b7a2f9c741d Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@hclcat.games>
Date: Sun, 12 Feb 2023 06:40:22 +0800
Subject: [PATCH] Update memory.h

---
 src/memory.h | 32 ++++++++++++++++++++++++++------
 1 file changed, 26 insertions(+), 6 deletions(-)

diff --git a/src/memory.h b/src/memory.h
index b228a6f5..a4561234 100644
--- a/src/memory.h
+++ b/src/memory.h
@@ -2,15 +2,35 @@
 
 #include "common.h"
 
-#define MALLOC(count) malloc(count)
-#define FREE(p) free(p)
-
 namespace pkpy{
+    const int kMaxMemPoolSize = 512;
+    static thread_local std::vector<int*> _mem_pool(kMaxMemPoolSize);
+
     template<typename T>
     struct sp_deleter {
         inline static void call(int* counter){
+            if constexpr(std::is_same_v<T, i64> || std::is_same_v<T, f64>){
+                if(_mem_pool.size() < kMaxMemPoolSize){
+                    _mem_pool.push_back(counter);
+                    return;
+                }
+            }
             ((T*)(counter + 1))->~T();
-            FREE(counter);
+            free(counter);
+        }
+    };
+
+    template<typename T>
+    struct sp_allocator {
+        inline static void* call(size_t size){
+            if constexpr(std::is_same_v<T, i64> || std::is_same_v<T, f64>){
+                if(!_mem_pool.empty()){
+                    int* p = _mem_pool.back();
+                    _mem_pool.pop_back();
+                    return p;
+                }
+            }
+            return malloc(size);
         }
     };
 
@@ -82,7 +102,7 @@ namespace pkpy{
     shared_ptr<T> make_shared(Args&&... args) {
         static_assert(std::is_base_of_v<T, U>, "U must be derived from T");
         static_assert(std::has_virtual_destructor_v<T>, "T must have virtual destructor");
-        int* p = (int*)MALLOC(sizeof(int) + sizeof(U));
+        int* p = (int*)sp_allocator<U>::call(sizeof(int) + sizeof(U));
         *p = 1;
         new(p+1) U(std::forward<Args>(args)...);
         return shared_ptr<T>(p);
@@ -90,7 +110,7 @@ namespace pkpy{
 
     template <typename T, typename... Args>
     shared_ptr<T> make_shared(Args&&... args) {
-        int* p = (int*)MALLOC(sizeof(int) + sizeof(T));
+        int* p = (int*)sp_allocator<T>::call(sizeof(int) + sizeof(T));
         *p = 1;
         new(p+1) T(std::forward<Args>(args)...);
         return shared_ptr<T>(p);