Add add_method_once and add_async_method_once UserData methods (experimental).

They will allow implementing userdata methods that can be called only once, destructing
userdata instance during the call.

Author: khvzak

src/userdata.rs

@@ -12,7 +12,7 @@ use crate::string::String;
 use crate::table::{Table, TablePairs};
 use crate::traits::{FromLua, FromLuaMulti, IntoLua, IntoLuaMulti};
 use crate::types::{MaybeSend, ValueRef};
-use crate::util::{check_stack, get_userdata, push_string, take_userdata, StackGuard};
+use crate::util::{check_stack, get_userdata, push_string, short_type_name, take_userdata, StackGuard};
 use crate::value::Value;
 
 #[cfg(feature = "async")]
@@ -273,6 +273,29 @@ pub trait UserDataMethods<T> {
         A: FromLuaMulti,
         R: IntoLuaMulti;
 
+    /// Add a method which accepts `T` as the first parameter.
+    ///
+    /// The userdata `T` will be moved out of the userdata container. This is useful for
+    /// methods that need to consume the userdata.
+    ///
+    /// The method can be called only once per userdata instance, subsequent calls will result in a
+    /// [`Error::UserDataDestructed`] error.
+    #[doc(hidden)]
+    fn add_method_once<M, A, R>(&mut self, name: impl Into<StdString>, method: M)
+    where
+        T: 'static,
+        M: Fn(&Lua, T, A) -> Result<R> + MaybeSend + 'static,
+        A: FromLuaMulti,
+        R: IntoLuaMulti,
+    {
+        let name = name.into();
+        let method_name = format!("{}.{name}", short_type_name::<T>());
+        self.add_function(name, move |lua, (ud, args): (AnyUserData, A)| {
+            let this = (ud.take()).map_err(|err| Error::bad_self_argument(&method_name, err))?;
+            method(lua, this, args)
+        });
+    }
+
     /// Add an async method which accepts a `&T` as the first parameter and returns [`Future`].
     ///
     /// Refer to [`add_method`] for more information about the implementation.
@@ -303,6 +326,34 @@ pub trait UserDataMethods<T> {
         MR: Future<Output = Result<R>> + MaybeSend + 'static,
         R: IntoLuaMulti;
 
+    /// Add an async method which accepts a `T` as the first parameter and returns [`Future`].
+    ///
+    /// The userdata `T` will be moved out of the userdata container. This is useful for
+    /// methods that need to consume the userdata.
+    ///
+    /// The method can be called only once per userdata instance, subsequent calls will result in a
+    /// [`Error::UserDataDestructed`] error.
+    #[cfg(feature = "async")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "async")))]
+    #[doc(hidden)]
+    fn add_async_method_once<M, A, MR, R>(&mut self, name: impl Into<StdString>, method: M)
+    where
+        T: 'static,
+        M: Fn(Lua, T, A) -> MR + MaybeSend + 'static,
+        A: FromLuaMulti,
+        MR: Future<Output = Result<R>> + MaybeSend + 'static,
+        R: IntoLuaMulti,
+    {
+        let name = name.into();
+        let method_name = format!("{}.{name}", short_type_name::<T>());
+        self.add_async_function(name, move |lua, (ud, args): (AnyUserData, A)| {
+            match (ud.take()).map_err(|err| Error::bad_self_argument(&method_name, err)) {
+                Ok(this) => either::Either::Left(method(lua, this, args)),
+                Err(err) => either::Either::Right(async move { Err(err) }),
+            }
+        });
+    }
+
     /// Add a regular method as a function which accepts generic arguments.
     ///
     /// The first argument will be a [`AnyUserData`] of type `T` if the method is called with Lua

tests/async.rs

@@ -423,9 +423,9 @@ async fn test_async_thread_pool() -> Result<()> {
 
 #[tokio::test]
 async fn test_async_userdata() -> Result<()> {
-    struct MyUserData(u64);
+    struct MyUserdata(u64);
 
-    impl UserData for MyUserData {
+    impl UserData for MyUserdata {
         fn add_methods<M: UserDataMethods<Self>>(methods: &mut M) {
             methods.add_async_method("get_value", |_, data, ()| async move {
                 sleep_ms(10).await;
@@ -438,6 +438,11 @@ async fn test_async_userdata() -> Result<()> {
                 Ok(())
             });
 
+            methods.add_async_method_once("take_value", |_, data, ()| async move {
+                sleep_ms(10).await;
+                Ok(data.0)
+            });
+
             methods.add_async_function("sleep", |_, n| async move {
                 sleep_ms(n).await;
                 Ok(format!("elapsed:{}ms", n))
@@ -479,7 +484,7 @@ async fn test_async_userdata() -> Result<()> {
     let lua = Lua::new();
     let globals = lua.globals();
 
-    let userdata = lua.create_userdata(MyUserData(11))?;
+    let userdata = lua.create_userdata(MyUserdata(11))?;
     globals.set("userdata", &userdata)?;
 
     lua.load(
@@ -518,6 +523,21 @@ async fn test_async_userdata() -> Result<()> {
     #[cfg(not(any(feature = "lua51", feature = "luau")))]
     assert_eq!(userdata.call_async::<String>(()).await?, "elapsed:24ms");
 
+    // Take value
+    let userdata2 = lua.create_userdata(MyUserdata(0))?;
+    globals.set("userdata2", userdata2)?;
+    lua.load("assert(userdata:take_value() == 24)")
+        .exec_async()
+        .await?;
+    match lua.load("userdata2.take_value(userdata)").exec_async().await {
+        Err(Error::CallbackError { cause, .. }) => {
+            let err = cause.to_string();
+            assert!(err.contains("bad argument `self` to `MyUserdata.take_value`"));
+            assert!(err.contains("userdata has been destructed"));
+        }
+        r => panic!("expected Err(CallbackError), got {r:?}"),
+    }
+
     Ok(())
 }
 

tests/userdata.rs

@@ -428,6 +428,39 @@ fn test_userdata_destroy() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_userdata_method_once() -> Result<()> {
+    struct MyUserdata(Arc<i64>);
+
+    impl UserData for MyUserdata {
+        fn add_methods<M: UserDataMethods<Self>>(methods: &mut M) {
+            methods.add_method_once("take_value", |_, this, ()| Ok(*this.0));
+        }
+    }
+
+    let lua = Lua::new();
+    let rc = Arc::new(42);
+    let userdata = lua.create_userdata(MyUserdata(rc.clone()))?;
+    lua.globals().set("userdata", &userdata)?;
+
+    // Control userdata
+    let userdata2 = lua.create_userdata(MyUserdata(rc.clone()))?;
+    lua.globals().set("userdata2", userdata2)?;
+
+    assert_eq!(lua.load("userdata:take_value()").eval::<i64>()?, 42);
+    match lua.load("userdata2.take_value(userdata)").eval::<i64>() {
+        Err(Error::CallbackError { cause, .. }) => {
+            let err = cause.to_string();
+            assert!(err.contains("bad argument `self` to `MyUserdata.take_value`"));
+            assert!(err.contains("userdata has been destructed"));
+        }
+        r => panic!("expected Err(CallbackError), got {r:?}"),
+    }
+    assert_eq!(Arc::strong_count(&rc), 2);
+
+    Ok(())
+}
+
 #[test]
 fn test_user_values() -> Result<()> {
     struct MyUserData;