Update "AnyUserData::take" to work on ref thread without need to push into stack.

Author: khvzak

src/scope.rs

@@ -8,9 +8,7 @@ use crate::state::{Lua, LuaGuard, RawLua};
 use crate::traits::{FromLuaMulti, IntoLuaMulti};
 use crate::types::{Callback, CallbackUpvalue, ScopedCallback, ValueRef};
 use crate::userdata::{AnyUserData, UserData, UserDataRegistry, UserDataStorage};
-use crate::util::{
-    self, assert_stack, check_stack, get_metatable_ptr, get_userdata, take_userdata, StackGuard,
-};
+use crate::util::{self, check_stack, get_metatable_ptr, get_userdata, take_userdata, StackGuard};
 
 /// Constructed by the [`Lua::scope`] method, allows temporarily creating Lua userdata and
 /// callbacks that are not required to be `Send` or `'static`.
@@ -284,22 +282,18 @@ impl<'scope, 'env: 'scope> Scope<'scope, 'env> {
     /// Shortens the lifetime of the userdata to the lifetime of the scope.
     fn seal_userdata<T: 'env>(&self, ud: &AnyUserData) {
         let destructor: DestructorCallback = Box::new(|rawlua, vref| unsafe {
-            let state = rawlua.state();
-            let _sg = StackGuard::new(state);
-            assert_stack(state, 2);
-
             // Ensure that userdata is not destructed
-            match rawlua.push_userdata_ref(&vref) {
+            match rawlua.get_userdata_ref_type_id(&vref) {
                 Ok(Some(_)) => {}
                 Ok(None) => {
                     // Deregister metatable
-                    let mt_ptr = get_metatable_ptr(state, -1);
+                    let mt_ptr = get_metatable_ptr(rawlua.ref_thread(), vref.index);
                     rawlua.deregister_userdata_metatable(mt_ptr);
                 }
                 Err(_) => return vec![],
             }
 
-            let data = take_userdata::<UserDataStorage<T>>(state);
+            let data = take_userdata::<UserDataStorage<T>>(rawlua.ref_thread(), vref.index);
             vec![Box::new(move || drop(data))]
         });
         self.destructors.0.borrow_mut().push((ud.0.clone(), destructor));

src/userdata.rs

@@ -683,22 +683,16 @@ impl AnyUserData {
     /// Keeps associated user values unchanged (they will be collected by Lua's GC).
     pub fn take<T: 'static>(&self) -> Result<T> {
         let lua = self.0.lua.lock();
-        let state = lua.state();
-        unsafe {
-            let _sg = StackGuard::new(state);
-            check_stack(state, 2)?;
-
-            let type_id = lua.push_userdata_ref(&self.0)?;
-            match type_id {
-                Some(type_id) if type_id == TypeId::of::<T>() => {
-                    if (*get_userdata::<UserDataStorage<T>>(state, -1)).has_exclusive_access() {
-                        take_userdata::<UserDataStorage<T>>(state).into_inner()
-                    } else {
-                        Err(Error::UserDataBorrowMutError)
-                    }
+        match lua.get_userdata_ref_type_id(&self.0)? {
+            Some(type_id) if type_id == TypeId::of::<T>() => unsafe {
+                let ref_thread = lua.ref_thread();
+                if (*get_userdata::<UserDataStorage<T>>(ref_thread, self.0.index)).has_exclusive_access() {
+                    take_userdata::<UserDataStorage<T>>(ref_thread, self.0.index).into_inner()
+                } else {
+                    Err(Error::UserDataBorrowMutError)
                 }
-                _ => Err(Error::UserDataTypeMismatch),
-            }
+            },
+            _ => Err(Error::UserDataTypeMismatch),
         }
     }
 

src/userdata/util.rs

@@ -455,9 +455,9 @@ pub(crate) unsafe extern "C" fn collect_userdata<T>(
 // It checks if the userdata is safe to destroy and sets the "destroyed" metatable
 // to prevent further GC collection.
 pub(super) unsafe extern "C-unwind" fn destroy_userdata_storage<T>(state: *mut ffi::lua_State) -> c_int {
-    let ud = get_userdata::<UserDataStorage<T>>(state, -1);
+    let ud = get_userdata::<UserDataStorage<T>>(state, 1);
     if (*ud).is_safe_to_destroy() {
-        take_userdata::<UserDataStorage<T>>(state);
+        take_userdata::<UserDataStorage<T>>(state, 1);
         ffi::lua_pushboolean(state, 1);
     } else {
         ffi::lua_pushboolean(state, 0);

src/util/userdata.rs

@@ -141,24 +141,27 @@ pub(crate) unsafe fn get_userdata<T>(state: *mut ffi::lua_State, index: c_int) -
     ud
 }
 
-// Pops the userdata off of the top of the stack and returns it to rust, invalidating the lua
-// userdata and gives it the special "destructed" userdata metatable. Userdata must not have been
-// previously invalidated, and this method does not check for this.
-// Uses 1 extra stack space and does not call checkstack.
-pub(crate) unsafe fn take_userdata<T>(state: *mut ffi::lua_State) -> T {
-    // We set the metatable of userdata on __gc to a special table with no __gc method and with
-    // metamethods that trigger an error on access. We do this so that it will not be double
-    // dropped, and also so that it cannot be used or identified as any particular userdata type
-    // after the first call to __gc.
+/// Unwraps `T` from the Lua userdata and invalidating it by setting the special "destructed"
+/// metatable.
+///
+/// This method does not check that userdata is of type `T` and was not previously invalidated.
+///
+/// Uses 1 extra stack space, does not call checkstack.
+pub(crate) unsafe fn take_userdata<T>(state: *mut ffi::lua_State, idx: c_int) -> T {
+    #[rustfmt::skip]
+    let idx = if idx < 0 { ffi::lua_absindex(state, idx) } else { idx };
+
+    // Update the metatable of this userdata to a special one with no `__gc` method and with
+    // metamethods that trigger an error on access.
+    // We do this so that it will not be double dropped or used after being dropped.
     get_destructed_userdata_metatable(state);
-    ffi::lua_setmetatable(state, -2);
-    let ud = get_userdata::<T>(state, -1);
+    ffi::lua_setmetatable(state, idx);
+    let ud = get_userdata::<T>(state, idx);
 
     // Update userdata tag to disable destructor and mark as destructed
     #[cfg(feature = "luau")]
-    ffi::lua_setuserdatatag(state, -1, 1);
+    ffi::lua_setuserdatatag(state, idx, 1);
 
-    ffi::lua_pop(state, 1);
     ptr::read(ud)
 }