Add Lua::(reset_)num_allocations API

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,8 +218,8 @@ impl fmt::Display for Error {
             Error::SafetyError(ref msg) => {
                 write!(fmt, "safety error: {msg}")
             },
-            Error::MemoryLimitNotAvailable => {
-                write!(fmt, "setting memory limit is not available")
+            Error::MemoryStatsNotAvailable => {
+                write!(fmt, "memory stats information is not available")
             }
             Error::RecursiveMutCallback => write!(fmt, "mutable callback called recursively"),
             Error::CallbackDestructed => write!(

src/lua.rs

@@ -1121,7 +1121,37 @@ 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),
+            }
+        }
+    }
+
+    /// 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),
             }
         }
     }

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;
     }