From b4ad911dce9b931d0813e62e0ebc2df3d8648e71 Mon Sep 17 00:00:00 2001
From: Nam Junghyun <13118114+cr0sh@users.noreply.github.com>
Date: Sat, 16 Sep 2023 16:03:34 +0900
Subject: [PATCH 1/3] Rename MemoryLimitNotAvailable to MemoryStatsNotAvailable

---
 src/error.rs | 6 +++---
 src/lua.rs   | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index 2a86eacc..3b1d451c 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -42,11 +42,11 @@ pub enum Error {
     GarbageCollectorError(StdString),
     /// Potentially unsafe action in safe mode.
     SafetyError(StdString),
-    /// Setting memory limit is not available.
+    /// Setting memory limit or getting/resetting the number of allocation is not available.
     ///
     /// This error can only happen when Lua state was not created by us and does not have the
     /// custom allocator attached.
-    MemoryLimitNotAvailable,
+    MemoryStatsNotAvailable,
     /// A mutable callback has triggered Lua code that has called the same mutable callback again.
     ///
     /// This is an error because a mutable callback can only be borrowed mutably once.
@@ -218,7 +218,7 @@ impl fmt::Display for Error {
             Error::SafetyError(ref msg) => {
                 write!(fmt, "safety error: {msg}")
             },
-            Error::MemoryLimitNotAvailable => {
+            Error::MemoryStatsNotAvailable => {
                 write!(fmt, "setting memory limit is not available")
             }
             Error::RecursiveMutCallback => write!(fmt, "mutable callback called recursively"),
diff --git a/src/lua.rs b/src/lua.rs
index 11765cee..df6aa933 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -1121,7 +1121,7 @@ impl Lua {
         unsafe {
             match (*self.extra.get()).mem_state.map(|mut x| x.as_mut()) {
                 Some(mem_state) => Ok(mem_state.set_memory_limit(limit)),
-                None => Err(Error::MemoryLimitNotAvailable),
+                None => Err(Error::MemoryStatsNotAvailable),
             }
         }
     }

From 28a6319423b4359546cc39bf94d188fd21c62aae Mon Sep 17 00:00:00 2001
From: Nam Junghyun <13118114+cr0sh@users.noreply.github.com>
Date: Sat, 16 Sep 2023 16:04:14 +0900
Subject: [PATCH 2/3] Introduce Lua::(reset_)num_allocations

---
 src/lua.rs    | 30 ++++++++++++++++++++++++++++++
 src/memory.rs | 12 ++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/src/lua.rs b/src/lua.rs
index df6aa933..012b14bd 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -1126,6 +1126,36 @@ impl Lua {
         }
     }
 
+    /// Returns the count of memory allocations invoked by this Lua runtime.
+    ///
+    /// Note that in case of overflow, the counter resets to zero(simply speaking, it 'wraps'),
+    /// and the overflowing threshold above the maximum depends on the platform.
+    ///
+    /// Returns [`Error::MemoryStatsNotAvailable`] if Lua state is managed externally.
+    pub fn num_allocations(&self) -> Result<usize> {
+        unsafe {
+            match (*self.extra.get()).mem_state.map(|x| x.as_ref()) {
+                Some(mem_state) => Ok(mem_state.num_allocations()),
+                None => Err(Error::MemoryStatsNotAvailable),
+            }
+        }
+    }
+
+    /// Resets the count of memory allocations to zero.
+    ///
+    /// Returns [`Error::MemoryStatsNotAvailable`] if Lua state is managed externally.
+    pub fn reset_num_allocations(&self) -> Result<()> {
+        unsafe {
+            match (*self.extra.get()).mem_state.map(|mut x| x.as_mut()) {
+                Some(mem_state) => {
+                    mem_state.reset_num_allocations();
+                    Ok(())
+                }
+                None => Err(Error::MemoryStatsNotAvailable),
+            }
+        }
+    }
+
     /// Returns true if the garbage collector is currently running automatically.
     ///
     /// Requires `feature = "lua54/lua53/lua52/luau"`
diff --git a/src/memory.rs b/src/memory.rs
index e199a759..d489f24b 100644
--- a/src/memory.rs
+++ b/src/memory.rs
@@ -17,6 +17,7 @@ pub(crate) struct MemoryState {
     // Indicates that the memory limit was reached on the last allocation.
     #[cfg(feature = "luau")]
     limit_reached: bool,
+    num_allocations: usize,
 }
 
 impl MemoryState {
@@ -37,6 +38,16 @@ impl MemoryState {
         prev_limit as usize
     }
 
+    #[inline]
+    pub(crate) fn num_allocations(&self) -> usize {
+        self.num_allocations
+    }
+
+    #[inline]
+    pub(crate) fn reset_num_allocations(&mut self) {
+        self.num_allocations = 0;
+    }
+
     // This function is used primarily for calling `lua_pushcfunction` in lua5.1/jit
     // to bypass the memory limit (if set).
     #[cfg(any(feature = "lua51", feature = "luajit"))]
@@ -139,6 +150,7 @@ unsafe extern "C-unwind" fn allocator(
         if new_ptr.is_null() {
             alloc::handle_alloc_error(new_layout);
         }
+        mem_state.num_allocations = mem_state.num_allocations.wrapping_add(1);
         return new_ptr;
     }
 

From bae1a20fd88338aee5a122e8e94f16743cf1630b Mon Sep 17 00:00:00 2001
From: Nam Junghyun <13118114+cr0sh@users.noreply.github.com>
Date: Sat, 23 Sep 2023 18:16:59 +0900
Subject: [PATCH 3/3] Fix MemoryStatsNotAvailable message

---
 src/error.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/error.rs b/src/error.rs
index 3b1d451c..f4a339ce 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -219,7 +219,7 @@ impl fmt::Display for Error {
                 write!(fmt, "safety error: {msg}")
             },
             Error::MemoryStatsNotAvailable => {
-                write!(fmt, "setting memory limit is not available")
+                write!(fmt, "memory stats information is not available")
             }
             Error::RecursiveMutCallback => write!(fmt, "mutable callback called recursively"),
             Error::CallbackDestructed => write!(