Write __gc metamethods for userdata, fix other memory leaks, add Table::{has_key, raw_has_key}

src/main.rs

@@ -121,7 +121,7 @@ fn main() -> Result<()> {
             f2(v1, v2)
         end
 
-        f3(v1, v2)
+        --f3(v1, v2)
     "#)?;
 
     // I don't care about the result of this execution, so I set the result as a

src/state.rs

@@ -81,6 +81,12 @@ impl<const N: usize> From<&[StdLibrary; N]> for StdLibraries {
     }
 }
 
+/// A struct that is used as an upvalue for creating Lua callable functions
+struct ClosureData<'a> {
+    wrapper_fn: Box<dyn Fn(&'a State, i32) -> i32>,
+    state: &'a State,
+}
+
 #[derive(Default)]
 struct ExtraSpace {
     userdata_metatables: HashMap<TypeId, LuaRef>,
@@ -94,9 +100,12 @@ impl Drop for State {
     fn drop(&mut self) {
         unsafe { 
             let extra = self.get_extra_space_ptr();
-            extra.drop_in_place();
         
             lua::lua_close(self.lua.as_ptr()); 
+            extra.drop_in_place();
+
+            // must be dealloced since it wasn't memory created from lua (like userdata)
+            alloc::dealloc(extra.cast(), Layout::new::<ExtraSpace>());
         }
     }
 }
@@ -282,6 +291,15 @@ impl State {
         }
     }
 
+    /// A function registered to the __gc metamethod of created ClosureData and drops ClosureData
+    unsafe extern "C-unwind" fn gc_drop_closure_data(s: *mut lua::lua_State) -> i32 {
+        let cptr = lua::lua_touserdata(s, -1)
+            .cast::<ClosureData>();
+        cptr.drop_in_place();
+
+        0
+    }
+
     pub fn create_function<'a, A, R, F>(&self, f: F) -> Result<Function>
     where
         A: FromLuaVec<'a>,
@@ -323,11 +341,6 @@ impl State {
             }
         }
 
-        struct ClosureData<'a> {
-            wrapper_fn: Box<dyn Fn(&'a State, i32) -> i32>,
-            state: &'a State,
-        }
-
         let wrapper_fn = move |lua: &State, _narg: i32| -> i32 {
             unsafe {
                 let lua: &State = mem::transmute(lua); // transmute lifetimes
@@ -365,6 +378,16 @@ impl State {
             let ptr = ptr.cast();
             ptr::write(ptr, data);
 
+            // create a metatable that is made to drop the ClosureData
+            let mt = Table::new_meta_table(self, "ClosureData")?;
+            mt.push_to_lua_stack(self)?;
+            lua::lua_pushliteral(s, "__gc");
+            lua::lua_pushcfunction(s, Self::gc_drop_closure_data);
+            lua::lua_rawset(s, -3);
+
+            // set the mt to the userdata above
+            lua::lua_setmetatable(s, -2);
+
             lua::lua_pushcclosure(s, rust_closure, 1);
             let lref = LuaRef::from_stack(self)?;
 
@@ -471,6 +494,26 @@ impl State {
             mt.set_meta(mm_name, mm_func)?;
         }
 
+        // if a Gc metamethod was not manually defined, set it
+        if !mt.raw_has_key(MetaMethod::Gc)? {
+            let gc_func = self.create_function(move |lua: &State, ud: AnyUserdata| {
+                unsafe {
+                    // if this panics, there's a weird bug.
+                    let ud_ptr = ud.as_ptr_unchecked::<T>().unwrap();
+                    ud_ptr.drop_in_place();
+
+                    lua::luaL_unref(lua.state_ptr(), lua::LUA_REGISTRYINDEX, *ud.lref.0);
+
+                    let extra = lua.get_extra_space();
+                    extra.userdata_metatables.remove(&TypeId::of::<T>());
+                }
+
+                Ok(())
+            })?;
+            
+            mt.set_meta(MetaMethod::Gc, gc_func)?;
+        }
+
         let extra = self.get_extra_space();
         extra.userdata_metatables.insert(TypeId::of::<T>(), mt.lref.clone());
 

src/table.rs

@@ -147,6 +147,28 @@ impl<'a> Table<'a> {
         }
     }
 
+    /// Returns a boolean indicating if this table has a key
+    pub fn has_key<K>(&self, key: K) -> Result<bool>
+    where
+        K: PushToLuaStack<'a>,
+    {
+        let s = self.state.state_ptr();
+        unsafe {
+            self.state.ensure_stack(2)?;
+            let _g = StackGuard::new(self.state);
+            
+            self.lref.push_to_lua_stack(self.state)?;
+            key.push_to_lua_stack(self.state)?;
+
+            // table[key] is at top of stack
+            if lua::lua_gettable(s, -2) == lua::LUA_TNIL {
+                Ok(true)
+            } else {
+                Ok(false)
+            }
+        }
+    }
+
     /// Set a key to a value in the table without calling any meta methods.
     pub fn raw_set<K, V>(&self, key: K, val: V) -> Result<()>
     where
@@ -201,6 +223,23 @@ impl<'a> Table<'a> {
         }
     }
 
+    /// Returns a boolean indicating if this table has a key without calling any meta methods.
+    pub fn raw_has_key<K>(&self, key: K) -> Result<bool>
+    where
+        K: PushToLuaStack<'a>,
+    {
+        let s = self.state.state_ptr();
+        unsafe {
+            self.state.ensure_stack(2)?;
+            let _g = StackGuard::new(self.state);
+
+            self.lref.push_to_lua_stack(self.state)?;
+            key.push_to_lua_stack(self.state)?;
+            // table[key] is at top of stack
+            Ok(lua::lua_rawget(s, -2) != lua::LUA_TNIL)
+        }
+    }
+
     /// Set something in the metatable
     /// 
     /// Does nothing if this table is not a metatable

src/userdata.rs

@@ -77,6 +77,14 @@ impl AsRef<str> for MetaMethod {
     }
 }
 
