Add AnyUserData::destroy method

Author: khvzak

src/state/raw.rs

@@ -27,7 +27,7 @@ use crate::util::{
     assert_stack, check_stack, get_destructed_userdata_metatable, get_internal_userdata, get_main_state,
     get_metatable_ptr, get_userdata, init_error_registry, init_internal_metatable, init_userdata_metatable,
     pop_error, push_internal_userdata, push_string, push_table, rawset_field, safe_pcall, safe_xpcall,
-    short_type_name, StackGuard, WrappedFailure,
+    short_type_name, take_userdata, StackGuard, WrappedFailure,
 };
 use crate::value::{Nil, Value};
 
@@ -960,21 +960,26 @@ impl RawLua {
             }
         }
 
-        #[cfg(feature = "luau")]
-        let extra_init = None;
-        #[cfg(not(feature = "luau"))]
-        let extra_init: Option<fn(*mut ffi::lua_State, c_int) -> Result<()>> = Some(|state, mt_idx| {
-            ffi::lua_pushcfunction(state, crate::util::userdata_destructor::<UserDataStorage<T>>);
-            rawset_field(state, mt_idx, "__gc")
-        });
+        unsafe extern "C-unwind" fn userdata_destructor<T>(state: *mut ffi::lua_State) -> c_int {
+            let ud = get_userdata::<UserDataStorage<T>>(state, -1);
+            if !(*ud).is_borrowed() {
+                take_userdata::<UserDataStorage<T>>(state);
+                ffi::lua_pushboolean(state, 1);
+            } else {
+                ffi::lua_pushboolean(state, 0);
+            }
+            1
+        }
+
+        ffi::lua_pushcfunction(state, userdata_destructor::<T>);
+        rawset_field(state, metatable_index, "__gc")?;
 
         init_userdata_metatable(
             state,
             metatable_index,
             field_getters_index,
             field_setters_index,
             methods_index,
-            extra_init,
         )?;
 
         // Update stack guard to keep metatable after return

src/userdata.rs

@@ -696,6 +696,31 @@ impl AnyUserData {
         }
     }
 
+    /// Destroys this userdata.
+    ///
+    /// This is similar to [`AnyUserData::take`], but it doesn't require a type.
+    ///
+    /// This method works for non-scoped userdata only.
+    pub fn destroy(&self) -> Result<()> {
+        let lua = self.0.lua.lock();
+        let state = lua.state();
+        unsafe {
+            let _sg = StackGuard::new(state);
+            check_stack(state, 3)?;
+
+            lua.push_userdata_ref(&self.0)?;
+            protect_lua!(state, 1, 1, fn(state) {
+                if ffi::luaL_callmeta(state, -1, cstr!("__gc")) == 0 {
+                    ffi::lua_pushboolean(state, 0);
+                }
+            })?;
+            if ffi::lua_isboolean(state, -1) != 0 && ffi::lua_toboolean(state, -1) != 0 {
+                return Ok(());
+            }
+            Err(Error::UserDataBorrowMutError)
+        }
+    }
+
     /// Sets an associated value to this [`AnyUserData`].
     ///
     /// The value may be any Lua value whatsoever, and can be retrieved with [`user_value`].

src/userdata/cell.rs

@@ -1,5 +1,5 @@
 use std::any::{type_name, TypeId};
-use std::cell::{RefCell, UnsafeCell};
+use std::cell::{Cell, RefCell, UnsafeCell};
 use std::fmt;
 use std::ops::{Deref, DerefMut};
 use std::os::raw::c_int;
@@ -99,6 +99,15 @@ impl<T> UserDataVariant<T> {
         }
     }
 
+    #[inline(always)]
+    fn borrow_count(&self) -> &Cell<usize> {
+        match self {
+            Self::Default(inner) => &inner.borrow_count,
+            #[cfg(feature = "serialize")]
+            Self::Serializable(inner) => &inner.borrow_count,
+        }
+    }
+
     #[inline(always)]
     fn as_ptr(&self) -> *mut T {
         match self {
@@ -130,6 +139,7 @@ impl Serialize for UserDataStorage<()> {
 /// A type that provides interior mutability for a userdata value (thread-safe).
 pub(crate) struct UserDataCell<T> {
     raw_lock: RawLock,
+    borrow_count: Cell<usize>,
     value: UnsafeCell<T>,
 }
 
@@ -143,6 +153,7 @@ impl<T> UserDataCell<T> {
     fn new(value: T) -> Self {
         UserDataCell {
             raw_lock: RawLock::INIT,
+            borrow_count: Cell::new(0),
             value: UnsafeCell::new(value),
         }
     }
@@ -291,7 +302,10 @@ pub(crate) struct UserDataBorrowRef<'a, T>(&'a UserDataVariant<T>);
 impl<T> Drop for UserDataBorrowRef<'_, T> {
     #[inline]
     fn drop(&mut self) {
-        unsafe { self.0.raw_lock().unlock_shared() };
+        unsafe {
+            self.0.borrow_count().set(self.0.borrow_count().get() - 1);
+            self.0.raw_lock().unlock_shared();
+        }
     }
 }
 
@@ -317,6 +331,7 @@ impl<'a, T> TryFrom<&'a UserDataVariant<T>> for UserDataBorrowRef<'a, T> {
         if !variant.raw_lock().try_lock_shared() {
             return Err(Error::UserDataBorrowError);
         }
+        variant.borrow_count().set(variant.borrow_count().get() + 1);
         Ok(UserDataBorrowRef(variant))
     }
 }
