From a9ead56505d41cfdd87f050a6d0f7cf038de436d Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@foxmail.com>
Date: Thu, 23 Jan 2025 14:35:40 +0800
Subject: [PATCH] ...

---
 include/pocketpy/interpreter/objectpool.h |  2 ++
 src/interpreter/objectpool.c              | 32 ++++++++++++++++++++++-
 src/modules/pkpy.c                        | 22 ++++++++++++++++
 3 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/include/pocketpy/interpreter/objectpool.h b/include/pocketpy/interpreter/objectpool.h
index 751c4975..c9700583 100644
--- a/include/pocketpy/interpreter/objectpool.h
+++ b/include/pocketpy/interpreter/objectpool.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include "pocketpy/common/vector.h"
+#include "pocketpy/common/str.h"
 
 #define kPoolArenaSize (120 * 1024)
 #define kMultiPoolCount 5
@@ -28,3 +29,4 @@ void* MultiPool__alloc(MultiPool* self, int size);
 int MultiPool__sweep_dealloc(MultiPool* self);
 void MultiPool__ctor(MultiPool* self);
 void MultiPool__dtor(MultiPool* self);
+c11_string* MultiPool__summary(MultiPool* self);
\ No newline at end of file
diff --git a/src/interpreter/objectpool.c b/src/interpreter/objectpool.c
index 3c49a529..686ca5b7 100644
--- a/src/interpreter/objectpool.c
+++ b/src/interpreter/objectpool.c
@@ -2,6 +2,7 @@
 
 #include "pocketpy/config.h"
 #include "pocketpy/objects/object.h"
+#include "pocketpy/common/sstream.h"
 
 #include <assert.h>
 #include <stdbool.h>
@@ -173,4 +174,33 @@ void MultiPool__dtor(MultiPool* self) {
     for(int i = 0; i < kMultiPoolCount; i++) {
         Pool__dtor(&self->pools[i]);
     }
-}
\ No newline at end of file
+}
+
+c11_string* MultiPool__summary(MultiPool* self) {
+    c11_sbuf sbuf;
+    c11_sbuf__ctor(&sbuf);
+    for(int i = 0; i < kMultiPoolCount; i++) {
+        Pool* item = &self->pools[i];
+        int total_bytes = (item->arenas.length + item->no_free_arenas.length) * kPoolArenaSize;
+        int used_bytes = 0;
+        for(int j = 0; j < item->arenas.length; j++) {
+            PoolArena* arena = c11__getitem(PoolArena*, &item->arenas, j);
+            used_bytes += (arena->block_count - arena->unused_length) * arena->block_size;
+        }
+        used_bytes += item->no_free_arenas.length * kPoolArenaSize;
+        float used_pct = (float)used_bytes / total_bytes * 100;
+        char buf[256];
+        snprintf(buf,
+                 sizeof(buf),
+                 "Pool<%d>: len(arenas)=%d, len(no_free_arenas)=%d, %d/%d (%.1f%% used)",
+                 item->block_size,
+                 item->arenas.length,
+                 item->no_free_arenas.length,
+                 used_bytes,
+                 total_bytes,
+                 used_pct);
+        c11_sbuf__write_cstr(&sbuf, buf);
+        c11_sbuf__write_char(&sbuf, '\n');
+    }
+    return c11_sbuf__submit(&sbuf);
+}
diff --git a/src/modules/pkpy.c b/src/modules/pkpy.c
index 4706c564..bc58b19c 100644
--- a/src/modules/pkpy.c
+++ b/src/modules/pkpy.c
@@ -1,3 +1,5 @@
+#include "pocketpy/interpreter/objectpool.h"
+#include "pocketpy/objects/base.h"
 #include "pocketpy/pocketpy.h"
 
 #include "pocketpy/common/utils.h"
@@ -33,6 +35,24 @@ DEF_TVALUE_METHODS(float, _f64)
 DEF_TVALUE_METHODS(vec2, _vec2)
 DEF_TVALUE_METHODS(vec2i, _vec2i)
 
+static bool pkpy_memory_usage(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(0);
+    ManagedHeap* heap = &pk_current_vm->heap;
+    c11_string* small_objects_usage = MultiPool__summary(&heap->small_objects);
+    int large_object_count = heap->large_objects.length;
+    c11_sbuf buf;
+    c11_sbuf__ctor(&buf);
+    c11_sbuf__write_cstr(&buf, "== heap.small_objects ==\n");
+    c11_sbuf__write_cstr(&buf, small_objects_usage->data);
+    c11_sbuf__write_cstr(&buf, "== heap.large_objects ==\n");
+    c11_sbuf__write_cstr(&buf, "len(large_objects)=\n");
+    c11_sbuf__write_int(&buf, large_object_count);
+    // c11_sbuf__write_cstr(&buf, "== vm.pool_frame ==\n");
+    c11_sbuf__py_submit(&buf, py_retval());
+    c11_string__delete(small_objects_usage);
+    return true;
+}
+
 void pk__add_module_pkpy() {
     py_Ref mod = py_newmodule("pkpy");
 
@@ -66,6 +86,8 @@ void pk__add_module_pkpy() {
 
     py_setdict(mod, py_name("TValue"), TValue_dict);
     py_pop();
+
+    py_bindfunc(mod, "memory_usage", pkpy_memory_usage);
 }
 
 #undef DEF_TVALUE_METHODS
\ No newline at end of file