+impl<'a> AsLua<'a> for MetaMethod {
+    fn as_lua(&self, _lua: &'a State) -> crate::Result<Value<'a>> {
+        let s = self.as_ref();
+
+        Ok(Value::String(s.to_string()))
+    }
+}
+
 impl<'a> PushToLuaStack<'a> for MetaMethod {
     unsafe fn push_to_lua_stack(&self, state: &'a State) -> crate::Result<()> {
         let s = self.as_ref().to_string();
@@ -247,7 +255,7 @@ pub trait Userdata: Sized {
 /// A handle to some userdata on the stack
 #[derive(Clone)]
 pub struct AnyUserdata<'a> {
-    lref: LuaRef,
+    pub(crate) lref: LuaRef,
     state: &'a State,
 }
 
@@ -259,6 +267,7 @@ impl<'a> AnyUserdata<'a> {
         }
     }
 
+    /// Returns a borrow to the userdata.
     pub fn as_ref<T: Userdata + 'static>(&self) -> crate::Result<&'a T> {
         unsafe {
             self.state.ensure_stack(3)?;
@@ -283,6 +292,7 @@ impl<'a> AnyUserdata<'a> {
         }
     }
 
+    /// Returns a mutable reference to the userdata.
     pub fn as_mut<T: Userdata + 'static>(&self) -> crate::Result<&'a mut T> {
         unsafe {
             self.state.ensure_stack(3)?;
@@ -307,6 +317,50 @@ impl<'a> AnyUserdata<'a> {
         }
     }
 
+    /// Returns a mutable pointer to the userdata **WITHOUT verifying the type of it**.
+    /// 
+    /// # Safety
+    /// * You must be certain that the type `T` is the same type that this userdata has a handle to.
+    ///   There is a blind cast from `void*` to `T*`
+    /// 
+    /// If there is a possibility that these types do not match, use [`AnyUserdata::as_ptr`]
+    /// which does verify the types.
+    pub unsafe fn as_ptr_unchecked<T: Userdata + 'static>(&self) -> crate::Result<*mut T> {
+        self.state.ensure_stack(1)?;
+        let _g = StackGuard::new(self.state);
+        let s = self.state.state_ptr();
+
+        self.lref.push_to_lua_stack(self.state)?;
+        let cptr = lua::lua_touserdata(s, -3);
+
+        Ok(cptr.cast())
+    }
+
+    /// Returns a mutable pointer to the userdata.
+    /// 
+    /// This function ensures that the type of the userdata this struct has a handle to is the
+    /// same as `T`. If it isn't, a `UserdataMismatch` error will be returned.
+    pub unsafe fn as_ptr<T: Userdata + 'static>(&self) -> crate::Result<*mut T> {
+        let _g = StackGuard::new(self.state);
+        let s = self.state.state_ptr();
+
+        self.lref.push_to_lua_stack(self.state)?;
+        
+        if lua::lua_getmetatable(s, -1) == 0 {
+            return Err(crate::Error::MissingMetatable);
+        }
+
+        self.state.get_userdata_metatable::<T>()
+            .push_to_lua_stack(self.state)?;
+
+        if lua::lua_rawequal(s, -2, -1) == 1 {
+            let cptr = lua::lua_touserdata(s, -3);
+            Ok(cptr.cast::<T>())
+        } else {
+            Err(crate::Error::UserdataMismatch)
+        }
+    }
+
     /// Returns the name of the userdata by accessing the metatable
     pub fn name(&self) -> crate::Result<String> {
         unsafe {