From a3dbe826130a68d1de7a869ded5890449ce1df56 Mon Sep 17 00:00:00 2001
From: SeanOMik <seanomik@gmail.com>
Date: Mon, 29 Jan 2024 10:38:46 -0500
Subject: [PATCH] unref LuaRefs when they are dropped

---
 src/function.rs |  4 ++--
 src/main.rs     | 45 +++++++++++++++++++++++++++++----------------
 src/state.rs    |  4 ++--
 src/table.rs    |  4 ++--
 src/userdata.rs |  4 ++--
 5 files changed, 37 insertions(+), 24 deletions(-)

diff --git a/src/function.rs b/src/function.rs
index eff8ee9..dbb5659 100755
--- a/src/function.rs
+++ b/src/function.rs
@@ -7,7 +7,7 @@ use mlua_sys as lua;
 #[derive(Clone)]
 pub struct Function<'a> {
     state: &'a State,
-    lref: LuaRef,
+    lref: LuaRef<'a>,
     pub(crate) error_handler: Option<Arc<Function<'a>>>,
 }
 
@@ -28,7 +28,7 @@ impl<'a> PushToLuaStack<'a> for Function<'a> {
 }
 
 impl<'a> Function<'a> {
-    pub fn from_ref(state: &'a State, lref: LuaRef) -> Self {
+    pub fn from_ref(state: &'a State, lref: LuaRef<'a>) -> Self {
         Self {
             state,
             lref,
diff --git a/src/main.rs b/src/main.rs
index 93ef9fe..d6ad586 100755
--- a/src/main.rs
+++ b/src/main.rs
@@ -143,42 +143,55 @@ fn main() -> Result<()> {
     Ok(())
 }
 
+/// A LuaRef is a a handle to something in Lua.
+/// 
+/// These are created with `luaL_ref`, they can be created from anything on the top of the stack
+/// with [`LuaRef::from_stack`]. The references are automatically freed with `luaL_unref` when
+/// the inner Arc detects a single strong count.
 #[derive(Clone)]
-pub struct LuaRef(Arc<i32>);
+pub struct LuaRef<'a>(Arc<i32>, &'a State);
 
-impl From<i32> for LuaRef {
-    fn from(value: i32) -> Self {
-        Self(Arc::new(value))
+impl<'a> Drop for LuaRef<'a> {
+    fn drop(&mut self) {
+        unsafe {
+            let s = self.1.state_ptr();
+            
+            if Arc::strong_count(&self.0) == 1 {
+                lua::luaL_unref(s, lua::LUA_REGISTRYINDEX, *self.0);
+            }
+        }
     }
 }
 
-impl LuaRef {
+impl<'a> LuaRef<'a> {
+    pub fn new(lua_ref: i32, state: &'a State) -> Self {
+        Self(Arc::new(lua_ref), state)
+    }
+
     /// Creates a reference to what is at the top of the stack.
-    pub unsafe fn from_stack(state: &State) -> Result<Self> {
+    pub unsafe fn from_stack(state: &'a State) -> Result<Self> {
         let s = state.state_ptr();
 
         let r = lua::luaL_ref(s, lua::LUA_REGISTRYINDEX);
         if r == lua::LUA_REFNIL {
             Err(Error::Nil)
         } else {
-            Ok(LuaRef::from(r))
+            Ok(LuaRef::new(r, state))
         }
     }
 }
 
-impl<'a> PushToLuaStack<'a> for LuaRef {
+impl<'a> PushToLuaStack<'a> for LuaRef<'a> {
     unsafe fn push_to_lua_stack(&self, state: &State) -> Result<()> {
         let s = state.state_ptr();
 
-        unsafe {
-            state.ensure_stack(1)?;
-            let top = lua::lua_gettop(s);
-            let ty = lua::lua_rawgeti(s, lua::LUA_REGISTRYINDEX, *self.0 as i64);
-            let new_top = lua::lua_gettop(s);
+        state.ensure_stack(1)?;
+        let top = lua::lua_gettop(s);
+        let ty = lua::lua_rawgeti(s, lua::LUA_REGISTRYINDEX, *self.0 as i64);
+        let new_top = lua::lua_gettop(s);
 
-            if ty == lua::LUA_TNIL || ty == lua::LUA_TNONE || top == new_top {
-                return Err(Error::Nil);
-            }
+        if ty == lua::LUA_TNIL || ty == lua::LUA_TNONE || top == new_top {
+            return Err(Error::Nil);
         }
 
         Ok(())
diff --git a/src/state.rs b/src/state.rs
index 3fcfb2d..94431ec 100755
--- a/src/state.rs
+++ b/src/state.rs
@@ -88,8 +88,8 @@ struct ClosureData<'a> {
 }
 
 #[derive(Default)]
-struct ExtraSpace {
-    userdata_metatables: HashMap<TypeId, LuaRef>,
+struct ExtraSpace<'a> {
+    userdata_metatables: HashMap<TypeId, LuaRef<'a>>,
 }
 
 pub struct State {
diff --git a/src/table.rs b/src/table.rs
index 77b4cd0..ceb810b 100755
--- a/src/table.rs
+++ b/src/table.rs
@@ -6,7 +6,7 @@ use crate::{ensure_type, FromLuaStack, LuaRef, PushToLuaStack, Result, StackGuar
 #[derive(Clone)]
 pub struct Table<'a> {
     state: &'a State,
-    pub(crate) lref: LuaRef,
+    pub(crate) lref: LuaRef<'a>,
     /// a boolean indicating if this Table is a Metatable
     mt_marker: bool,
 }
@@ -59,7 +59,7 @@ impl<'a> Table<'a> {
     }
 
     /// Construct a table with a lua reference to one.
-    pub fn with_ref(state: &'a State, lua_ref: LuaRef, is_metatable: bool) -> Result<Self> {
+    pub fn with_ref(state: &'a State, lua_ref: LuaRef<'a>, is_metatable: bool) -> Result<Self> {
         let s = state.state_ptr();
         
         unsafe {
diff --git a/src/userdata.rs b/src/userdata.rs
index 5f82b7f..5f85b2e 100755
--- a/src/userdata.rs
+++ b/src/userdata.rs
@@ -255,12 +255,12 @@ pub trait Userdata: Sized {
 /// A handle to some userdata on the stack
 #[derive(Clone)]
 pub struct AnyUserdata<'a> {
-    pub(crate) lref: LuaRef,
+    pub(crate) lref: LuaRef<'a>,
     state: &'a State,
 }
 
 impl<'a> AnyUserdata<'a> {
-    pub fn from_ref(state: &'a State, lref: LuaRef) -> Self {
+    pub fn from_ref(state: &'a State, lref: LuaRef<'a>) -> Self {
         Self {
             lref,
             state,