@@ -326,7 +341,10 @@ pub(crate) struct UserDataBorrowMut<'a, T>(&'a UserDataVariant<T>);
 impl<T> Drop for UserDataBorrowMut<'_, T> {
     #[inline]
     fn drop(&mut self) {
-        unsafe { self.0.raw_lock().unlock_exclusive() };
+        unsafe {
+            self.0.borrow_count().set(self.0.borrow_count().get() - 1);
+            self.0.raw_lock().unlock_exclusive();
+        }
     }
 }
 
@@ -354,6 +372,7 @@ impl<'a, T> TryFrom<&'a UserDataVariant<T>> for UserDataBorrowMut<'a, T> {
         if !variant.raw_lock().try_lock_exclusive() {
             return Err(Error::UserDataBorrowMutError);
         }
+        variant.borrow_count().set(variant.borrow_count().get() + 1);
         Ok(UserDataBorrowMut(variant))
     }
 }
@@ -470,6 +489,14 @@ impl<T> UserDataStorage<T> {
         Self::Scoped(ScopedUserDataVariant::Boxed(RefCell::new(data)))
     }
 
+    #[inline(always)]
+    pub(crate) fn is_borrowed(&self) -> bool {
+        match self {
+            Self::Owned(variant) => variant.borrow_count().get() > 0,
+            Self::Scoped(_) => true,
+        }
+    }
+
     #[inline]
     pub(crate) fn try_borrow_scoped<R>(&self, f: impl FnOnce(&T) -> R) -> Result<R> {
         match self {

src/util/mod.rs

@@ -21,9 +21,6 @@ pub(crate) use userdata::{
 pub(crate) use userdata::push_uninit_userdata;
 pub(crate) use userdata::push_userdata;
 
-#[cfg(not(feature = "luau"))]
-pub(crate) use userdata::userdata_destructor;
-
 // Checks that Lua has enough free stack space for future stack operations. On failure, this will
 // panic with an internal error message.
 #[inline]

src/util/userdata.rs

@@ -152,7 +152,6 @@ pub(crate) unsafe fn init_userdata_metatable(
     field_getters: Option<c_int>,
     field_setters: Option<c_int>,
     methods: Option<c_int>,
-    extra_init: Option<fn(*mut ffi::lua_State, c_int) -> Result<()>>,
 ) -> Result<()> {
     if field_getters.is_some() || methods.is_some() {
         // Push `__index` generator function
@@ -195,11 +194,6 @@ pub(crate) unsafe fn init_userdata_metatable(
         rawset_field(state, metatable, "__newindex")?;
     }
 
-    // Additional initialization
-    if let Some(extra_init) = extra_init {
-        extra_init(state, metatable)?;
-    }
-
     ffi::lua_pushboolean(state, 0);
     rawset_field(state, metatable, "__metatable")?;
 
@@ -345,7 +339,7 @@ unsafe fn init_userdata_metatable_newindex(state: *mut ffi::lua_State) -> Result
 }
 
 #[cfg(not(feature = "luau"))]
-pub(crate) unsafe extern "C-unwind" fn userdata_destructor<T>(state: *mut ffi::lua_State) -> c_int {
+unsafe extern "C-unwind" fn userdata_destructor<T>(state: *mut ffi::lua_State) -> c_int {
     // It's probably NOT a good idea to catch Rust panics in finalizer
     // Lua 5.4 ignores it, other versions generates `LUA_ERRGCMM` without calling message handler
     take_userdata::<T>(state);

tests/userdata.rs

@@ -376,7 +376,18 @@ fn test_userdata_take() -> Result<()> {
 fn test_userdata_destroy() -> Result<()> {
     struct MyUserdata(#[allow(unused)] Arc<()>);
 
-    impl UserData for MyUserdata {}
+    impl UserData for MyUserdata {
+        fn add_methods<M: UserDataMethods<Self>>(methods: &mut M) {
+            methods.add_method("try_destroy", |lua, _this, ()| {
+                let ud = lua.globals().get::<AnyUserData>("ud")?;
+                match ud.destroy() {
+                    Err(Error::UserDataBorrowMutError) => {}
+                    r => panic!("expected `UserDataBorrowMutError` error, got {:?}", r),
+                }
+                Ok(())
+            });
+        }
+    }
 
     let rc = Arc::new(());
 
@@ -394,6 +405,23 @@ fn test_userdata_destroy() -> Result<()> {
 
     assert_eq!(Arc::strong_count(&rc), 1);
 
+    let ud = lua.create_userdata(MyUserdata(rc.clone()))?;
+    assert_eq!(Arc::strong_count(&rc), 2);
+    let ud_ref = ud.borrow::<MyUserdata>()?;
+    // With active `UserDataRef` this methods only marks userdata as destructed
+    // without running destructor
+    ud.destroy()?;
+    assert_eq!(Arc::strong_count(&rc), 2);
+    drop(ud_ref);
+    assert_eq!(Arc::strong_count(&rc), 1);
+
+    // We cannot destroy (internally) borrowed userdata
+    let ud = lua.create_userdata(MyUserdata(rc.clone()))?;
+    lua.globals().set("ud", &ud)?;
+    lua.load("ud:try_destroy()").exec().unwrap();
+    ud.destroy()?;
+    assert_eq!(Arc::strong_count(&rc), 1);
+
     Ok(())
 }