From ed44520b4294610e9ca69c0c024198fdcf1ba103 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@hclcat.games>
Date: Sun, 12 Feb 2023 17:22:08 +0800
Subject: [PATCH] up

---
 src/memory.h  | 40 ++++++++-----------------------
 src/obj.h     | 65 ++++++++++++++++++++++++++++++++++++++++++++++++---
 src/safestl.h |  6 ++---
 3 files changed, 75 insertions(+), 36 deletions(-)

diff --git a/src/memory.h b/src/memory.h
index a4561234..921f0929 100644
--- a/src/memory.h
+++ b/src/memory.h
@@ -3,44 +3,26 @@
 #include "common.h"
 
 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;
-                }
-            }
+    struct SpAllocator {
+        template<typename U>
+        inline static int* alloc(){
+            return (int*)malloc(sizeof(int) + sizeof(U));
+        }
+
+        inline static void dealloc(int* counter){
             ((T*)(counter + 1))->~T();
             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);
-        }
-    };
-
     template <typename T>
     class shared_ptr {
         int* counter = nullptr;
 
 #define _t() ((T*)(counter + 1))
 #define _inc_counter() if(counter) ++(*counter)
-#define _dec_counter() if(counter && --(*counter) == 0){ pkpy::sp_deleter<T>::call(counter); }
+#define _dec_counter() if(counter && --(*counter) == 0){ SpAllocator<T>::dealloc(counter); }
 
     public:
         shared_ptr() {}
@@ -102,16 +84,14 @@ 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*)sp_allocator<U>::call(sizeof(int) + sizeof(U));
-        *p = 1;
+        int* p = SpAllocator<T>::template alloc<U>(); *p = 1;
         new(p+1) U(std::forward<Args>(args)...);
         return shared_ptr<T>(p);
     }
 
     template <typename T, typename... Args>
     shared_ptr<T> make_shared(Args&&... args) {
-        int* p = (int*)sp_allocator<T>::call(sizeof(int) + sizeof(T));
-        *p = 1;
+        int* p = SpAllocator<T>::template alloc<T>(); *p = 1;
         new(p+1) T(std::forward<Args>(args)...);
         return shared_ptr<T>(p);
     }
diff --git a/src/obj.h b/src/obj.h
index 37b48915..d52ee69d 100644
--- a/src/obj.h
+++ b/src/obj.h
@@ -77,7 +77,7 @@ public:
 struct PyObject {
     Type type;
     pkpy::NameDict* _attr;
-    //void* _tid;
+    void* _tid;
     inline bool is_attr_valid() const noexcept { return _attr != nullptr; }
     inline pkpy::NameDict& attr() noexcept { return *_attr; }
     inline PyVar& attr(const Str& name) noexcept { return (*_attr)[name]; }
@@ -85,7 +85,7 @@ struct PyObject {
     inline bool is_type(Type type) const noexcept{ return this->type == type; }
     virtual void* value() = 0;
 
-    PyObject(Type type) : type(type) {}
+    PyObject(Type type, void* _tid) : type(type), _tid(_tid) {}
     virtual ~PyObject() { delete _attr; }
 };
 
@@ -93,7 +93,7 @@ template <typename T>
 struct Py_ : PyObject {
     T _value;
 
-    Py_(Type type, T val) : PyObject(type), _value(val) {
+    Py_(Type type, T val) : PyObject(type, tid<T>()), _value(val) {
         if constexpr (std::is_same_v<T, Dummy> || std::is_same_v<T, Type>
         || std::is_same_v<T, pkpy::Function_> || std::is_same_v<T, pkpy::NativeFunc>) {
             _attr = new pkpy::NameDict();
@@ -111,3 +111,62 @@ struct Py_ : PyObject {
     inline static Type _type(VM* vm) { return OBJ_GET(Type, vm->_modules[#mod]->attr(#name)); } \
     inline static const char* _mod() { return #mod; } \
     inline static const char* _name() { return #name; }
+
+
+namespace pkpy {
+    template<int N>
+    struct MemBlock {
+        std::vector<void*> a;
+        int block_size;
+
+        MemBlock(int block_size) : block_size(block_size) {
+            new_block();
+        }
+
+        void new_block(){
+            int8_t* total = (int8_t*)malloc(N * block_size);
+            for(int i = 0; i < block_size; ++i){
+                a.push_back((void*)(total + i * N));
+            }
+        }
+
+        inline void* alloc(){
+            if(a.empty()) new_block();
+            void* p = a.back();
+            a.pop_back();
+            return p;
+        }
+
+        inline void dealloc(void* p) noexcept{
+            a.push_back(p);
+        }
+
+        ~MemBlock(){
+            free(a[0]);
+        }
+    };
+
+    static_assert(sizeof(i64) == sizeof(f64));
+    static thread_local MemBlock<sizeof(int)+sizeof(Py_<i64>)> _mem_i64(512);
+
+    template<>
+    struct SpAllocator<PyObject> {
+        template<typename U>
+        inline static int* alloc(){
+            if constexpr (std::is_same_v<U, Py_<i64>> || std::is_same_v<U, Py_<f64>>) {
+                return (int*)_mem_i64.alloc();
+            }
+            return (int*)malloc(sizeof(int) + sizeof(U));
+        }
+
+        inline static void dealloc(int* counter){
+            PyObject* obj = (PyObject*)(counter + 1);
+            obj->~PyObject();
+            if(obj->_tid == tid<i64>() || obj->_tid == tid<f64>()){
+                _mem_i64.dealloc(counter);
+            }else{
+                free(counter);
+            }
+        }
+    };
+}
\ No newline at end of file
diff --git a/src/safestl.h b/src/safestl.h
index b461343f..66b4a780 100644
--- a/src/safestl.h
+++ b/src/safestl.h
@@ -61,7 +61,7 @@ namespace pkpy {
             }
         }
 
-        void _free(){
+        void _dealloc(){
             if(_size == 0 || _args == nullptr) return;
             if(_size >= kMaxPoolSize || _args_pool[_size].size() > 32){
                 delete[] _args;
@@ -102,7 +102,7 @@ namespace pkpy {
         const PyVar& operator[](int i) const { return _args[i]; }
 
         Args& operator=(Args&& other) noexcept {
-            _free();
+            _dealloc();
             this->_args = other._args;
             this->_size = other._size;
             other._args = nullptr;
@@ -135,7 +135,7 @@ namespace pkpy {
             }
         }
 
-        ~Args(){ _free(); }
+        ~Args(){ _dealloc(); }
     };
 
     static const Args _zero(0);