Delay "any" userdata metatable creation until first instance

khvzak · khvzak · commit 00e9dd689f1c · 2024-11-12T22:20:26.000Z
diff --git a/src/scope.rs b/src/scope.rs
@@ -174,7 +174,7 @@ impl<'scope, 'env: 'scope> Scope<'scope, 'env> {
             let ud_ptr = util::push_uninit_userdata::<UserDataStorage<T>>(state, protect)?;
 
             // Push the metatable and register it with no TypeId
-            let mut registry = UserDataRegistry::new_unique(ud_ptr as *mut _);
+            let mut registry = UserDataRegistry::new_unique(self.lua.lua(), ud_ptr as *mut _);
             T::register(&mut registry);
             self.lua.push_userdata_metatable(registry.into_raw())?;
             let mt_ptr = ffi::lua_topointer(state, -1);
@@ -222,7 +222,7 @@ impl<'scope, 'env: 'scope> Scope<'scope, 'env> {
             let ud_ptr = util::push_uninit_userdata::<UserDataStorage<T>>(state, protect)?;
 
             // Push the metatable and register it with no TypeId
-            let mut registry = UserDataRegistry::new_unique(ud_ptr as *mut _);
+            let mut registry = UserDataRegistry::new_unique(self.lua.lua(), ud_ptr as *mut _);
             register(&mut registry);
             self.lua.push_userdata_metatable(registry.into_raw())?;
             let mt_ptr = ffi::lua_topointer(state, -1);
diff --git a/src/state.rs b/src/state.rs
@@ -1306,7 +1306,7 @@ impl Lua {
     /// This methods provides a way to add fields or methods to userdata objects of a type `T`.
     pub fn register_userdata_type<T: 'static>(&self, f: impl FnOnce(&mut UserDataRegistry<T>)) -> Result<()> {
         let type_id = TypeId::of::<T>();
-        let mut registry = UserDataRegistry::new(type_id);
+        let mut registry = UserDataRegistry::new(self, type_id);
         f(&mut registry);
 
         let lua = self.lock();
@@ -1316,8 +1316,8 @@ impl Lua {
                 ffi::luaL_unref(lua.state(), ffi::LUA_REGISTRYINDEX, table_id);
             }
 
-            // Register the type
-            lua.create_userdata_metatable(registry.into_raw())?;
+            // Add to "pending" registration map
+            ((*lua.extra.get()).pending_userdata_reg).insert(type_id, registry.into_raw());
         }
         Ok(())
     }
diff --git a/src/state/extra.rs b/src/state/extra.rs
@@ -13,6 +13,7 @@ use crate::error::Result;
 use crate::state::RawLua;
 use crate::stdlib::StdLib;
 use crate::types::{AppData, ReentrantMutex, XRc};
+use crate::userdata::RawUserDataRegistry;
 use crate::util::{get_internal_metatable, push_internal_userdata, TypeKey, WrappedFailure};
 
 #[cfg(any(feature = "luau", doc))]
@@ -35,6 +36,7 @@ pub(crate) struct ExtraData {
     pub(super) weak: MaybeUninit<WeakLua>,
     pub(super) owned: bool,
 
+    pub(super) pending_userdata_reg: FxHashMap<TypeId, RawUserDataRegistry>,
     pub(super) registered_userdata_t: FxHashMap<TypeId, c_int>,
     pub(super) registered_userdata_mt: FxHashMap<*const c_void, Option<TypeId>>,
     pub(super) last_checked_userdata_mt: (*const c_void, Option<TypeId>),
@@ -144,6 +146,7 @@ impl ExtraData {
             lua: MaybeUninit::uninit(),
             weak: MaybeUninit::uninit(),
             owned,
+            pending_userdata_reg: FxHashMap::default(),
             registered_userdata_t: FxHashMap::default(),
             registered_userdata_mt: FxHashMap::default(),
             last_checked_userdata_mt: (ptr::null(), None),
diff --git a/src/state/raw.rs b/src/state/raw.rs
@@ -756,7 +756,7 @@ impl RawLua {
             }
 
             // Create a new metatable from `UserData` definition
-            let mut registry = UserDataRegistry::new(type_id);
+            let mut registry = UserDataRegistry::new(self.lua(), type_id);
             T::register(&mut registry);
 
             self.create_userdata_metatable(registry.into_raw())
@@ -774,9 +774,12 @@ impl RawLua {
                 return Ok(table_id as Integer);
             }
 
-            // Create an empty metatable
-            let registry = UserDataRegistry::<T>::new(type_id);
-            self.create_userdata_metatable(registry.into_raw())
+            // Check if metatable creation is pending or create an empty metatable otherwise
+            let registry = match (*self.extra.get()).pending_userdata_reg.remove(&type_id) {
+                Some(registry) => registry,
+                None => UserDataRegistry::<T>::new(self.lua(), type_id).into_raw(),
+            };
+            self.create_userdata_metatable(registry)
         })
     }
 
