From 04e30b9c7f32088a79a53af340367b969e20bd43 Mon Sep 17 00:00:00 2001
From: ValKmjolnir <lhk101lhk101@qq.com>
Date: Wed, 5 Feb 2025 19:43:02 +0800
Subject: [PATCH] :art: add memory usage calc

---
 src/main.cpp     |  6 +++++-
 src/nasal_gc.cpp | 29 ++++++++++++++++++++++++++---
 src/nasal_gc.h   |  6 ++++++
 src/nasal_vm.h   |  4 ++++
 4 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/src/main.cpp b/src/main.cpp
index 24b1419..7ba5e9d 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -75,6 +75,7 @@ void execute(const nasal::cli::cli_config& config) {
     // run
     const auto start = clk::now();
     double gc_time_ms = 0.0;
+    double gc_total_memory = 0.0;
     if (config.has(option::cli_debug_mode)) {
         auto debugger = std::make_unique<nasal::dbg>();
         debugger->run(
@@ -85,6 +86,7 @@ void execute(const nasal::cli::cli_config& config) {
             config.has(option::cli_profile_all)
         );
         gc_time_ms = debugger->get_gc_time_ms();
+        gc_total_memory = debugger->get_total_memory();
     } else if (config.has(option::cli_show_execute_time) ||
                config.has(option::cli_detail_info) ||
                config.has(option::cli_limit_mode) ||
@@ -94,6 +96,7 @@ void execute(const nasal::cli::cli_config& config) {
         runtime->set_limit_mode_flag(config.has(option::cli_limit_mode));
         runtime->run(gen, ld, config.nasal_vm_args);
         gc_time_ms = runtime->get_gc_time_ms();
+        gc_total_memory = runtime->get_total_memory();
     }
 
     // get running time
@@ -104,7 +107,8 @@ void execute(const nasal::cli::cli_config& config) {
         std::clog << "process exited after ";
         std::clog << execute_time_sec << "s, gc time: ";
         std::clog << gc_time_sec << "s (";
-        std::clog << gc_time_sec / execute_time_sec * 100.0 << "%)\n\n";
+        std::clog << gc_time_sec / execute_time_sec * 100.0 << "%), ";
+        std::clog << "memory usage: " << gc_total_memory << "MB\n\n";
     }
 }
 
diff --git a/src/nasal_gc.cpp b/src/nasal_gc.cpp
index 6192849..b565c7e 100644
--- a/src/nasal_gc.cpp
+++ b/src/nasal_gc.cpp
@@ -202,9 +202,28 @@ void gc::extend(const vm_type type) {
         memory.push_back(tmp);
         unused[index].push_back(tmp);
     }
+    switch(type) {
+        case vm_type::vm_str:
+            total_memory_usage += incr[index] * sizeof(std::string); break;
+        case vm_type::vm_vec:
+            total_memory_usage += incr[index] * sizeof(nas_vec); break;
+        case vm_type::vm_hash:
+            total_memory_usage += incr[index] * sizeof(nas_hash); break;
+        case vm_type::vm_func:
+            total_memory_usage += incr[index] * sizeof(nas_func); break;
+        case vm_type::vm_upval:
+            total_memory_usage += incr[index] * sizeof(nas_upval); break;
+        case vm_type::vm_ghost:
+            total_memory_usage += incr[index] * sizeof(nas_ghost); break;
+        case vm_type::vm_co:
+            total_memory_usage += incr[index] * sizeof(nas_co); break;
+        case vm_type::vm_map:
+            total_memory_usage += incr[index] * sizeof(nas_map); break;
+        default: break;
+    }
 
     // if incr[index] = 1, this will always be 1
-    incr[index] = incr[index]+incr[index]/2;
+    incr[index] = incr[index] + incr[index]/2;
 }
 
 void gc::init(const std::vector<std::string>& constant_strings,
@@ -220,7 +239,7 @@ void gc::init(const std::vector<std::string>& constant_strings,
 
     // init constant strings
     strs.resize(constant_strings.size());
-    for(u64 i = 0; i<strs.size(); ++i) {
+    for (u64 i = 0; i < strs.size(); ++i) {
         // incremental initialization, avoid memory leak in repl mode
         if (strs[i].is_str() && strs[i].str()==constant_strings[i]) {
             continue;
@@ -228,11 +247,13 @@ void gc::init(const std::vector<std::string>& constant_strings,
         strs[i] = var::gcobj(new nas_val(vm_type::vm_str));
         strs[i].val.gcobj->immutable = 1;
         strs[i].str() = constant_strings[i];
+        total_memory_usage += strs[i].str().size();
+        total_memory_usage += sizeof(std::string);
     }
 
     // record arguments
     env_argv.resize(argv.size());
-    for(u64 i = 0; i<argv.size(); ++i) {
+    for (u64 i = 0; i < argv.size(); ++i) {
         // incremental initialization, avoid memory leak in repl mode
         if (env_argv[i].is_str() && env_argv[i].str()==argv[i]) {
             continue;
@@ -240,6 +261,8 @@ void gc::init(const std::vector<std::string>& constant_strings,
         env_argv[i] = var::gcobj(new nas_val(vm_type::vm_str));
         env_argv[i].val.gcobj->immutable = 1;
         env_argv[i].str() = argv[i];
+        total_memory_usage += env_argv[i].str().size();
+        total_memory_usage += sizeof(std::string);
     }
 }
 
diff --git a/src/nasal_gc.h b/src/nasal_gc.h
index 4827608..76631aa 100644
--- a/src/nasal_gc.h
+++ b/src/nasal_gc.h
@@ -52,6 +52,8 @@ struct gc {
         16,  // vm_co
         2,   // vm_map
     };
+    // total memory usage, not very accurate
+    u64 total_memory_usage = 0;
 
     /* values for analysis */
     u64 size[gc_type_size];
@@ -100,6 +102,10 @@ public:
         return worktime * 1.0 / den * 1000.0;
     }
 
+    double get_total_memory() const {
+        return total_memory_usage * 1.0 / 1024.0 / 1024.0;
+    }
+
 public:
     var newstr(char c) {
         var s = alloc(vm_type::vm_str);
diff --git a/src/nasal_vm.h b/src/nasal_vm.h
index 43e6644..0335c14 100644
--- a/src/nasal_vm.h
+++ b/src/nasal_vm.h
@@ -325,6 +325,10 @@ public:
         return ngc.get_gc_time_ms();
     }
 
+    auto get_total_memory() const {
+        return ngc.get_total_memory();
+    }
+
     void set_interrupt_ptr(std::atomic<bool>* p) {
         interrupt_ptr = p;
     }