mlua_derive(macros): Allow taking both Lua and &Lua in async functions.

Async functions/methods by default receive owned `Lua` variant.
When user want to take `&Lua`, the code fails with "arguments to this function are incorrect".
We can dynamically detect "owned" or "ref" varians and pass corresponding type.

Closes #713

Author: khvzak

mlua_derive/src/userdata/userdata_impl.rs

@@ -38,7 +38,7 @@ struct ArgInfo {
 
 struct MethodInfo {
     self_kind: SelfKind,
-    has_lua: bool,
+    lua: Option<LuaArg>,
     args: Vec<ArgInfo>,
 }
 
@@ -50,18 +50,31 @@ fn ref_inner_type(ty: &Type) -> Type {
     }
 }
 
-/// Check if the type is `&Lua` or `&mlua::Lua`.
-fn is_lua_ref(ty: &Type) -> bool {
-    let Type::Reference(ref_ty) = ty else { return false };
-    match &*ref_ty.elem {
-        Type::Path(p) if p.path.segments.len() == 1 => p.path.segments[0].ident == "Lua",
-        Type::Path(p) if p.path.segments.len() == 2 => {
-            p.path.segments[0].ident == "mlua" && p.path.segments[1].ident == "Lua"
-        }
+/// How the `Lua` context parameter is passed to the method.
+enum LuaArg {
+    Ref,
+    Owned,
+}
+
+/// Check if the type is `Lua` or `mlua::Lua`.
+fn is_lua_type(ty: &Type) -> bool {
+    let Type::Path(p) = ty else { return false };
+    match p.path.segments.len() {
+        1 => p.path.segments[0].ident == "Lua",
+        2 => p.path.segments[0].ident == "mlua" && p.path.segments[1].ident == "Lua",
         _ => false,
     }
 }
 
+/// Classify the method's `Lua` context parameter, if present.
+fn lua_arg_kind(ty: &Type) -> Option<LuaArg> {
+    match ty {
+        Type::Reference(r) if r.mutability.is_none() && is_lua_type(&r.elem) => Some(LuaArg::Ref),
+        ty if is_lua_type(ty) => Some(LuaArg::Owned),
+        _ => None,
+    }
+}
+
 /// Classify a `&[mut] T` parameter, returning the callback wrapper type.
 ///
 /// Known borrow types come from the mapping table `BORROW_WRAPPERS`.
@@ -128,10 +141,10 @@ fn try_unwrap_option(ty: &Type) -> Option<&Type> {
 /// Analyze method signature.
 ///
 /// Determine `self` kind and collect the callback arguments.
-/// Auto-detects `&Lua` as the first non-self parameter.
+/// Auto-detects `Lua` (owned or reference) as the first non-self parameter.
 fn analyze_self_and_args(sig: &Signature) -> syn::Result<MethodInfo> {
     let mut self_kind = SelfKind::None;
-    let mut has_lua = false;
+    let mut lua = None;
     let mut args = Vec::new();
     let mut check_first_typed = true;
 
@@ -147,12 +160,14 @@ fn analyze_self_and_args(sig: &Signature) -> syn::Result<MethodInfo> {
                 self_kind = SelfKind::Owned;
             }
             FnArg::Typed(typed) => {
-                if check_first_typed && is_lua_ref(&typed.ty) {
-                    has_lua = true;
+                if check_first_typed {
                     check_first_typed = false;
-                    continue;
+                    if let Some(kind) = lua_arg_kind(&typed.ty) {
+                        lua = Some(kind);
+                        continue;
+                    }
                 }
-                check_first_typed = false;
+
                 if let syn::Pat::Ident(pat_ident) = &*typed.pat {
                     let arg_type = &*typed.ty;
                     let mut option_inner = None;
@@ -202,11 +217,7 @@ fn analyze_self_and_args(sig: &Signature) -> syn::Result<MethodInfo> {
         }
     }
 
-    Ok(MethodInfo {
-        self_kind,
-        has_lua,
-        args,
-    })
+    Ok(MethodInfo { self_kind, lua, args })
 }
 
 fn strip_item_attrs(attrs: &[Attribute]) -> Vec<Attribute> {
@@ -343,6 +354,16 @@ pub fn userdata_impl(attr: TokenStream, item: TokenStream) -> TokenStream {
                 let info = try_compile!(analyze_self_and_args(&method.sig));
                 let is_async = method.sig.asyncness.is_some();
 
+                // Owned `Lua` is only available to async callbacks.
+                if !is_async && matches!(info.lua, Some(LuaArg::Owned)) {
+                    return syn::Error::new_spanned(
+                        &method.sig,
+                        "owned `Lua` parameter is only supported for `async` methods (use `&Lua` instead)",
+                    )
+                    .to_compile_error()
+                    .into();
+                }
+
                 if lua_attr.getter {
                     if is_async {
                         return syn::Error::new_spanned(&method.sig, "async field getter is not supported")
@@ -520,7 +541,7 @@ fn gen_call_args(info: &MethodInfo) -> TokenStream2 {
         _ => call_args.push(quote! { this }),
     }
 
-    if info.has_lua {
+    if info.lua.is_some() {
         call_args.push(quote! { lua });
     }
 
@@ -542,8 +563,10 @@ fn gen_async_call_args(info: &MethodInfo) -> TokenStream2 {
         SelfKind::Owned => call_args.push(quote! { this }),
     }
 
-    if info.has_lua {
-        call_args.push(quote! { lua });
+    match info.lua {
+        Some(LuaArg::Ref) => call_args.push(quote! { &lua }),
+        Some(LuaArg::Owned) => call_args.push(quote! { lua }),
+        None => {}
     }
 
     for arg in &info.args {

tests/userdata_macro.rs

@@ -452,6 +452,14 @@ mod async_tests {
             Ok(42)
         }
 
+        async fn lua_version(lua: &Lua, extra: Option<&str>) -> Result<String> {
+            (lua.globals().get("_VERSION")).map(|s: String| s + extra.unwrap_or(""))
+        }
+
+        async fn lua_version_owned(lua: Lua) -> Result<String> {
+            lua.globals().get("_VERSION")
+        }
+
         #[cfg(not(any(feature = "lua51", feature = "luau")))]
         #[lua(meta)]
         async fn __tostring(&self) -> Result<String> {
@@ -483,6 +491,26 @@ mod async_tests {
         .unwrap();
     }
 
+    #[tokio::test]
+    async fn test_async_lua_param() {
+        let lua = Lua::new();
+        lua.globals()
+            .set("AsyncCounter", lua.create_proxy::<AsyncCounter>().unwrap())
+            .unwrap();
+
+        lua.load(
+            r#"
+            local c = AsyncCounter.new()
+            assert(type(AsyncCounter.lua_version()) == "string")
+            assert(string.sub(AsyncCounter.lua_version(" extra"), -5) == "extra", "expected 'extra' at the end")
+            assert(type(AsyncCounter.lua_version_owned()) == "string")
+        "#,
+        )
+        .exec_async()
+        .await
+        .unwrap();
+    }
+
     #[tokio::test]
     async fn test_async_consume() {
         let lua = Lua::new();