mlua_codemp_patch/userdata/
registry.rs

1#![allow(clippy::await_holding_refcell_ref, clippy::await_holding_lock)]
2
3use std::any::TypeId;
4use std::cell::RefCell;
5use std::marker::PhantomData;
6use std::os::raw::c_int;
7use std::string::String as StdString;
8
9use crate::error::{Error, Result};
10use crate::state::{Lua, RawLua};
11use crate::types::{Callback, MaybeSend};
12use crate::userdata::{
13    AnyUserData, MetaMethod, UserData, UserDataFields, UserDataMethods, UserDataRef, UserDataRefMut,
14};
15use crate::util::{get_userdata, short_type_name};
16use crate::value::{FromLua, FromLuaMulti, IntoLua, IntoLuaMulti, Value};
17
18use super::cell::{UserDataBorrowMut, UserDataBorrowRef, UserDataVariant};
19
20#[cfg(feature = "async")]
21use {
22    crate::types::AsyncCallback,
23    std::future::{self, Future},
24};
25
26type StaticFieldCallback = Box<dyn FnOnce(&RawLua) -> Result<()> + 'static>;
27
28/// Handle to registry for userdata methods and metamethods.
29pub struct UserDataRegistry<T: 'static> {
30    // Fields
31    pub(crate) fields: Vec<(String, StaticFieldCallback)>,
32    pub(crate) field_getters: Vec<(String, Callback)>,
33    pub(crate) field_setters: Vec<(String, Callback)>,
34    pub(crate) meta_fields: Vec<(String, StaticFieldCallback)>,
35
36    // Methods
37    pub(crate) methods: Vec<(String, Callback)>,
38    #[cfg(feature = "async")]
39    pub(crate) async_methods: Vec<(String, AsyncCallback)>,
40    pub(crate) meta_methods: Vec<(String, Callback)>,
41    #[cfg(feature = "async")]
42    pub(crate) async_meta_methods: Vec<(String, AsyncCallback)>,
43
44    _type: PhantomData<T>,
45}
46
47impl<T: 'static> UserDataRegistry<T> {
48    pub(crate) const fn new() -> Self {
49        UserDataRegistry {
50            fields: Vec::new(),
51            field_getters: Vec::new(),
52            field_setters: Vec::new(),
53            meta_fields: Vec::new(),
54            methods: Vec::new(),
55            #[cfg(feature = "async")]
56            async_methods: Vec::new(),
57            meta_methods: Vec::new(),
58            #[cfg(feature = "async")]
59            async_meta_methods: Vec::new(),
60            _type: PhantomData,
61        }
62    }
63
64    fn box_method<M, A, R>(name: &str, method: M) -> Callback
65    where
66        M: Fn(&Lua, &T, A) -> Result<R> + MaybeSend + 'static,
67        A: FromLuaMulti,
68        R: IntoLuaMulti,
69    {
70        let name = get_function_name::<T>(name);
71        macro_rules! try_self_arg {
72            ($res:expr) => {
73                $res.map_err(|err| Error::bad_self_argument(&name, err))?
74            };
75        }
76
77        Box::new(move |rawlua, nargs| unsafe {
78            if nargs == 0 {
79                let err = Error::from_lua_conversion("missing argument", "userdata", None);
80                try_self_arg!(Err(err));
81            }
82            let state = rawlua.state();
83            // Find absolute "self" index before processing args
84            let self_index = ffi::lua_absindex(state, -nargs);
85            // Self was at position 1, so we pass 2 here
86            let args = A::from_stack_args(nargs - 1, 2, Some(&name), rawlua);
87
88            match try_self_arg!(rawlua.get_userdata_type_id(self_index)) {
89                Some(id) if id == TypeId::of::<T>() => {
90                    let ud = try_self_arg!(borrow_userdata_ref::<T>(state, self_index));
91                    method(rawlua.lua(), &ud, args?)?.push_into_stack_multi(rawlua)
92                }
93                _ => Err(Error::bad_self_argument(&name, Error::UserDataTypeMismatch)),
94            }
95        })
96    }
97
98    fn box_method_mut<M, A, R>(name: &str, method: M) -> Callback
99    where
100        M: FnMut(&Lua, &mut T, A) -> Result<R> + MaybeSend + 'static,
101        A: FromLuaMulti,
102        R: IntoLuaMulti,
103    {
104        let name = get_function_name::<T>(name);
105        macro_rules! try_self_arg {
106            ($res:expr) => {
107                $res.map_err(|err| Error::bad_self_argument(&name, err))?
108            };
109        }
110
111        let method = RefCell::new(method);
112        Box::new(move |rawlua, nargs| unsafe {
113            let mut method = method.try_borrow_mut().map_err(|_| Error::RecursiveMutCallback)?;
114            if nargs == 0 {
115                let err = Error::from_lua_conversion("missing argument", "userdata", None);
116                try_self_arg!(Err(err));
117            }
118            let state = rawlua.state();
119            // Find absolute "self" index before processing args
120            let self_index = ffi::lua_absindex(state, -nargs);
121            // Self was at position 1, so we pass 2 here
122            let args = A::from_stack_args(nargs - 1, 2, Some(&name), rawlua);
123
124            match try_self_arg!(rawlua.get_userdata_type_id(self_index)) {
125                Some(id) if id == TypeId::of::<T>() => {
126                    let mut ud = try_self_arg!(borrow_userdata_mut::<T>(state, self_index));
127                    method(rawlua.lua(), &mut ud, args?)?.push_into_stack_multi(rawlua)
128                }
129                _ => Err(Error::bad_self_argument(&name, Error::UserDataTypeMismatch)),
130            }
131        })
132    }
133
134    #[cfg(feature = "async")]
135    fn box_async_method<M, A, MR, R>(name: &str, method: M) -> AsyncCallback
136    where
137        M: Fn(Lua, UserDataRef<T>, A) -> MR + MaybeSend + 'static,
138        A: FromLuaMulti,
139        MR: Future<Output = Result<R>> + MaybeSend + 'static,
140        R: IntoLuaMulti,
141    {
142        let name = get_function_name::<T>(name);
143        macro_rules! try_self_arg {
144            ($res:expr) => {
145                match $res {
146                    Ok(res) => res,
147                    Err(err) => return Box::pin(future::ready(Err(Error::bad_self_argument(&name, err)))),
148                }
149            };
150        }
151
152        Box::new(move |rawlua, nargs| unsafe {
153            if nargs == 0 {
154                let err = Error::from_lua_conversion("missing argument", "userdata", None);
155                try_self_arg!(Err(err));
156            }
157            // Stack will be empty when polling the future, keep `self` on the ref thread
158            let self_ud = try_self_arg!(AnyUserData::from_stack(-nargs, rawlua));
159            let args = A::from_stack_args(nargs - 1, 2, Some(&name), rawlua);
160
161            let self_ud = try_self_arg!(self_ud.borrow());
162            let args = match args {
163                Ok(args) => args,
164                Err(e) => return Box::pin(future::ready(Err(e))),
165            };
166            let lua = rawlua.lua();
167            let fut = method(lua.clone(), self_ud, args);
168            // Lua is locked when the future is polled
169            Box::pin(async move { fut.await?.push_into_stack_multi(lua.raw_lua()) })
170        })
171    }
172
173    #[cfg(feature = "async")]
174    fn box_async_method_mut<M, A, MR, R>(name: &str, method: M) -> AsyncCallback
175    where
176        M: Fn(Lua, UserDataRefMut<T>, A) -> MR + MaybeSend + 'static,
177        A: FromLuaMulti,
178        MR: Future<Output = Result<R>> + MaybeSend + 'static,
179        R: IntoLuaMulti,
180    {
181        let name = get_function_name::<T>(name);
182        macro_rules! try_self_arg {
183            ($res:expr) => {
184                match $res {
185                    Ok(res) => res,
186                    Err(err) => return Box::pin(future::ready(Err(Error::bad_self_argument(&name, err)))),
187                }
188            };
189        }
190
191        Box::new(move |rawlua, nargs| unsafe {
192            if nargs == 0 {
193                let err = Error::from_lua_conversion("missing argument", "userdata", None);
194                try_self_arg!(Err(err));
195            }
196            // Stack will be empty when polling the future, keep `self` on the ref thread
197            let self_ud = try_self_arg!(AnyUserData::from_stack(-nargs, rawlua));
198            let args = A::from_stack_args(nargs - 1, 2, Some(&name), rawlua);
199
200            let self_ud = try_self_arg!(self_ud.borrow_mut());
201            let args = match args {
202                Ok(args) => args,
203                Err(e) => return Box::pin(future::ready(Err(e))),
204            };
205            let lua = rawlua.lua();
206            let fut = method(lua.clone(), self_ud, args);
207            // Lua is locked when the future is polled
208            Box::pin(async move { fut.await?.push_into_stack_multi(lua.raw_lua()) })
209        })
210    }
211
212    fn box_function<F, A, R>(name: &str, function: F) -> Callback
213    where
214        F: Fn(&Lua, A) -> Result<R> + MaybeSend + 'static,
215        A: FromLuaMulti,
216        R: IntoLuaMulti,
217    {
218        let name = get_function_name::<T>(name);
219        Box::new(move |lua, nargs| unsafe {
220            let args = A::from_stack_args(nargs, 1, Some(&name), lua)?;
221            function(lua.lua(), args)?.push_into_stack_multi(lua)
222        })
223    }
224
225    fn box_function_mut<F, A, R>(name: &str, function: F) -> Callback
226    where
227        F: FnMut(&Lua, A) -> Result<R> + MaybeSend + 'static,
228        A: FromLuaMulti,
229        R: IntoLuaMulti,
230    {
231        let name = get_function_name::<T>(name);
232        let function = RefCell::new(function);
233        Box::new(move |lua, nargs| unsafe {
234            let function = &mut *function
235                .try_borrow_mut()
236                .map_err(|_| Error::RecursiveMutCallback)?;
237            let args = A::from_stack_args(nargs, 1, Some(&name), lua)?;
238            function(lua.lua(), args)?.push_into_stack_multi(lua)
239        })
240    }
241
242    #[cfg(feature = "async")]
243    fn box_async_function<F, A, FR, R>(name: &str, function: F) -> AsyncCallback
244    where
245        F: Fn(Lua, A) -> FR + MaybeSend + 'static,
246        A: FromLuaMulti,
247        FR: Future<Output = Result<R>> + MaybeSend + 'static,
248        R: IntoLuaMulti,
249    {
250        let name = get_function_name::<T>(name);
251        Box::new(move |rawlua, nargs| unsafe {
252            let args = match A::from_stack_args(nargs, 1, Some(&name), rawlua) {
253                Ok(args) => args,
254                Err(e) => return Box::pin(future::ready(Err(e))),
255            };
256            let lua = rawlua.lua();
257            let fut = function(lua.clone(), args);
258            Box::pin(async move { fut.await?.push_into_stack_multi(lua.raw_lua()) })
259        })
260    }
261
262    pub(crate) fn check_meta_field(lua: &Lua, name: &str, value: impl IntoLua) -> Result<Value> {
263        let value = value.into_lua(lua)?;
264        if name == MetaMethod::Index || name == MetaMethod::NewIndex {
265            match value {
266                Value::Nil | Value::Table(_) | Value::Function(_) => {}
267                _ => {
268                    return Err(Error::MetaMethodTypeError {
269                        method: name.to_string(),
270                        type_name: value.type_name(),
271                        message: Some("expected nil, table or function".to_string()),
272                    })
273                }
274            }
275        }
276        value.into_lua(lua)
277    }
278}
279
280// Returns function name for the type `T`, without the module path
281fn get_function_name<T>(name: &str) -> StdString {
282    format!("{}.{name}", short_type_name::<T>())
283}
284
285impl<T: 'static> UserDataFields<T> for UserDataRegistry<T> {
286    fn add_field<V>(&mut self, name: impl ToString, value: V)
287    where
288        V: IntoLua + 'static,
289    {
290        let name = name.to_string();
291        self.fields.push((
292            name,
293            Box::new(move |rawlua| unsafe { value.push_into_stack(rawlua) }),
294        ));
295    }
296
297    fn add_field_method_get<M, R>(&mut self, name: impl ToString, method: M)
298    where
299        M: Fn(&Lua, &T) -> Result<R> + MaybeSend + 'static,
300        R: IntoLua,
301    {
302        let name = name.to_string();
303        let callback = Self::box_method(&name, move |lua, data, ()| method(lua, data));
304        self.field_getters.push((name, callback));
305    }
306
307    fn add_field_method_set<M, A>(&mut self, name: impl ToString, method: M)
308    where
309        M: FnMut(&Lua, &mut T, A) -> Result<()> + MaybeSend + 'static,
310        A: FromLua,
311    {
312        let name = name.to_string();
313        let callback = Self::box_method_mut(&name, method);
314        self.field_setters.push((name, callback));
315    }
316
317    fn add_field_function_get<F, R>(&mut self, name: impl ToString, function: F)
318    where
319        F: Fn(&Lua, AnyUserData) -> Result<R> + MaybeSend + 'static,
320        R: IntoLua,
321    {
322        let name = name.to_string();
323        let callback = Self::box_function(&name, function);
324        self.field_getters.push((name, callback));
325    }
326
327    fn add_field_function_set<F, A>(&mut self, name: impl ToString, mut function: F)
328    where
329        F: FnMut(&Lua, AnyUserData, A) -> Result<()> + MaybeSend + 'static,
330        A: FromLua,
331    {
332        let name = name.to_string();
333        let callback = Self::box_function_mut(&name, move |lua, (data, val)| function(lua, data, val));
334        self.field_setters.push((name, callback));
335    }
336
337    fn add_meta_field<V>(&mut self, name: impl ToString, value: V)
338    where
339        V: IntoLua + 'static,
340    {
341        let name = name.to_string();
342        self.meta_fields.push((
343            name.clone(),
344            Box::new(move |rawlua| unsafe {
345                Self::check_meta_field(rawlua.lua(), &name, value)?.push_into_stack(rawlua)
346            }),
347        ));
348    }
349
350    fn add_meta_field_with<F, R>(&mut self, name: impl ToString, f: F)
351    where
352        F: FnOnce(&Lua) -> Result<R> + 'static,
353        R: IntoLua,
354    {
355        let name = name.to_string();
356        self.meta_fields.push((
357            name.clone(),
358            Box::new(move |rawlua| unsafe {
359                let lua = rawlua.lua();
360                Self::check_meta_field(lua, &name, f(lua)?)?.push_into_stack(rawlua)
361            }),
362        ));
363    }
364}
365
366impl<T: 'static> UserDataMethods<T> for UserDataRegistry<T> {
367    fn add_method<M, A, R>(&mut self, name: impl ToString, method: M)
368    where
369        M: Fn(&Lua, &T, A) -> Result<R> + MaybeSend + 'static,
370        A: FromLuaMulti,
371        R: IntoLuaMulti,
372    {
373        let name = name.to_string();
374        let callback = Self::box_method(&name, method);
375        self.methods.push((name, callback));
376    }
377
378    fn add_method_mut<M, A, R>(&mut self, name: impl ToString, method: M)
379    where
380        M: FnMut(&Lua, &mut T, A) -> Result<R> + MaybeSend + 'static,
381        A: FromLuaMulti,
382        R: IntoLuaMulti,
383    {
384        let name = name.to_string();
385        let callback = Self::box_method_mut(&name, method);
386        self.methods.push((name, callback));
387    }
388
389    #[cfg(feature = "async")]
390    fn add_async_method<M, A, MR, R>(&mut self, name: impl ToString, method: M)
391    where
392        M: Fn(Lua, UserDataRef<T>, A) -> MR + MaybeSend + 'static,
393        A: FromLuaMulti,
394        MR: Future<Output = Result<R>> + MaybeSend + 'static,
395        R: IntoLuaMulti,
396    {
397        let name = name.to_string();
398        let callback = Self::box_async_method(&name, method);
399        self.async_methods.push((name, callback));
400    }
401
402    #[cfg(feature = "async")]
403    fn add_async_method_mut<M, A, MR, R>(&mut self, name: impl ToString, method: M)
404    where
405        M: Fn(Lua, UserDataRefMut<T>, A) -> MR + MaybeSend + 'static,
406        A: FromLuaMulti,
407        MR: Future<Output = Result<R>> + MaybeSend + 'static,
408        R: IntoLuaMulti,
409    {
410        let name = name.to_string();
411        let callback = Self::box_async_method_mut(&name, method);
412        self.async_methods.push((name, callback));
413    }
414
415    fn add_function<F, A, R>(&mut self, name: impl ToString, function: F)
416    where
417        F: Fn(&Lua, A) -> Result<R> + MaybeSend + 'static,
418        A: FromLuaMulti,
419        R: IntoLuaMulti,
420    {
421        let name = name.to_string();
422        let callback = Self::box_function(&name, function);
423        self.methods.push((name, callback));
424    }
425
426    fn add_function_mut<F, A, R>(&mut self, name: impl ToString, function: F)
427    where
428        F: FnMut(&Lua, A) -> Result<R> + MaybeSend + 'static,
429        A: FromLuaMulti,
430        R: IntoLuaMulti,
431    {
432        let name = name.to_string();
433        let callback = Self::box_function_mut(&name, function);
434        self.methods.push((name, callback));
435    }
436
437    #[cfg(feature = "async")]
438    fn add_async_function<F, A, FR, R>(&mut self, name: impl ToString, function: F)
439    where
440        F: Fn(Lua, A) -> FR + MaybeSend + 'static,
441        A: FromLuaMulti,
442        FR: Future<Output = Result<R>> + MaybeSend + 'static,
443        R: IntoLuaMulti,
444    {
445        let name = name.to_string();
446        let callback = Self::box_async_function(&name, function);
447        self.async_methods.push((name, callback));
448    }
449
450    fn add_meta_method<M, A, R>(&mut self, name: impl ToString, method: M)
451    where
452        M: Fn(&Lua, &T, A) -> Result<R> + MaybeSend + 'static,
453        A: FromLuaMulti,
454        R: IntoLuaMulti,
455    {
456        let name = name.to_string();
457        let callback = Self::box_method(&name, method);
458        self.meta_methods.push((name, callback));
459    }
460
461    fn add_meta_method_mut<M, A, R>(&mut self, name: impl ToString, method: M)
462    where
463        M: FnMut(&Lua, &mut T, A) -> Result<R> + MaybeSend + 'static,
464        A: FromLuaMulti,
465        R: IntoLuaMulti,
466    {
467        let name = name.to_string();
468        let callback = Self::box_method_mut(&name, method);
469        self.meta_methods.push((name, callback));
470    }
471
472    #[cfg(all(feature = "async", not(any(feature = "lua51", feature = "luau"))))]
473    fn add_async_meta_method<M, A, MR, R>(&mut self, name: impl ToString, method: M)
474    where
475        M: Fn(Lua, UserDataRef<T>, A) -> MR + MaybeSend + 'static,
476        A: FromLuaMulti,
477        MR: Future<Output = Result<R>> + MaybeSend + 'static,
478        R: IntoLuaMulti,
479    {
480        let name = name.to_string();
481        let callback = Self::box_async_method(&name, method);
482        self.async_meta_methods.push((name, callback));
483    }
484
485    #[cfg(all(feature = "async", not(any(feature = "lua51", feature = "luau"))))]
486    fn add_async_meta_method_mut<M, A, MR, R>(&mut self, name: impl ToString, method: M)
487    where
488        M: Fn(Lua, UserDataRefMut<T>, A) -> MR + MaybeSend + 'static,
489        A: FromLuaMulti,
490        MR: Future<Output = Result<R>> + MaybeSend + 'static,
491        R: IntoLuaMulti,
492    {
493        let name = name.to_string();
494        let callback = Self::box_async_method_mut(&name, method);
495        self.async_meta_methods.push((name, callback));
496    }
497
498    fn add_meta_function<F, A, R>(&mut self, name: impl ToString, function: F)
499    where
500        F: Fn(&Lua, A) -> Result<R> + MaybeSend + 'static,
501        A: FromLuaMulti,
502        R: IntoLuaMulti,
503    {
504        let name = name.to_string();
505        let callback = Self::box_function(&name, function);
506        self.meta_methods.push((name, callback));
507    }
508
509    fn add_meta_function_mut<F, A, R>(&mut self, name: impl ToString, function: F)
510    where
511        F: FnMut(&Lua, A) -> Result<R> + MaybeSend + 'static,
512        A: FromLuaMulti,
513        R: IntoLuaMulti,
514    {
515        let name = name.to_string();
516        let callback = Self::box_function_mut(&name, function);
517        self.meta_methods.push((name, callback));
518    }
519
520    #[cfg(all(feature = "async", not(any(feature = "lua51", feature = "luau"))))]
521    fn add_async_meta_function<F, A, FR, R>(&mut self, name: impl ToString, function: F)
522    where
523        F: Fn(Lua, A) -> FR + MaybeSend + 'static,
524        A: FromLuaMulti,
525        FR: Future<Output = Result<R>> + MaybeSend + 'static,
526        R: IntoLuaMulti,
527    {
528        let name = name.to_string();
529        let callback = Self::box_async_function(&name, function);
530        self.async_meta_methods.push((name, callback));
531    }
532}
533
534// Borrow the userdata in-place from the Lua stack
535#[inline(always)]
536unsafe fn borrow_userdata_ref<'a, T>(
537    state: *mut ffi::lua_State,
538    index: c_int,
539) -> Result<UserDataBorrowRef<'a, T>> {
540    let ud = get_userdata::<UserDataVariant<T>>(state, index);
541    (*ud).try_borrow()
542}
543
544// Borrow the userdata mutably in-place from the Lua stack
545#[inline(always)]
546unsafe fn borrow_userdata_mut<'a, T>(
547    state: *mut ffi::lua_State,
548    index: c_int,
549) -> Result<UserDataBorrowMut<'a, T>> {
550    let ud = get_userdata::<UserDataVariant<T>>(state, index);
551    (*ud).try_borrow_mut()
552}
553
554macro_rules! lua_userdata_impl {
555    ($type:ty) => {
556        impl<T: UserData + 'static> UserData for $type {
557            fn register(registry: &mut UserDataRegistry<Self>) {
558                let mut orig_registry = UserDataRegistry::new();
559                T::register(&mut orig_registry);
560
561                // Copy all fields, methods, etc. from the original registry
562                registry.fields.extend(orig_registry.fields);
563                registry.field_getters.extend(orig_registry.field_getters);
564                registry.field_setters.extend(orig_registry.field_setters);
565                registry.meta_fields.extend(orig_registry.meta_fields);
566                registry.methods.extend(orig_registry.methods);
567                #[cfg(feature = "async")]
568                registry.async_methods.extend(orig_registry.async_methods);
569                registry.meta_methods.extend(orig_registry.meta_methods);
570                #[cfg(feature = "async")]
571                registry
572                    .async_meta_methods
573                    .extend(orig_registry.async_meta_methods);
574            }
575        }
576    };
577}
578
579// A special proxy object for UserData
580pub(crate) struct UserDataProxy<T>(pub(crate) PhantomData<T>);
581
582lua_userdata_impl!(UserDataProxy<T>);