Make Lua reference values cheap to clone

Instead of locking the VM and making a copy on auxiliary thread, track number of references using Rust ref counter.
This should also help reducing number of used references (they are limited to to 1M usually) on auxiliary thread.

Author: khvzak

benches/benchmark.rs

@@ -128,6 +128,22 @@ fn table_traversal_sequence(c: &mut Criterion) {
     });
 }
 
+fn table_ref_clone(c: &mut Criterion) {
+    let lua = Lua::new();
+
+    let t = lua.create_table().unwrap();
+
+    c.bench_function("table [ref clone]", |b| {
+        b.iter_batched(
+            || collect_gc_twice(&lua),
+            |_| {
+                let _t2 = t.clone();
+            },
+            BatchSize::SmallInput,
+        );
+    });
+}
+
 fn function_create(c: &mut Criterion) {
     let lua = Lua::new();
 
@@ -399,6 +415,7 @@ criterion_group! {
         table_traversal_pairs,
         table_traversal_for_each,
         table_traversal_sequence,
+        table_ref_clone,
 
         function_create,
         function_call_sum,

src/state/extra.rs

@@ -12,7 +12,7 @@ use rustc_hash::FxHashMap;
 use crate::error::Result;
 use crate::state::RawLua;
 use crate::stdlib::StdLib;
-use crate::types::{AppData, ReentrantMutex, XRc};
+use crate::types::{AppData, ReentrantMutex, ValueRefIndex, XRc};
 use crate::userdata::RawUserDataRegistry;
 use crate::util::{get_internal_metatable, push_internal_userdata, TypeKey, WrappedFailure};
 
@@ -64,7 +64,7 @@ pub(crate) struct ExtraData {
     pub(super) wrapped_failure_top: usize,
     // Pool of `Thread`s (coroutines) for async execution
     #[cfg(feature = "async")]
-    pub(super) thread_pool: Vec<c_int>,
+    pub(super) thread_pool: Vec<ValueRefIndex>,
 
     // Address of `WrappedFailure` metatable
     pub(super) wrapped_failure_mt_ptr: *const c_void,

src/state/raw.rs

@@ -624,7 +624,7 @@ impl RawLua {
     #[cfg(feature = "async")]
     pub(crate) unsafe fn create_recycled_thread(&self, func: &Function) -> Result<Thread> {
         if let Some(index) = (*self.extra.get()).thread_pool.pop() {
-            let thread_state = ffi::lua_tothread(self.ref_thread(), index);
+            let thread_state = ffi::lua_tothread(self.ref_thread(), *index.0);
             ffi::lua_xpush(self.ref_thread(), thread_state, func.0.index);
 
             #[cfg(feature = "luau")]
@@ -645,8 +645,9 @@ impl RawLua {
     pub(crate) unsafe fn recycle_thread(&self, thread: &mut Thread) {
         let extra = &mut *self.extra.get();
         if extra.thread_pool.len() < extra.thread_pool.capacity() {
-            extra.thread_pool.push(thread.0.index);
-            thread.0.drop = false; // Prevent thread from being garbage collected
+            if let Some(index) = thread.0.index_count.take() {
+                extra.thread_pool.push(index);
+            }
         }
     }
 
@@ -827,13 +828,6 @@ impl RawLua {
         ValueRef::new(self, index)
     }
 
-    #[inline]
-    pub(crate) unsafe fn clone_ref(&self, vref: &ValueRef) -> ValueRef {
-        ffi::lua_pushvalue(self.ref_thread(), vref.index);
-        let index = (*self.extra.get()).ref_stack_pop();
-        ValueRef::new(self, index)
-    }
-
     pub(crate) unsafe fn drop_ref(&self, vref: &ValueRef) {
         let ref_thread = self.ref_thread();
         mlua_debug_assert!(

src/table.rs

@@ -884,7 +884,7 @@ impl ObjectLike for Table {
         R: FromLuaMulti,
     {
         // Convert table to a function and call via pcall that respects the `__call` metamethod.
-        Function(self.0.copy()).call(args)
+        Function(self.0.clone()).call(args)
     }
 
     #[cfg(feature = "async")]
@@ -893,7 +893,7 @@ impl ObjectLike for Table {
     where
         R: FromLuaMulti,
     {
-        Function(self.0.copy()).call_async(args)
+        Function(self.0.clone()).call_async(args)
     }
 
     #[inline]
@@ -941,7 +941,7 @@ impl ObjectLike for Table {
 
     #[inline]
     fn to_string(&self) -> Result<StdString> {
-        Value::Table(Table(self.0.copy())).to_string()
+        Value::Table(Table(self.0.clone())).to_string()
     }
 }
 

src/types.rs

@@ -18,7 +18,7 @@ pub(crate) type BoxFuture<'a, T> = futures_util::future::LocalBoxFuture<'a, T>;
 pub use app_data::{AppData, AppDataRef, AppDataRefMut};
 pub use either::Either;
 pub use registry_key::RegistryKey;
-pub(crate) use value_ref::ValueRef;
+pub(crate) use value_ref::{ValueRef, ValueRefIndex};
 
 /// Type of Lua integer numbers.
 pub type Integer = ffi::lua_Integer;

src/types/value_ref.rs

@@ -1,22 +1,39 @@
 use std::fmt;
 use std::os::raw::{c_int, c_void};
 
+use super::XRc;
 use crate::state::{RawLua, WeakLua};
 
 /// A reference to a Lua (complex) value stored in the Lua auxiliary thread.
+#[derive(Clone)]
 pub struct ValueRef {
     pub(crate) lua: WeakLua,
+    // Keep index separate to avoid additional indirection when accessing it.
     pub(crate) index: c_int,
-    pub(crate) drop: bool,
+    // If `index_count` is `None`, the value does not need to be destroyed.
+    pub(crate) index_count: Option<ValueRefIndex>,
+}
+
+/// A reference to a Lua value index in the auxiliary thread.
+/// It's cheap to clone and can be used to track the number of references to a value.
+#[derive(Clone)]
+pub(crate) struct ValueRefIndex(pub(crate) XRc<c_int>);
+
+impl From<c_int> for ValueRefIndex {
+    #[inline]
+    fn from(index: c_int) -> Self {
+        ValueRefIndex(XRc::new(index))
+    }
 }
 
 impl ValueRef {
     #[inline]
-    pub(crate) fn new(lua: &RawLua, index: c_int) -> Self {
+    pub(crate) fn new(lua: &RawLua, index: impl Into<ValueRefIndex>) -> Self {
+        let index = index.into();
         ValueRef {
             lua: lua.weak().clone(),
-            index,
-            drop: true,
+            index: *index.0,
+            index_count: Some(index),
         }
     }
 
@@ -25,16 +42,6 @@ impl ValueRef {
         let lua = self.lua.lock();
         unsafe { ffi::lua_topointer(lua.ref_thread(), self.index) }
     }
-
-    /// Returns a copy of the value, which is valid as long as the original value is held.
-    #[inline]
-    pub(crate) fn copy(&self) -> Self {
-        ValueRef {
-            lua: self.lua.clone(),
-            index: self.index,
-            drop: false,
-        }
-    }
 }
 
 impl fmt::Debug for ValueRef {
@@ -43,17 +50,15 @@ impl fmt::Debug for ValueRef {
     }
 }
 
-impl Clone for ValueRef {
-    fn clone(&self) -> Self {
-        unsafe { self.lua.lock().clone_ref(self) }
-    }
-}
-
 impl Drop for ValueRef {
     fn drop(&mut self) {
-        if self.drop {
-            if let Some(lua) = self.lua.try_lock() {
-                unsafe { lua.drop_ref(self) };
+        if let Some(ValueRefIndex(index)) = self.index_count.take() {
+            // It's guaranteed that the inner value returns exactly once.
+            // This means in particular that the value is not dropped.
+            if XRc::into_inner(index).is_some() {
+                if let Some(lua) = self.lua.try_lock() {
+                    unsafe { lua.drop_ref(self) };
+                }
             }
         }
     }

src/userdata/object.rs

@@ -15,22 +15,22 @@ impl ObjectLike for AnyUserData {
     fn get<V: FromLua>(&self, key: impl IntoLua) -> Result<V> {
         // `lua_gettable` method used under the hood can work with any Lua value
         // that has `__index` metamethod
-        Table(self.0.copy()).get_protected(key)
+        Table(self.0.clone()).get_protected(key)
     }
 
     #[inline]
     fn set(&self, key: impl IntoLua, value: impl IntoLua) -> Result<()> {
         // `lua_settable` method used under the hood can work with any Lua value
         // that has `__newindex` metamethod
-        Table(self.0.copy()).set_protected(key, value)
+        Table(self.0.clone()).set_protected(key, value)
     }
 
     #[inline]
     fn call<R>(&self, args: impl IntoLuaMulti) -> Result<R>
     where
         R: FromLuaMulti,
     {
-        Function(self.0.copy()).call(args)
+        Function(self.0.clone()).call(args)
     }
 
     #[cfg(feature = "async")]
@@ -39,7 +39,7 @@ impl ObjectLike for AnyUserData {
     where
         R: FromLuaMulti,
     {
-        Function(self.0.copy()).call_async(args)
+        Function(self.0.clone()).call_async(args)
     }
 
     #[inline]
@@ -88,6 +88,6 @@ impl ObjectLike for AnyUserData {
 
     #[inline]
     fn to_string(&self) -> Result<StdString> {
-        Value::UserData(AnyUserData(self.0.copy())).to_string()
+        Value::UserData(AnyUserData(self.0.clone())).to_string()
     }
 }