diff --git a/src/lua.rs b/src/lua.rs index 38d7c42..d1001de 100644 --- a/src/lua.rs +++ b/src/lua.rs @@ -88,6 +88,8 @@ pub(crate) struct ExtraData { registered_userdata: FxHashMap, registered_userdata_mt: FxHashMap<*const c_void, Option>, + + // When Lua instance dropped, setting `None` would prevent collecting `RegistryKey`s registry_unref_list: Arc>>>, #[cfg(not(feature = "send"))] @@ -2060,6 +2062,12 @@ impl Lua { /// [`RegistryKey`]: crate::RegistryKey pub fn create_registry_value<'lua, T: ToLua<'lua>>(&'lua self, t: T) -> Result { let t = t.to_lua(self)?; + if t == Value::Nil { + // Special case to skip calling `luaL_ref` and use `LUA_REFNIL` instead + let unref_list = unsafe { (*self.extra.get()).registry_unref_list.clone() }; + return Ok(RegistryKey::new(ffi::LUA_REFNIL, unref_list)); + } + unsafe { let _sg = StackGuard::new(self.state); check_stack(self.state, 4)?; @@ -2067,27 +2075,20 @@ impl Lua { let unref_list = (*self.extra.get()).registry_unref_list.clone(); self.push_value(t)?; - // Try to reuse previously allocated RegistryKey + // Try to reuse previously allocated slot let unref_list2 = unref_list.clone(); let mut unref_list2 = mlua_expect!(unref_list2.lock(), "unref list poisoned"); if let Some(registry_id) = unref_list2.as_mut().and_then(|x| x.pop()) { // It must be safe to replace the value without triggering memory error ffi::lua_rawseti(self.state, ffi::LUA_REGISTRYINDEX, registry_id as Integer); - return Ok(RegistryKey { - registry_id, - unref_list, - }); + return Ok(RegistryKey::new(registry_id, unref_list)); } // Allocate a new RegistryKey let registry_id = protect_lua!(self.state, 1, 0, |state| { ffi::luaL_ref(state, ffi::LUA_REGISTRYINDEX) })?; - - Ok(RegistryKey { - registry_id, - unref_list, - }) + Ok(RegistryKey::new(registry_id, unref_list)) } } @@ -2102,13 +2103,16 @@ impl Lua { return Err(Error::MismatchedRegistryKey); } - let value = unsafe { - let _sg = StackGuard::new(self.state); - check_stack(self.state, 1)?; + let value = match key.is_nil() { + true => Value::Nil, + false => unsafe { + let _sg = StackGuard::new(self.state); + check_stack(self.state, 1)?; - let id = key.registry_id as Integer; - ffi::lua_rawgeti(self.state, ffi::LUA_REGISTRYINDEX, id); - self.pop_value() + let id = key.registry_id as Integer; + ffi::lua_rawgeti(self.state, ffi::LUA_REGISTRYINDEX, id); + self.pop_value() + }, }; T::from_lua(value, self) } @@ -2147,13 +2151,28 @@ impl Lua { } let t = t.to_lua(self)?; + if t == Value::Nil && key.is_nil() { + // Nothing to replace + return Ok(()); + } else if t != Value::Nil && key.registry_id == ffi::LUA_REFNIL { + // We cannot update `LUA_REFNIL` slot + let err = "cannot replace nil value with non-nil".to_string(); + return Err(Error::RuntimeError(err)); + } + unsafe { let _sg = StackGuard::new(self.state); check_stack(self.state, 2)?; - self.push_value(t)?; - // It must be safe to replace the value without triggering memory error let id = key.registry_id as Integer; + if t == Value::Nil { + self.push_value(Value::Integer(id))?; + key.set_nil(true); + } else { + self.push_value(t)?; + key.set_nil(false); + } + // It must be safe to replace the value without triggering memory error ffi::lua_rawseti(self.state, ffi::LUA_REGISTRYINDEX, id); } Ok(()) diff --git a/src/types.rs b/src/types.rs index 7b1f653..9404d10 100644 --- a/src/types.rs +++ b/src/types.rs @@ -1,6 +1,7 @@ use std::cell::UnsafeCell; use std::hash::{Hash, Hasher}; use std::os::raw::{c_int, c_void}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use std::{fmt, mem, ptr}; @@ -104,6 +105,7 @@ pub(crate) struct DestructedUserdata; /// [`AnyUserData::get_user_value`]: crate::AnyUserData::get_user_value pub struct RegistryKey { pub(crate) registry_id: c_int, + pub(crate) is_nil: AtomicBool, pub(crate) unref_list: Arc>>>, } @@ -129,15 +131,27 @@ impl Eq for RegistryKey {} impl Drop for RegistryKey { fn drop(&mut self) { - let mut unref_list = mlua_expect!(self.unref_list.lock(), "unref list poisoned"); - if let Some(list) = unref_list.as_mut() { - list.push(self.registry_id); + // We don't need to collect nil slot + if self.registry_id > ffi::LUA_REFNIL { + let mut unref_list = mlua_expect!(self.unref_list.lock(), "unref list poisoned"); + if let Some(list) = unref_list.as_mut() { + list.push(self.registry_id); + } } } } impl RegistryKey { - // Destroys the RegistryKey without adding to the drop list + // Creates a new instance of `RegistryKey` + pub(crate) const fn new(id: c_int, unref_list: Arc>>>) -> Self { + RegistryKey { + registry_id: id, + is_nil: AtomicBool::new(id == ffi::LUA_REFNIL), + unref_list, + } + } + + // Destroys the `RegistryKey` without adding to the unref list pub(crate) fn take(self) -> c_int { let registry_id = self.registry_id; unsafe { @@ -146,6 +160,21 @@ impl RegistryKey { } registry_id } + + // Returns true if this `RegistryKey` holds a nil value + #[inline(always)] + pub(crate) fn is_nil(&self) -> bool { + self.is_nil.load(Ordering::Relaxed) + } + + // Marks value of this `RegistryKey` as `Nil` + #[inline(always)] + pub(crate) fn set_nil(&self, enabled: bool) { + // We cannot replace previous value with nil in as this will break + // Lua mechanism to find free keys. + // Instead, we set a special flag to mark value as nil. + self.is_nil.store(enabled, Ordering::Relaxed); + } } pub(crate) struct LuaRef<'lua> { diff --git a/tests/tests.rs b/tests/tests.rs index 59fb28f..a71691e 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -793,6 +793,17 @@ fn test_replace_registry_value() -> Result<()> { let key = lua.create_registry_value::(42)?; lua.replace_registry_value(&key, "new value")?; assert_eq!(lua.registry_value::(&key)?, "new value"); + lua.replace_registry_value(&key, Value::Nil)?; + assert_eq!(lua.registry_value::(&key)?, Value::Nil); + lua.replace_registry_value(&key, 123)?; + assert_eq!(lua.registry_value::(&key)?, 123); + + // It should be impossible to replace (initial) nil value with non-nil + let key2 = lua.create_registry_value(Value::Nil)?; + match lua.replace_registry_value(&key2, "abc") { + Err(Error::RuntimeError(_)) => {} + r => panic!("expected RuntimeError, got {r:?}"), + } Ok(()) } @@ -844,6 +855,28 @@ fn test_mismatched_registry_key() -> Result<()> { Ok(()) } +#[test] +fn test_registry_value_reuse() -> Result<()> { + let lua = Lua::new(); + + let r1 = lua.create_registry_value("value1")?; + let r1_slot = format!("{r1:?}"); + drop(r1); + + // Previous slot must not be reused by nil value + let r2 = lua.create_registry_value(Value::Nil)?; + let r2_slot = format!("{r2:?}"); + assert_ne!(r1_slot, r2_slot); + drop(r2); + + // But should be reused by non-nil value + let r3 = lua.create_registry_value("value3")?; + let r3_slot = format!("{r3:?}"); + assert_eq!(r1_slot, r3_slot); + + Ok(()) +} + #[test] fn test_application_data() -> Result<()> { let lua = Lua::new();