Skip to content

Commit 3f4a741

Browse files
khvzakkhvzak
committed
mlua_derive: Support Option<&[mut] T> callback parameters in userdata_impl macro
Closes #709
1 parent 4882011 commit 3f4a741

3 files changed

Lines changed: 99 additions & 25 deletions

File tree

‎mlua_derive/src/chunk/token.rs‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ use std::vec::IntoIter;
55

66
use proc_macro::{Delimiter, Span, TokenStream, TokenTree};
77
use proc_macro2::{Span as Span2, TokenStream as TokenStream2};
8-
use syn;
98

109
#[derive(Clone, Copy, Debug)]
1110
pub(crate) struct Pos {

‎mlua_derive/src/userdata/userdata_impl.rs‎

Lines changed: 67 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ enum SelfKind {
2626
enum RefKind {
2727
Ref,
2828
Mut,
29+
OptionRef,
30+
OptionMut,
2931
}
3032

3133
struct ArgInfo {
@@ -104,6 +106,25 @@ fn classify_ref_type(ty: &Type) -> Option<Type> {
104106
}
105107
}
106108

109+
/// If `ty` is `Option<InnerType>`, return the inner type.
110+
fn try_unwrap_option(ty: &Type) -> Option<&Type> {
111+
let Type::Path(type_path) = ty else { return None };
112+
let segment = type_path.path.segments.last()?;
113+
if segment.ident != "Option" {
114+
return None;
115+
}
116+
let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
117+
return None;
118+
};
119+
if args.args.len() != 1 {
120+
return None;
121+
}
122+
let syn::GenericArgument::Type(inner) = &args.args[0] else {
123+
return None;
124+
};
125+
Some(inner)
126+
}
127+
107128
/// Analyze method signature.
108129
///
109130
/// Determine `self` kind and collect the callback arguments.
@@ -134,12 +155,32 @@ fn analyze_self_and_args(sig: &Signature) -> syn::Result<MethodInfo> {
134155
check_first_typed = false;
135156
if let syn::Pat::Ident(pat_ident) = &*typed.pat {
136157
let arg_type = &*typed.ty;
158+
let mut option_inner = None;
137159
let ref_kind = match arg_type {
138160
Type::Reference(r) if r.mutability.is_some() => Some(RefKind::Mut),
139161
Type::Reference(_) => Some(RefKind::Ref),
140-
_ => None,
162+
_ => {
163+
// Check if it's `Option<&T>` or `Option<&mut T>`
164+
option_inner = try_unwrap_option(arg_type);
165+
option_inner.and_then(|inner| match inner {
166+
Type::Reference(r) if r.mutability.is_some() => Some(RefKind::OptionMut),
167+
Type::Reference(_) => Some(RefKind::OptionRef),
168+
_ => None,
169+
})
170+
}
141171
};
142172
let callback_type = match &ref_kind {
173+
Some(RefKind::OptionRef | RefKind::OptionMut) => {
174+
match classify_ref_type(option_inner.unwrap()) {
175+
Some(ty) => parse_quote! { Option<#ty> },
176+
None => {
177+
return Err(syn::Error::new_spanned(
178+
arg_type,
179+
"this reference type is not supported as a callback parameter",
180+
));
181+
}
182+
}
183+
}
143184
Some(_) => match classify_ref_type(arg_type) {
144185
Some(ty) => ty,
145186
None => {
@@ -449,7 +490,7 @@ fn gen_closure_destructure(info: &MethodInfo) -> TokenStream2 {
449490
.iter()
450491
.map(|a| {
451492
let ident = &a.ident;
452-
if matches!(a.userdata_ref, Some(RefKind::Mut)) {
493+
if matches!(a.userdata_ref, Some(RefKind::Mut | RefKind::OptionMut)) {
453494
quote! { mut #ident }
454495
} else {
455496
quote! { #ident }
@@ -460,6 +501,18 @@ fn gen_closure_destructure(info: &MethodInfo) -> TokenStream2 {
460501
quote! { (#(#idents),*): (#(#types),*) }
461502
}
462503

504+
/// Generate the call-site expression for a single argument.
505+
fn gen_arg_token(arg: &ArgInfo) -> TokenStream2 {
506+
let ident = &arg.ident;
507+
match arg.userdata_ref {
508+
Some(RefKind::Ref) => quote! { &*#ident },
509+
Some(RefKind::Mut) => quote! { &mut *#ident },
510+
Some(RefKind::OptionRef) => quote! { #ident.as_ref().map(|r| &**r) },
511+
Some(RefKind::OptionMut) => quote! { #ident.as_mut().map(|r| &mut **r) },
512+
None => quote! { #ident },
513+
}
514+
}
515+
463516
/// Generate call arguments for invoking the original method.
464517
fn gen_call_args(info: &MethodInfo) -> TokenStream2 {
465518
let mut call_args: Vec<TokenStream2> = Vec::new();
@@ -474,12 +527,7 @@ fn gen_call_args(info: &MethodInfo) -> TokenStream2 {
474527
}
475528

476529
for arg in &info.args {
477-
let ident = &arg.ident;
478-
match arg.userdata_ref {
479-
Some(RefKind::Ref) => call_args.push(quote! { &*#ident }),
480-
Some(RefKind::Mut) => call_args.push(quote! { &mut *#ident }),
481-
None => call_args.push(quote! { #ident }),
482-
}
530+
call_args.push(gen_arg_token(arg));
483531
}
484532

485533
quote! { #(#call_args),* }
@@ -491,8 +539,8 @@ fn gen_async_call_args(info: &MethodInfo) -> TokenStream2 {
491539

492540
match info.self_kind {
493541
SelfKind::None => {}
494-
SelfKind::Ref(RefKind::Ref) => call_args.push(quote! { &this }),
495-
SelfKind::Ref(RefKind::Mut) => call_args.push(quote! { &mut this }),
542+
SelfKind::Ref(RefKind::Mut | RefKind::OptionMut) => call_args.push(quote! { &mut this }),
543+
SelfKind::Ref(_) => call_args.push(quote! { &this }),
496544
SelfKind::Owned => call_args.push(quote! { this }),
497545
}
498546

@@ -501,12 +549,7 @@ fn gen_async_call_args(info: &MethodInfo) -> TokenStream2 {
501549
}
502550

503551
for arg in &info.args {
504-
let ident = &arg.ident;
505-
match arg.userdata_ref {
506-
Some(RefKind::Ref) => call_args.push(quote! { &*#ident }),
507-
Some(RefKind::Mut) => call_args.push(quote! { &mut *#ident }),
508-
None => call_args.push(quote! { #ident }),
509-
}
552+
call_args.push(gen_arg_token(arg));
510553
}
511554

512555
quote! { #(#call_args),* }
@@ -640,12 +683,12 @@ fn gen_regular_method(
640683
quote! { #fn_path(#call_args) }
641684
};
642685
match info.self_kind {
643-
SelfKind::Ref(RefKind::Ref) => quote! {
644-
registry.add_method(#lua_name, #closure_params { #body });
645-
},
646-
SelfKind::Ref(RefKind::Mut) => quote! {
686+
SelfKind::Ref(RefKind::Mut | RefKind::OptionMut) => quote! {
647687
registry.add_method_mut(#lua_name, #closure_params { #body });
648688
},
689+
SelfKind::Ref(_) => quote! {
690+
registry.add_method(#lua_name, #closure_params { #body });
691+
},
649692
SelfKind::Owned => quote! {
650693
registry.add_method_once(#lua_name, #closure_params { #body });
651694
},
@@ -672,12 +715,12 @@ fn gen_async_regular_method(
672715
quote! { async move { #fn_path(#call_args).await } }
673716
};
674717
match info.self_kind {
675-
SelfKind::Ref(RefKind::Ref) => quote! {
676-
registry.add_async_method(#lua_name, #closure_params #body);
677-
},
678-
SelfKind::Ref(RefKind::Mut) => quote! {
718+
SelfKind::Ref(RefKind::Mut | RefKind::OptionMut) => quote! {
679719
registry.add_async_method_mut(#lua_name, #closure_params #body);
680720
},
721+
SelfKind::Ref(_) => quote! {
722+
registry.add_async_method(#lua_name, #closure_params #body);
723+
},
681724
SelfKind::Owned => quote! {
682725
registry.add_async_method_once(#lua_name, #closure_params #body);
683726
},

‎tests/userdata_macro.rs‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,18 @@ impl Rectangle {
7979
}
8080
}
8181

82+
#[lua(infallible)]
83+
fn maybe_add(&self, other: Option<&Rectangle>) -> Rectangle {
84+
match other {
85+
Some(other) => Rectangle {
86+
length: self.length + other.length,
87+
width: self.width + other.width,
88+
..Default::default()
89+
},
90+
None => self.clone(),
91+
}
92+
}
93+
8294
#[lua(meta, field, name = "__answer")]
8395
fn answer() -> u32 {
8496
42
@@ -113,6 +125,13 @@ impl Rectangle {
113125
Ok(format!("Hello, {name}!"))
114126
}
115127

128+
fn maybe_greet(&self, name: Option<&str>) -> Result<String> {
129+
match name {
130+
Some(name) => Ok(format!("Hello, {name}!")),
131+
None => Ok("Hello!".to_string()),
132+
}
133+
}
134+
116135
fn transfer_length(&mut self, other: &mut Rectangle) -> Result<()> {
117136
other.length += self.length;
118137
self.length = 0;
@@ -190,6 +209,14 @@ fn test_rectangle() {
190209
assert(r4.length == 8, "__add length should be 5 + 3 = 8")
191210
assert(r4.width == 14, "__add width should be 10 + 4 = 14")
192211
212+
-- Option<&T> wrapped parameter
213+
local r5 = r1:maybe_add(r3)
214+
assert(r5.length == 8, "maybe_add with arg length should be 5 + 3 = 8")
215+
assert(r5.width == 14, "maybe_add with arg width should be 10 + 4 = 14")
216+
local r6 = r1:maybe_add()
217+
assert(r6.length == 5, "maybe_add with nil is no-op")
218+
assert(r6.width == 10, "maybe_add with nil is no-op")
219+
193220
-- method with &mut self and &mut Rectangle param
194221
rect = Rectangle.new(5, 10, 3)
195222
other = Rectangle.new(2, 3, 0)
@@ -212,6 +239,11 @@ fn test_rectangle() {
212239
assert(h == 10, "into_tuple height should be 7")
213240
local ok, err = pcall(function() rect:area() end)
214241
assert(not ok and tostring(err):match("userdata has been destructed"), "rect should be consumed and unusable after into_tuple")
242+
243+
-- Custom methods
244+
assert(other:greet("User") == "Hello, User!", "greet should return 'Hello, User!'")
245+
assert(other:maybe_greet("User") == "Hello, User!", "maybe_greet with arg should return 'Hello, User!'")
246+
assert(other:maybe_greet() == "Hello!", "maybe_greet with nil should return 'Hello!'")
215247
"#,
216248
)
217249
.exec()

0 commit comments

Comments
 (0)