Add UserDataOwned<T> wrapper to take ownership of userdata T

It implements `FromLua` and takes ownership of a Lua userdata value.
The semantics is similar to `AnyUserData::take`, preventing any
further use from Lua.

Closes #686

Author: khvzak

src/lib.rs

@@ -142,8 +142,8 @@ pub use crate::table::{TablePairs, TableSequence};
 pub use crate::thread::ThreadStatus;
 #[doc(hidden)]
 pub use crate::userdata::{
-    MetaMethod, UserData, UserDataFields, UserDataMetatable, UserDataMethods, UserDataRef, UserDataRefMut,
-    UserDataRegistry,
+    MetaMethod, UserData, UserDataFields, UserDataMetatable, UserDataMethods, UserDataOwned, UserDataRef,
+    UserDataRefMut, UserDataRegistry,
 };
 
 #[cfg(not(feature = "luau"))]

src/prelude.rs

@@ -9,7 +9,7 @@ pub use crate::{
     ObjectLike as LuaObjectLike, RegistryKey as LuaRegistryKey, Result as LuaResult, StdLib as LuaStdLib,
     Table as LuaTable, Thread as LuaThread, UserData as LuaUserData, UserDataFields as LuaUserDataFields,
     UserDataMetatable as LuaUserDataMetatable, UserDataMethods as LuaUserDataMethods,
-    UserDataRef as LuaUserDataRef, UserDataRefMut as LuaUserDataRefMut,
+    UserDataOwned as LuaUserDataOwned, UserDataRef as LuaUserDataRef, UserDataRefMut as LuaUserDataRefMut,
     UserDataRegistry as LuaUserDataRegistry, Value as LuaValue, Variadic as LuaVariadic,
     VmState as LuaVmState, WeakLua, chunk::AsChunk as AsLuaChunk, chunk::Chunk as LuaChunk,
     chunk::ChunkMode as LuaChunkMode, error::ErrorContext as LuaErrorContext,

src/userdata.rs

@@ -30,7 +30,7 @@ use {
 
 // Re-export for convenience
 pub(crate) use cell::UserDataStorage;
-pub use r#ref::{UserDataRef, UserDataRefMut};
+pub use r#ref::{UserDataOwned, UserDataRef, UserDataRefMut};
 pub use registry::UserDataRegistry;
 pub(crate) use registry::{RawUserDataRegistry, UserDataProxy};
 pub(crate) use util::{

src/userdata/ref.rs

@@ -7,7 +7,7 @@ use crate::error::{Error, Result};
 use crate::state::{Lua, RawLua};
 use crate::traits::FromLua;
 use crate::userdata::AnyUserData;
-use crate::util::get_userdata;
+use crate::util::{check_stack, get_userdata, take_userdata};
 use crate::value::Value;
 
 use super::cell::{UserDataStorage, UserDataVariant};
@@ -440,6 +440,66 @@ impl<T> DerefMut for UserDataRefMutInner<T> {
     }
 }
 
+/// A wrapper type that takes ownership of a userdata value.
+///
+/// It implements [`FromLua`] and can be used to receive a typed userdata from Lua by taking
+/// ownership of it.
+/// The original Lua userdata is marked as destructed and cannot be used further.
+pub struct UserDataOwned<T>(pub T);
+
+impl<T> Deref for UserDataOwned<T> {
+    type Target = T;
+
+    #[inline]
+    fn deref(&self) -> &T {
+        &self.0
+    }
+}
+
+impl<T> DerefMut for UserDataOwned<T> {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut T {
+        &mut self.0
+    }
+}
+
+impl<T: fmt::Debug> fmt::Debug for UserDataOwned<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        (**self).fmt(f)
+    }
+}
+
+impl<T: fmt::Display> fmt::Display for UserDataOwned<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        (**self).fmt(f)
+    }
+}
+
+impl<T: 'static> FromLua for UserDataOwned<T> {
+    fn from_lua(value: Value, _: &Lua) -> Result<Self> {
+        try_value_to_userdata::<T>(value)?.take().map(UserDataOwned)
+    }
+
+    unsafe fn from_stack(idx: c_int, lua: &RawLua) -> Result<Self> {
+        let state = lua.state();
+        let type_id = lua.get_userdata_type_id::<T>(state, idx)?;
+        match type_id {
+            Some(type_id) if type_id == TypeId::of::<T>() => {
+                let ud = get_userdata::<UserDataStorage<T>>(state, idx);
+                if (*ud).has_exclusive_access() {
+                    check_stack(state, 1)?;
+                    take_userdata::<UserDataStorage<T>>(state, idx)
+                        .into_inner()
+                        .map(UserDataOwned)
+                } else {
+                    Err(Error::UserDataBorrowMutError)
+                }
+            }
+            _ => Err(Error::UserDataTypeMismatch),
+        }
+    }
+}
+
 #[inline]
 fn try_value_to_userdata<T>(value: Value) -> Result<AnyUserData> {
     match value {
@@ -464,6 +524,10 @@ mod assertions {
     static_assertions::assert_impl_all!(UserDataRefMut<()>: Sync, Send);
     #[cfg(feature = "send")]
     static_assertions::assert_not_impl_all!(UserDataRefMut<std::rc::Rc<()>>: Send, Sync);
+    #[cfg(feature = "send")]
+    static_assertions::assert_impl_all!(UserDataOwned<()>: Send, Sync);
+    #[cfg(feature = "send")]
+    static_assertions::assert_not_impl_all!(UserDataOwned<std::rc::Rc<()>>: Send, Sync);
 
     #[cfg(not(feature = "send"))]
     static_assertions::assert_not_impl_all!(UserDataRef<()>: Send, Sync);

tests/userdata.rs

@@ -7,7 +7,7 @@ use std::sync::atomic::{AtomicI64, Ordering};
 
 use mlua::{
     AnyUserData, Error, ExternalError, Function, Lua, LuaString, MetaMethod, Nil, ObjectLike, Result,
-    UserData, UserDataFields, UserDataMethods, UserDataRef, UserDataRegistry, Value, Variadic,
+    UserData, UserDataFields, UserDataMethods, UserDataOwned, UserDataRef, UserDataRegistry, Value, Variadic,
 };
 
 #[test]
@@ -1459,3 +1459,49 @@ fn test_userdata_get_path() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_userdata_owned() -> Result<()> {
+    #[derive(Debug)]
+    struct MyUserdata(Arc<i64>);
+
+    impl UserData for MyUserdata {
+        fn register(registry: &mut UserDataRegistry<Self>) {
+            registry.add_method("num", |_, this, ()| Ok(*this.0));
+        }
+    }
+
+    let lua = Lua::new();
+    let rc = Arc::new(42);
+
+    // It takes ownership and destructs the Lua userdata
+    let ud = lua.create_userdata(MyUserdata(rc.clone()))?;
+    assert_eq!(Arc::strong_count(&rc), 2);
+    let owned: UserDataOwned<MyUserdata> = lua.convert(&ud)?;
+    assert_eq!(*owned.0.0, 42);
+    drop(owned);
+    assert_eq!(Arc::strong_count(&rc), 1);
+    match ud.borrow::<MyUserdata>() {
+        Err(Error::UserDataDestructed) => {}
+        r => panic!("expected UserDataDestructed, got {:?}", r),
+    }
+
+    // Cannot take while borrowed
+    let rc = Arc::new(7);
+    let ud = lua.create_userdata(MyUserdata(rc.clone()))?;
+    let borrowed = ud.borrow::<MyUserdata>()?;
+    match lua.convert::<UserDataOwned<MyUserdata>>(&ud) {
+        Err(Error::UserDataBorrowMutError) => {}
+        r => panic!("expected UserDataBorrowMutError, got {:?}", r),
+    }
+    drop(borrowed);
+
+    // Works as a function parameter
+    let f = lua.create_function(|_, owned: UserDataOwned<MyUserdata>| Ok(*owned.0.0))?;
+    let rc = Arc::new(55);
+    let ud = lua.create_userdata(MyUserdata(rc.clone()))?;
+    assert_eq!(f.call::<i64>(ud)?, 55);
+    assert_eq!(Arc::strong_count(&rc), 1); // dropped after call
+
+    Ok(())
+}