Fix bug when recycled Registry slot can be set to Nil.

This can result in allocating the same slot twice and rewriting old value.
Lua uses (registry) table length to find next free slot and having Nil in the middle of the table can impact length calculation.
With this fix we ensure that Nil values uses a special LUA_REFNIL slot.

src/lua.rs

@@ -88,6 +88,8 @@ pub(crate) struct ExtraData {
 
     registered_userdata: FxHashMap<TypeId, c_int>,
     registered_userdata_mt: FxHashMap<*const c_void, Option<TypeId>>,
+
+    // When Lua instance dropped, setting `None` would prevent collecting `RegistryKey`s
     registry_unref_list: Arc<Mutex<Option<Vec<c_int>>>>,
 
     #[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<RegistryKey> {
         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(())

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<Mutex<Option<Vec<c_int>>>>,
 }
 
@@ -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<Mutex<Option<Vec<c_int>>>>) -> 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> {

tests/tests.rs

@@ -793,6 +793,17 @@ fn test_replace_registry_value() -> Result<()> {
     let key = lua.create_registry_value::<i32>(42)?;
     lua.replace_registry_value(&key, "new value")?;
     assert_eq!(lua.registry_value::<String>(&key)?, "new value");
+    lua.replace_registry_value(&key, Value::Nil)?;
+    assert_eq!(lua.registry_value::<Value>(&key)?, Value::Nil);
+    lua.replace_registry_value(&key, 123)?;
+    assert_eq!(lua.registry_value::<i32>(&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();