Optimize metatable pointer lookup for userdata (Luau)

Author: khvzak

Cargo.toml

@@ -55,7 +55,7 @@ serde-value = { version = "0.7", optional = true }
 parking_lot = { version = "0.12", features = ["arc_lock"] }
 anyhow = { version = "1.0", optional = true }
 
-ffi = { package = "mlua-sys", version = "0.6.3", path = "mlua-sys" }
+ffi = { package = "mlua-sys", version = "0.6.4", path = "mlua-sys" }
 
 [target.'cfg(unix)'.dependencies]
 libloading = { version = "0.8", optional = true }

src/scope.rs

@@ -9,7 +9,9 @@ use crate::state::{Lua, LuaGuard, RawLua};
 use crate::traits::{FromLuaMulti, IntoLuaMulti};
 use crate::types::{Callback, CallbackUpvalue, ScopedCallback, ValueRef};
 use crate::userdata::{AnyUserData, UserData, UserDataRegistry, UserDataStorage};
-use crate::util::{self, assert_stack, check_stack, get_userdata, take_userdata, StackGuard};
+use crate::util::{
+    self, assert_stack, check_stack, get_metatable_ptr, get_userdata, take_userdata, StackGuard,
+};
 
 /// Constructed by the [`Lua::scope`] method, allows temporarily creating Lua userdata and
 /// callbacks that are not required to be `Send` or `'static`.
@@ -199,9 +201,7 @@ impl<'scope, 'env: 'scope> Scope<'scope, 'env> {
                 }
 
                 // Deregister metatable
-                ffi::lua_getmetatable(state, -1);
-                let mt_ptr = ffi::lua_topointer(state, -1);
-                ffi::lua_pop(state, 1);
+                let mt_ptr = get_metatable_ptr(state, -1);
                 rawlua.deregister_userdata_metatable(mt_ptr);
 
                 let ud = take_userdata::<UserDataStorage<T>>(state);

src/state/raw.rs

@@ -24,9 +24,9 @@ use crate::types::{
 use crate::userdata::{AnyUserData, MetaMethod, UserData, UserDataRegistry, UserDataStorage};
 use crate::util::{
     assert_stack, check_stack, get_destructed_userdata_metatable, get_internal_userdata, get_main_state,
-    get_userdata, init_error_registry, init_internal_metatable, init_userdata_metatable, pop_error,
-    push_internal_userdata, push_string, push_table, rawset_field, safe_pcall, safe_xpcall, short_type_name,
-    StackGuard, WrappedFailure,
+    get_metatable_ptr, get_userdata, init_error_registry, init_internal_metatable, init_userdata_metatable,
+    pop_error, push_internal_userdata, push_string, push_table, rawset_field, safe_pcall, safe_xpcall,
+    short_type_name, StackGuard, WrappedFailure,
 };
 use crate::value::{Nil, Value};
 
@@ -1026,11 +1026,10 @@ impl RawLua {
         state: *mut ffi::lua_State,
         idx: c_int,
     ) -> Result<Option<TypeId>> {
-        if ffi::lua_getmetatable(state, idx) == 0 {
+        let mt_ptr = get_metatable_ptr(state, idx);
+        if mt_ptr.is_null() {
             return Err(Error::UserDataTypeMismatch);
         }
-        let mt_ptr = ffi::lua_topointer(state, -1);
-        ffi::lua_pop(state, 1);
 
         // Fast path to skip looking up the metatable in the map
         let (last_mt, last_type_id) = (*self.extra.get()).last_checked_userdata_mt;

src/util/mod.rs

@@ -1,6 +1,6 @@
 use std::borrow::Cow;
 use std::ffi::CStr;
-use std::os::raw::{c_char, c_int};
+use std::os::raw::{c_char, c_int, c_void};
 use std::{ptr, slice, str};
 
 use crate::error::{Error, Result};
@@ -282,6 +282,21 @@ pub(crate) unsafe fn to_string(state: *mut ffi::lua_State, index: c_int) -> Stri
     }
 }
 
+#[inline(always)]
+pub(crate) unsafe fn get_metatable_ptr(state: *mut ffi::lua_State, index: c_int) -> *const c_void {
+    #[cfg(feature = "luau")]
+    return ffi::lua_getmetatablepointer(state, index);
+
+    #[cfg(not(feature = "luau"))]
+    if ffi::lua_getmetatable(state, index) == 0 {
+        ptr::null()
+    } else {
+        let p = ffi::lua_topointer(state, -1);
+        ffi::lua_pop(state, 1);
+        p
+    }
+}
+
 pub(crate) unsafe fn ptr_to_str<'a>(input: *const c_char) -> Option<&'a str> {
     if input.is_null() {
         return None;

src/util/userdata.rs

@@ -3,7 +3,7 @@ use std::os::raw::{c_int, c_void};
 use std::{ptr, str};
 
 use crate::error::Result;
-use crate::util::{check_stack, push_string, push_table, rawset_field, TypeKey};
+use crate::util::{check_stack, get_metatable_ptr, push_string, push_table, rawset_field, TypeKey};
 
 // Pushes the userdata and attaches a metatable with __gc method.
 // Internally uses 3 stack spaces, does not call checkstack.
@@ -58,25 +58,20 @@ pub(crate) unsafe fn init_internal_metatable<T: TypeKey>(
 pub(crate) unsafe fn get_internal_userdata<T: TypeKey>(
     state: *mut ffi::lua_State,
     index: c_int,
-    type_mt_ptr: *const c_void,
+    mut type_mt_ptr: *const c_void,
 ) -> *mut T {
     let ud = ffi::lua_touserdata(state, index) as *mut T;
-    if ud.is_null() || ffi::lua_getmetatable(state, index) == 0 {
+    if ud.is_null() {
         return ptr::null_mut();
     }
-    if !type_mt_ptr.is_null() {
-        let ud_mt_ptr = ffi::lua_topointer(state, -1);
-        ffi::lua_pop(state, 1);
-        if ud_mt_ptr != type_mt_ptr {
-            return ptr::null_mut();
-        }
-    } else {
+    let mt_ptr = get_metatable_ptr(state, index);
+    if type_mt_ptr.is_null() {
         get_internal_metatable::<T>(state);
-        let res = ffi::lua_rawequal(state, -1, -2);
-        ffi::lua_pop(state, 2);
-        if res == 0 {
-            return ptr::null_mut();
-        }
+        type_mt_ptr = ffi::lua_topointer(state, -1);
+        ffi::lua_pop(state, 1);
+    }
+    if mt_ptr != type_mt_ptr {
+        return ptr::null_mut();
     }
     ud
 }