Add optional __namecall optimization for Luau

Add `UserDataRegistry::enable_namecall()` hint to set `__namecall` metamethod to enable Luau-specific method resolution optimization.

Author: khvzak

benches/benchmark.rs

@@ -376,6 +376,9 @@ fn userdata_call_method_complex(c: &mut Criterion) {
                 this.0 += by;
                 Ok(this.0)
             });
+
+            #[cfg(feature = "luau")]
+            registry.enable_namecall();
         }
     }
 

src/state/raw.rs

@@ -28,8 +28,8 @@ use crate::userdata::{
 use crate::util::{
     assert_stack, check_stack, get_destructed_userdata_metatable, get_internal_userdata, get_main_state,
     get_metatable_ptr, get_userdata, init_error_registry, init_internal_metatable, pop_error,
-    push_internal_userdata, push_string, push_table, rawset_field, safe_pcall, safe_xpcall, short_type_name,
-    StackGuard, WrappedFailure,
+    push_internal_userdata, push_string, push_table, push_userdata, rawset_field, safe_pcall, safe_xpcall,
+    short_type_name, StackGuard, WrappedFailure,
 };
 use crate::value::{Nil, Value};
 
@@ -928,7 +928,7 @@ impl RawLua {
         // We generate metatable first to make sure it *always* available when userdata pushed
         let mt_id = get_metatable_id()?;
         let protect = !self.unlikely_memory_error();
-        crate::util::push_userdata(state, data, protect)?;
+        push_userdata(state, data, protect)?;
         ffi::lua_rawgeti(state, ffi::LUA_REGISTRYINDEX, mt_id);
         ffi::lua_setmetatable(state, -2);
 
@@ -1056,6 +1056,18 @@ impl RawLua {
             field_setters_index = Some(ffi::lua_absindex(state, -1));
         }
 
+        // Create methods namecall table
+        #[cfg_attr(not(feature = "luau"), allow(unused_mut))]
+        let mut methods_map = None;
+        #[cfg(feature = "luau")]
+        if registry.enable_namecall {
+            let map: &mut rustc_hash::FxHashMap<_, crate::types::CallbackPtr> =
+                methods_map.get_or_insert_with(Default::default);
+            for (k, m) in &registry.methods {
+                map.insert(k.as_bytes().to_vec(), &**m);
+            }
+        }
+
         let mut methods_index = None;
         let methods_nrec = registry.methods.len();
         #[cfg(feature = "async")]
@@ -1103,6 +1115,7 @@ impl RawLua {
             field_getters_index,
             field_setters_index,
             methods_index,
+            methods_map,
         )?;
 
         // Update stack guard to keep metatable after return

src/types.rs

@@ -38,10 +38,13 @@ unsafe impl Send for LightUserData {}
 unsafe impl Sync for LightUserData {}
 
 #[cfg(feature = "send")]
-pub(crate) type Callback = Box<dyn Fn(&RawLua, c_int) -> Result<c_int> + Send + 'static>;
+type CallbackFn<'a> = dyn Fn(&RawLua, c_int) -> Result<c_int> + Send + 'a;
 
 #[cfg(not(feature = "send"))]
-pub(crate) type Callback = Box<dyn Fn(&RawLua, c_int) -> Result<c_int> + 'static>;
+type CallbackFn<'a> = dyn Fn(&RawLua, c_int) -> Result<c_int> + 'a;
+
+pub(crate) type Callback = Box<CallbackFn<'static>>;
+pub(crate) type CallbackPtr = *const CallbackFn<'static>;
 
 pub(crate) type ScopedCallback<'s> = Box<dyn Fn(&RawLua, c_int) -> Result<c_int> + 's>;
 

src/userdata/registry.rs

@@ -56,6 +56,9 @@ pub(crate) struct RawUserDataRegistry {
     pub(crate) destructor: ffi::lua_CFunction,
     pub(crate) type_id: Option<TypeId>,
     pub(crate) type_name: StdString,
+
+    #[cfg(feature = "luau")]
+    pub(crate) enable_namecall: bool,
 }
 
 impl UserDataType {
@@ -100,6 +103,8 @@ impl<T> UserDataRegistry<T> {
             destructor: super::util::destroy_userdata_storage::<T>,
             type_id: r#type.type_id(),
             type_name: short_type_name::<T>(),
+            #[cfg(feature = "luau")]
+            enable_namecall: false,
         };
 
         UserDataRegistry {
@@ -110,6 +115,23 @@ impl<T> UserDataRegistry<T> {
         }
     }
 
+    /// Enables support for the namecall optimization in Luau.
+    ///
+    /// This enables methods resolution optimization in Luau for complex userdata types with methods
+    /// and field getters. When enabled, Luau will use a faster lookup path for method calls when a
+    /// specific syntax is used (e.g. `obj:method()`.
+    ///
+    /// This optimization does not play well with async methods, custom `__index` metamethod and
+    /// field getters as functions. So, it is disabled by default.
+    ///
+    /// Use with caution.
+    #[doc(hidden)]
+    #[cfg(feature = "luau")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
+    pub fn enable_namecall(&mut self) {
+        self.raw.enable_namecall = true;
+    }
+
     fn box_method<M, A, R>(&self, name: &str, method: M) -> Callback
     where
         M: Fn(&Lua, &T, A) -> Result<R> + MaybeSend + 'static,

src/userdata/util.rs

@@ -4,8 +4,11 @@ use std::marker::PhantomData;
 use std::os::raw::c_int;
 use std::ptr;
 
+use rustc_hash::FxHashMap;
+
 use super::UserDataStorage;
 use crate::error::{Error, Result};
+use crate::types::CallbackPtr;
 use crate::util::{get_userdata, rawget_field, rawset_field, take_userdata};
 
 // This is a trick to check if a type is `Sync` or not.
@@ -244,6 +247,7 @@ pub(crate) unsafe fn init_userdata_metatable(
     field_getters: Option<c_int>,
     field_setters: Option<c_int>,
     methods: Option<c_int>,
+    _methods_map: Option<FxHashMap<Vec<u8>, CallbackPtr>>, // Used only in Luau for `__namecall`
 ) -> Result<()> {
     if field_getters.is_some() || methods.is_some() {
         // Push `__index` generator function
@@ -267,6 +271,13 @@ pub(crate) unsafe fn init_userdata_metatable(
         }
 
         rawset_field(state, metatable, "__index")?;
+
+        #[cfg(feature = "luau")]
+        if let Some(methods_map) = _methods_map {
+            // In Luau we can speedup method calls by providing a dedicated `__namecall` metamethod
+            push_userdata_metatable_namecall(state, methods_map)?;
+            rawset_field(state, metatable, "__namecall")?;
+        }
     }
 
     if let Some(field_setters) = field_setters {
@@ -425,6 +436,36 @@ unsafe fn init_userdata_metatable_newindex(state: *mut ffi::lua_State) -> Result
     })
 }
 
+#[cfg(feature = "luau")]
+unsafe fn push_userdata_metatable_namecall(
+    state: *mut ffi::lua_State,
+    methods_map: FxHashMap<Vec<u8>, CallbackPtr>,
+) -> Result<()> {
+    unsafe extern "C-unwind" fn namecall(state: *mut ffi::lua_State) -> c_int {
+        let name = ffi::lua_namecallatom(state, ptr::null_mut());
+        if name.is_null() {
+            ffi::luaL_error(state, cstr!("attempt to call an unknown method"));
+        }
+        let name_cs = std::ffi::CStr::from_ptr(name);
+        let methods_map = get_userdata::<FxHashMap<Vec<u8>, CallbackPtr>>(state, ffi::lua_upvalueindex(1));
+        let callback_ptr = match (*methods_map).get(name_cs.to_bytes()) {
+            Some(ptr) => *ptr,
+            #[rustfmt::skip]
+            None => ffi::luaL_error(state, cstr!("attempt to call an unknown method '%s'"), name),
+        };
+        crate::state::callback_error_ext(state, ptr::null_mut(), true, |extra, nargs| {
+            let rawlua = (*extra).raw_lua();
+            (*callback_ptr)(rawlua, nargs)
+        })
+    }
+
+    // Automatic destructor is provided for any Luau userdata
+    crate::util::push_userdata(state, methods_map, true)?;
+    protect_lua!(state, 1, 1, |state| {
+        ffi::lua_pushcclosured(state, namecall, cstr!("__namecall"), 1);
+    })
+}
+
 // This method is called by Lua GC when it's time to collect the userdata.
 //
 // This method is usually used to collect internal userdata.

src/util/error.rs

@@ -402,6 +402,8 @@ pub(crate) unsafe fn init_error_registry(state: *mut ffi::lua_State) -> Result<(
         "__ipairs",
         #[cfg(feature = "luau")]
         "__iter",
+        #[cfg(feature = "luau")]
+        "__namecall",
         #[cfg(feature = "lua54")]
         "__close",
     ] {

tests/userdata.rs

@@ -1307,3 +1307,41 @@ fn test_userdata_wrappers() -> Result<()> {
 
     Ok(())
 }
+
+#[cfg(feature = "luau")]
+#[test]
+fn test_userdata_namecall() -> Result<()> {
+    let lua = Lua::new();
+
+    struct MyUserData;
+
+    impl UserData for MyUserData {
+        fn register(registry: &mut mlua::UserDataRegistry<Self>) {
+            registry.add_method("method", |_, _, ()| Ok("method called"));
+            registry.add_field_method_get("field", |_, _| Ok("field value"));
+
+            registry.add_meta_method(MetaMethod::Index, |_, _, key: StdString| Ok(key));
+
+            registry.enable_namecall();
+        }
+    }
+
+    let ud = lua.create_userdata(MyUserData)?;
+    lua.globals().set("ud", &ud)?;
+    lua.load(
+        r#"
+        assert(ud:method() == "method called")
+        assert(ud.field == "field value")
+        assert(ud.dynamic_field == "dynamic_field")
+        local ok, err = pcall(function() return ud:dynamic_field() end)
+        assert(tostring(err):find("attempt to call an unknown method 'dynamic_field'") ~= nil)
+        "#,
+    )
+    .exec()?;
+
+    ud.destroy()?;
+    let err = lua.load("ud:method()").exec().unwrap_err();
+    assert!(err.to_string().contains("userdata has been destructed"));
+
+    Ok(())
+}