@@ -851,9 +854,9 @@ impl RawLua {
             rawset_field(state, -2, MetaMethod::validate(&k)?)?;
         }
         let mut has_name = false;
-        for (k, push_field) in registry.meta_fields {
+        for (k, v) in registry.meta_fields {
             has_name = has_name || k == MetaMethod::Type;
-            push_field(self)?;
+            v?.push_into_stack(self)?;
             rawset_field(state, -2, MetaMethod::validate(&k)?)?;
         }
         // Set `__name/__type` if not provided
@@ -875,8 +878,8 @@ impl RawLua {
                         ffi::lua_pop(state, 1);
                         push_table(state, 0, fields_nrec, true)?;
                     }
-                    for (k, push_field) in mem::take(&mut registry.fields) {
-                        push_field(self)?;
+                    for (k, v) in mem::take(&mut registry.fields) {
+                        v?.push_into_stack(self)?;
                         rawset_field(state, -2, &k)?;
                     }
                     rawset_field(state, metatable_index, "__index")?;
@@ -896,12 +899,12 @@ impl RawLua {
                 self.push(self.create_callback(m)?)?;
                 rawset_field(state, -2, &k)?;
             }
-            for (k, push_field) in registry.fields {
+            for (k, v) in registry.fields {
                 unsafe extern "C-unwind" fn return_field(state: *mut ffi::lua_State) -> c_int {
                     ffi::lua_pushvalue(state, ffi::lua_upvalueindex(1));
                     1
                 }
-                push_field(self)?;
+                v?.push_into_stack(self)?;
                 protect_lua!(state, 1, 1, fn(state) {
                     ffi::lua_pushcclosure(state, return_field, 1);
                 })?;
diff --git a/src/userdata/registry.rs b/src/userdata/registry.rs
@@ -7,7 +7,7 @@ use std::os::raw::c_void;
 use std::string::String as StdString;
 
 use crate::error::{Error, Result};
-use crate::state::{Lua, RawLua};
+use crate::state::{Lua, LuaGuard};
 use crate::traits::{FromLua, FromLuaMulti, IntoLua, IntoLuaMulti};
 use crate::types::{Callback, MaybeSend};
 use crate::userdata::{AnyUserData, MetaMethod, UserData, UserDataFields, UserDataMethods, UserDataStorage};
@@ -26,8 +26,6 @@ use std::rc::Rc;
 #[cfg(feature = "userdata-wrappers")]
 use std::sync::{Arc, Mutex, RwLock};
 
-type StaticFieldCallback = Box<dyn FnOnce(&RawLua) -> Result<()> + 'static>;
-
 #[derive(Clone, Copy)]
 enum UserDataTypeId {
     Shared(TypeId),
@@ -51,17 +49,18 @@ enum UserDataTypeId {
 
 /// Handle to registry for userdata methods and metamethods.
 pub struct UserDataRegistry<T> {
+    lua: LuaGuard,
     raw: RawUserDataRegistry,
     ud_type_id: UserDataTypeId,
     _type: PhantomData<T>,
 }
 
 pub(crate) struct RawUserDataRegistry {
     // Fields
-    pub(crate) fields: Vec<(String, StaticFieldCallback)>,
+    pub(crate) fields: Vec<(String, Result<Value>)>,
     pub(crate) field_getters: Vec<(String, Callback)>,
     pub(crate) field_setters: Vec<(String, Callback)>,
-    pub(crate) meta_fields: Vec<(String, StaticFieldCallback)>,
+    pub(crate) meta_fields: Vec<(String, Result<Value>)>,
 
     // Methods
     pub(crate) methods: Vec<(String, Callback)>,
@@ -102,17 +101,17 @@ impl UserDataTypeId {
 
 impl<T> UserDataRegistry<T> {
     #[inline(always)]
-    pub(crate) fn new(type_id: TypeId) -> Self {
-        Self::with_type_id(UserDataTypeId::Shared(type_id))
+    pub(crate) fn new(lua: &Lua, type_id: TypeId) -> Self {
+        Self::with_type_id(lua, UserDataTypeId::Shared(type_id))
     }
 
     #[inline(always)]
-    pub(crate) fn new_unique(ud_ptr: *mut c_void) -> Self {
-        Self::with_type_id(UserDataTypeId::Unique(ud_ptr as usize))
+    pub(crate) fn new_unique(lua: &Lua, ud_ptr: *mut c_void) -> Self {
+        Self::with_type_id(lua, UserDataTypeId::Unique(ud_ptr as usize))
     }
 
     #[inline(always)]
-    fn with_type_id(ud_type_id: UserDataTypeId) -> Self {
+    fn with_type_id(lua: &Lua, ud_type_id: UserDataTypeId) -> Self {
         let raw = RawUserDataRegistry {
             fields: Vec::new(),
             field_getters: Vec::new(),
@@ -130,6 +129,7 @@ impl<T> UserDataRegistry<T> {
         };
 
         UserDataRegistry {
+            lua: lua.lock_arc(),
             raw,
             ud_type_id,
             _type: PhantomData,
@@ -548,10 +548,7 @@ impl<T> UserDataFields<T> for UserDataRegistry<T> {
         V: IntoLua + 'static,
     {
         let name = name.to_string();
-        self.raw.fields.push((
-            name,
-            Box::new(move |rawlua| unsafe { value.push_into_stack(rawlua) }),
-        ));
+        self.raw.fields.push((name, value.into_lua(self.lua.lua())));
     }
 
     fn add_field_method_get<M, R>(&mut self, name: impl ToString, method: M)
@@ -598,28 +595,21 @@ impl<T> UserDataFields<T> for UserDataRegistry<T> {
     where
         V: IntoLua + 'static,
     {
+        let lua = self.lua.lua();
         let name = name.to_string();
-        self.raw.meta_fields.push((
-            name.clone(),
-            Box::new(move |rawlua| unsafe {
-                Self::check_meta_field(rawlua.lua(), &name, value)?.push_into_stack(rawlua)
-            }),
-        ));
+        let field = Self::check_meta_field(lua, &name, value).and_then(|v| v.into_lua(lua));
+        self.raw.meta_fields.push((name, field));
     }
 
     fn add_meta_field_with<F, R>(&mut self, name: impl ToString, f: F)
     where
         F: FnOnce(&Lua) -> Result<R> + 'static,
         R: IntoLua,
     {
+        let lua = self.lua.lua();
         let name = name.to_string();
-        self.raw.meta_fields.push((
-            name.clone(),
-            Box::new(move |rawlua| unsafe {
-                let lua = rawlua.lua();
-                Self::check_meta_field(lua, &name, f(lua)?)?.push_into_stack(rawlua)
-            }),
-        ));
+        let field = f(lua).and_then(|v| Self::check_meta_field(lua, &name, v).and_then(|v| v.into_lua(lua)));
+        self.raw.meta_fields.push((name, field));
     }
 }
 
@@ -803,7 +793,7 @@ macro_rules! lua_userdata_impl {
     ($type:ty, $type_id:expr) => {
         impl<T: UserData + 'static> UserData for $type {
             fn register(registry: &mut UserDataRegistry<Self>) {
-                let mut orig_registry = UserDataRegistry::with_type_id($type_id);
+                let mut orig_registry = UserDataRegistry::with_type_id(registry.lua.lua(), $type_id);
                 T::register(&mut orig_registry);
 
                 // Copy all fields, methods, etc. from the original registry
@@ -841,3 +831,9 @@ lua_userdata_impl!(Arc<RwLock<T>> => ArcRwLock);
 lua_userdata_impl!(Arc<parking_lot::Mutex<T>> => ArcParkingLotMutex);
 #[cfg(feature = "userdata-wrappers")]
 lua_userdata_impl!(Arc<parking_lot::RwLock<T>> => ArcParkingLotRwLock);
+
+#[cfg(test)]
+mod assertions {
+    #[cfg(feature = "send")]
+    static_assertions::assert_impl_all!(super::RawUserDataRegistry: Send);
+}

Original file line number	Diff line number	Diff line change
`@@ -1306,7 +1306,7 @@ impl Lua {`
`1306`	`1306`	/// This methods provides a way to add fields or methods to userdata objects of a type `T`.
`1307`	`1307`	`pub fn register_userdata_type<T: 'static>(&self, f: impl FnOnce(&mut UserDataRegistry<T>)) -> Result<()> {`
`1308`	`1308`	`let type_id = TypeId::of::<T>();`
`1309`		`- let mut registry = UserDataRegistry::new(type_id);`
	`1309`	`+ let mut registry = UserDataRegistry::new(self, type_id);`
`1310`	`1310`	`f(&mut registry);`
`1311`	`1311`
`1312`	`1312`	`let lua = self.lock();`
`@@ -1316,8 +1316,8 @@ impl Lua {`
`1316`	`1316`	`ffi::luaL_unref(lua.state(), ffi::LUA_REGISTRYINDEX, table_id);`
`1317`	`1317`	`}`
`1318`	`1318`
`1319`		`- // Register the type`
`1320`		`- lua.create_userdata_metatable(registry.into_raw())?;`
	`1319`	`+ // Add to "pending" registration map`
	`1320`	`+ ((*lua.extra.get()).pending_userdata_reg).insert(type_id, registry.into_raw());`
`1321`	`1321`	`}`
`1322`	`1322`	`Ok(())`
`1323`	`1323`	`}`
Original file line number	Diff line number	Diff line change
`@@ -756,7 +756,7 @@ impl RawLua {`
`756`	`756`	`}`
`757`	`757`
`758`	`758`	// Create a new metatable from `UserData` definition
`759`		`- let mut registry = UserDataRegistry::new(type_id);`
	`759`	`+ let mut registry = UserDataRegistry::new(self.lua(), type_id);`
`760`	`760`	`T::register(&mut registry);`
`761`	`761`
`762`	`762`	`self.create_userdata_metatable(registry.into_raw())`
`@@ -774,9 +774,12 @@ impl RawLua {`
`774`	`774`	`return Ok(table_id as Integer);`
`775`	`775`	`}`
`776`	`776`
`777`		`- // Create an empty metatable`
`778`		`- let registry = UserDataRegistry::<T>::new(type_id);`
`779`		`- self.create_userdata_metatable(registry.into_raw())`
	`777`	`+ // Check if metatable creation is pending or create an empty metatable otherwise`
	`778`	`+ let registry = match (*self.extra.get()).pending_userdata_reg.remove(&type_id) {`
	`779`	`+ Some(registry) => registry,`
	`780`	`+ None => UserDataRegistry::<T>::new(self.lua(), type_id).into_raw(),`
	`781`	`+ };`
	`782`	`+ self.create_userdata_metatable(registry)`
`780`	`783`	`})`
`781`	`784`	`}`
`782`	`785`
`@@ -851,9 +854,9 @@ impl RawLua {`
`851`	`854`	`rawset_field(state, -2, MetaMethod::validate(&k)?)?;`
`852`	`855`	`}`
`853`	`856`	`let mut has_name = false;`
`854`		`- for (k, push_field) in registry.meta_fields {`
	`857`	`+ for (k, v) in registry.meta_fields {`
`855`	`858`	`has_name = has_name \|\| k == MetaMethod::Type;`
`856`		`- push_field(self)?;`
	`859`	`+ v?.push_into_stack(self)?;`
`857`	`860`	`rawset_field(state, -2, MetaMethod::validate(&k)?)?;`
`858`	`861`	`}`
`859`	`862`	// Set `__name/__type` if not provided
`@@ -875,8 +878,8 @@ impl RawLua {`
`875`	`878`	`ffi::lua_pop(state, 1);`
`876`	`879`	`push_table(state, 0, fields_nrec, true)?;`
`877`	`880`	`}`
`878`		`- for (k, push_field) in mem::take(&mut registry.fields) {`
`879`		`- push_field(self)?;`
	`881`	`+ for (k, v) in mem::take(&mut registry.fields) {`
	`882`	`+ v?.push_into_stack(self)?;`
`880`	`883`	`rawset_field(state, -2, &k)?;`
`881`	`884`	`}`
`882`	`885`	`rawset_field(state, metatable_index, "__index")?;`
`@@ -896,12 +899,12 @@ impl RawLua {`
`896`	`899`	`self.push(self.create_callback(m)?)?;`
`897`	`900`	`rawset_field(state, -2, &k)?;`
`898`	`901`	`}`
`899`		`- for (k, push_field) in registry.fields {`
	`902`	`+ for (k, v) in registry.fields {`
`900`	`903`	`unsafe extern "C-unwind" fn return_field(state: *mut ffi::lua_State) -> c_int {`
`901`	`904`	`ffi::lua_pushvalue(state, ffi::lua_upvalueindex(1));`
`902`	`905`	`1`
`903`	`906`	`}`
`904`		`- push_field(self)?;`
	`907`	`+ v?.push_into_stack(self)?;`
`905`	`908`	`protect_lua!(state, 1, 1, fn(state) {`
`906`	`909`	`ffi::lua_pushcclosure(state, return_field, 1);`
`907`	`910`	`})?;`