Optimize RegistryKey internals

- Store single `AtomicI32` field instead of pair i32,AtomicBool - Make creation faster by skipping intermediate `Value` layer Add new `RegistryKey::id()` method to return underlying identifier
mlua-rs · Jun 18, 2024 · c23fa5a · c23fa5a
1 parent 4f1d2ab
commit c23fa5a
Show file tree

Hide file tree

Showing 5 changed files with 89 additions and 76 deletions.
diff --git a/benches/benchmark.rs b/benches/benchmark.rs
@@ -268,6 +268,23 @@ fn registry_value_create(c: &mut Criterion) {
     });
 }
 
+fn registry_value_get(c: &mut Criterion) {
+    let lua = Lua::new();
+    lua.gc_stop();
+
+    let value = lua.create_registry_value("hello").unwrap();
+
+    c.bench_function("registry value [get]", |b| {
+        b.iter_batched(
+            || collect_gc_twice(&lua),
+            |_| {
+                assert_eq!(lua.registry_value::<LuaString>(&value).unwrap(), "hello");
+            },
+            BatchSize::SmallInput,
+        );
+    });
+}
+
 fn userdata_create(c: &mut Criterion) {
     struct UserData(#[allow(unused)] i64);
     impl LuaUserData for UserData {}
@@ -406,6 +423,7 @@ criterion_group! {
         function_async_call_sum,
 
         registry_value_create,
+        registry_value_get,
 
         userdata_create,
         userdata_call_index,

diff --git a/src/conversion.rs b/src/conversion.rs
@@ -453,10 +453,11 @@ impl<'lua> IntoLua<'lua> for &RegistryKey {
             return Err(Error::MismatchedRegistryKey);
         }
 
-        if self.is_nil() {
-            ffi::lua_pushnil(lua.state());
-        } else {
-            ffi::lua_rawgeti(lua.state(), ffi::LUA_REGISTRYINDEX, self.registry_id as _);
+        match self.id() {
+            ffi::LUA_REFNIL => ffi::lua_pushnil(lua.state()),
+            id => {
+                ffi::lua_rawgeti(lua.state(), ffi::LUA_REGISTRYINDEX, id as _);
+            }
         }
         Ok(())
     }

diff --git a/src/lua.rs b/src/lua.rs
@@ -2049,17 +2049,16 @@ impl Lua {
         T: FromLua<'lua>,
     {
         let state = self.state();
-        let value = unsafe {
+        unsafe {
             let _sg = StackGuard::new(state);
             check_stack(state, 3)?;
 
             let protect = !self.unlikely_memory_error();
             push_string(state, name.as_bytes(), protect)?;
             ffi::lua_rawget(state, ffi::LUA_REGISTRYINDEX);
 
-            self.pop_value()
-        };
-        T::from_lua(value, self)
+            T::from_stack(-1, self)
+        }
     }
 
     /// Removes a named value in the Lua registry.
@@ -2082,22 +2081,21 @@ impl Lua {
     ///
     /// [`RegistryKey`]: crate::RegistryKey
     pub fn create_registry_value<'lua, T: IntoLua<'lua>>(&'lua self, t: T) -> Result<RegistryKey> {
-        let t = t.into_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));
-        }
-
         let state = self.state();
         unsafe {
             let _sg = StackGuard::new(state);
             check_stack(state, 4)?;
 
-            self.push_value(t)?;
+            self.push(t)?;
 
-            // Try to reuse previously allocated slot
             let unref_list = (*self.extra.get()).registry_unref_list.clone();
+
+            // Check if the value is nil (no need to store it in the registry)
+            if ffi::lua_isnil(state, -1) != 0 {
+                return Ok(RegistryKey::new(ffi::LUA_REFNIL, unref_list));
+            }
+
+            // Try to reuse previously allocated slot
             let free_registry_id = mlua_expect!(unref_list.lock(), "unref list poisoned")
                 .as_mut()
                 .and_then(|x| x.pop());
@@ -2107,7 +2105,7 @@ impl Lua {
                 return Ok(RegistryKey::new(registry_id, unref_list));
             }
 
-            // Allocate a new RegistryKey
+            // Allocate a new RegistryKey slot
             let registry_id = if self.unlikely_memory_error() {
                 ffi::luaL_ref(state, ffi::LUA_REGISTRYINDEX)
             } else {
@@ -2131,18 +2129,16 @@ impl Lua {
         }
 
         let state = self.state();
-        let value = match key.is_nil() {
-            true => Value::Nil,
-            false => unsafe {
+        match key.id() {
+            ffi::LUA_REFNIL => T::from_lua(Value::Nil, self),
+            registry_id => unsafe {
                 let _sg = StackGuard::new(state);
                 check_stack(state, 1)?;
 
-                let id = key.registry_id as Integer;
-                ffi::lua_rawgeti(state, ffi::LUA_REGISTRYINDEX, id);
-                self.pop_value()
+                ffi::lua_rawgeti(state, ffi::LUA_REGISTRYINDEX, registry_id as Integer);
+                T::from_stack(-1, self)
             },
-        };
-        T::from_lua(value, self)
+        }
     }
 
     /// Removes a value from the Lua registry.
@@ -2180,29 +2176,32 @@ impl Lua {
         }
 
         let t = t.into_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
-            return Err(Error::runtime("cannot replace nil value with non-nil"));
-        }
 
         let state = self.state();
         unsafe {
             let _sg = StackGuard::new(state);
             check_stack(state, 2)?;
 
-            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);
+            match (t, key.id()) {
+                (Value::Nil, ffi::LUA_REFNIL) => {
+                    // Do nothing, no need to replace nil with nil
+                }
+                (Value::Nil, registry_id) => {
+                    // Remove the value
+                    ffi::luaL_unref(state, ffi::LUA_REGISTRYINDEX, registry_id);
+                    key.set_id(ffi::LUA_REFNIL);
+                }
+                (value, ffi::LUA_REFNIL) => {
+                    // Allocate a new `RegistryKey`
+                    let new_key = self.create_registry_value(value)?;
+                    key.set_id(new_key.take());
+                }
+                (value, registry_id) => {
+                    // It must be safe to replace the value without triggering memory error
+                    self.push_value(value)?;
+                    ffi::lua_rawseti(state, ffi::LUA_REGISTRYINDEX, registry_id as Integer);
+                }
             }
-            // It must be safe to replace the value without triggering memory error
-            ffi::lua_rawseti(state, ffi::LUA_REGISTRYINDEX, id);
         }
         Ok(())
     }

diff --git a/src/types.rs b/src/types.rs
@@ -4,7 +4,7 @@ use std::hash::{Hash, Hasher};
 use std::ops::{Deref, DerefMut};
 use std::os::raw::{c_int, c_void};
 use std::result::Result as StdResult;
-use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::atomic::{AtomicI32, Ordering};
 use std::sync::{Arc, Mutex};
 use std::{fmt, mem, ptr};
 
@@ -204,76 +204,72 @@ pub(crate) struct DestructedUserdata;
 /// [`AnyUserData::set_user_value`]: crate::AnyUserData::set_user_value
 /// [`AnyUserData::user_value`]: crate::AnyUserData::user_value
 pub struct RegistryKey {
-    pub(crate) registry_id: c_int,
-    pub(crate) is_nil: AtomicBool,
+    pub(crate) registry_id: AtomicI32,
     pub(crate) unref_list: Arc<Mutex<Option<Vec<c_int>>>>,
 }
 
 impl fmt::Debug for RegistryKey {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "RegistryKey({})", self.registry_id)
+        write!(f, "RegistryKey({})", self.id())
     }
 }
 
 impl Hash for RegistryKey {
     fn hash<H: Hasher>(&self, state: &mut H) {
-        self.registry_id.hash(state)
+        self.id().hash(state)
     }
 }
 
 impl PartialEq for RegistryKey {
     fn eq(&self, other: &RegistryKey) -> bool {
-        self.registry_id == other.registry_id && Arc::ptr_eq(&self.unref_list, &other.unref_list)
+        self.id() == other.id() && Arc::ptr_eq(&self.unref_list, &other.unref_list)
     }
 }
 
 impl Eq for RegistryKey {}
 
 impl Drop for RegistryKey {
     fn drop(&mut self) {
+        let registry_id = self.id();
         // We don't need to collect nil slot
-        if self.registry_id > ffi::LUA_REFNIL {
+        if 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);
+                list.push(registry_id);
             }
         }
     }
 }
 
 impl RegistryKey {
-    // Creates a new instance of `RegistryKey`
+    /// 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),
+            registry_id: AtomicI32::new(id),
             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 {
-            ptr::read(&self.unref_list);
-            mem::forget(self);
-        }
-        registry_id
+    /// Returns the underlying Lua reference of this `RegistryKey`
+    #[inline(always)]
+    pub fn id(&self) -> c_int {
+        self.registry_id.load(Ordering::Relaxed)
     }
 
-    // Returns true if this `RegistryKey` holds a nil value
+    /// Sets the unique Lua reference key of this `RegistryKey`
     #[inline(always)]
-    pub(crate) fn is_nil(&self) -> bool {
-        self.is_nil.load(Ordering::Relaxed)
+    pub(crate) fn set_id(&self, id: c_int) {
+        self.registry_id.store(id, 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);
+    /// Destroys the `RegistryKey` without adding to the unref list
+    pub(crate) fn take(self) -> i32 {
+        let registry_id = self.id();
+        unsafe {
+            ptr::read(&self.unref_list);
+            mem::forget(self);
+        }
+        registry_id
     }
 }
 

diff --git a/tests/tests.rs b/tests/tests.rs
@@ -775,12 +775,11 @@ fn test_replace_registry_value() -> Result<()> {
     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:?}"),
-    }
+    lua.replace_registry_value(&key2, Value::Nil)?;
+    assert_eq!(lua.registry_value::<Value>(&key2)?, Value::Nil);
+    lua.replace_registry_value(&key2, "abc")?;
+    assert_eq!(lua.registry_value::<String>(&key2)?, "abc");
 
     Ok(())
 }