Optimize (and simplify) protected mode for Rust function calls

mlua-rs · Dec 7, 2024 · 91e069a · 91e069a
1 parent cacd3dc
commit 91e069a
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 37 deletions.
diff --git a/src/state/extra.rs b/src/state/extra.rs
@@ -27,7 +27,7 @@ use super::{Lua, WeakLua};
 // Unique key to store `ExtraData` in the registry
 static EXTRA_REGISTRY_KEY: u8 = 0;
 
-const WRAPPED_FAILURE_POOL_SIZE: usize = 64;
+const WRAPPED_FAILURE_POOL_DEFAULT_CAPACITY: usize = 64;
 const REF_STACK_RESERVE: c_int = 1;
 
 /// Data associated with the Lua state.
@@ -60,6 +60,7 @@ pub(crate) struct ExtraData {
 
     // Pool of `WrappedFailure` enums in the ref thread (as userdata)
     pub(super) wrapped_failure_pool: Vec<c_int>,
+    pub(super) wrapped_failure_top: usize,
     // Pool of `Thread`s (coroutines) for async execution
     #[cfg(feature = "async")]
     pub(super) thread_pool: Vec<c_int>,
@@ -160,7 +161,8 @@ impl ExtraData {
             ref_stack_size: ffi::LUA_MINSTACK - REF_STACK_RESERVE,
             ref_stack_top: ffi::lua_gettop(ref_thread),
             ref_free: Vec::new(),
-            wrapped_failure_pool: Vec::with_capacity(WRAPPED_FAILURE_POOL_SIZE),
+            wrapped_failure_pool: Vec::with_capacity(WRAPPED_FAILURE_POOL_DEFAULT_CAPACITY),
+            wrapped_failure_top: 0,
             #[cfg(feature = "async")]
             thread_pool: Vec::new(),
             wrapped_failure_mt_ptr,

diff --git a/src/state/util.rs b/src/state/util.rs
@@ -7,8 +7,6 @@ use crate::error::{Error, Result};
 use crate::state::{ExtraData, RawLua};
 use crate::util::{self, get_internal_metatable, WrappedFailure};
 
-const WRAPPED_FAILURE_POOL_SIZE: usize = 64;
-
 pub(super) struct StateGuard<'a>(&'a RawLua, *mut ffi::lua_State);
 
 impl<'a> StateGuard<'a> {
@@ -42,39 +40,40 @@ where
 
     enum PreallocatedFailure {
         New(*mut WrappedFailure),
-        Existing(i32),
+        Reserved,
     }
 
     impl PreallocatedFailure {
         unsafe fn reserve(state: *mut ffi::lua_State, extra: *mut ExtraData) -> Self {
-            match (*extra).wrapped_failure_pool.pop() {
-                Some(index) => PreallocatedFailure::Existing(index),
-                None => {
-                    // We need to check stack for Luau in case when callback is called from interrupt
-                    // See https://github.com/Roblox/luau/issues/446 and mlua #142 and #153
-                    #[cfg(feature = "luau")]
-                    ffi::lua_rawcheckstack(state, 2);
-                    // Place it to the beginning of the stack
-                    let ud = WrappedFailure::new_userdata(state);
-                    ffi::lua_insert(state, 1);
-                    PreallocatedFailure::New(ud)
-                }
+            if (*extra).wrapped_failure_top > 0 {
+                (*extra).wrapped_failure_top -= 1;
+                return PreallocatedFailure::Reserved;
             }
+
+            // We need to check stack for Luau in case when callback is called from interrupt
+            // See https://github.com/Roblox/luau/issues/446 and mlua #142 and #153
+            #[cfg(feature = "luau")]
+            ffi::lua_rawcheckstack(state, 2);
+            // Place it to the beginning of the stack
+            let ud = WrappedFailure::new_userdata(state);
+            ffi::lua_insert(state, 1);
+            PreallocatedFailure::New(ud)
         }
 
+        #[cold]
         unsafe fn r#use(&self, state: *mut ffi::lua_State, extra: *mut ExtraData) -> *mut WrappedFailure {
             let ref_thread = (*extra).ref_thread;
             match *self {
                 PreallocatedFailure::New(ud) => {
                     ffi::lua_settop(state, 1);
                     ud
                 }
-                PreallocatedFailure::Existing(index) => {
+                PreallocatedFailure::Reserved => {
+                    let index = (*extra).wrapped_failure_pool.pop().unwrap();
                     ffi::lua_settop(state, 0);
                     #[cfg(feature = "luau")]
                     ffi::lua_rawcheckstack(state, 2);
-                    ffi::lua_pushvalue(ref_thread, index);
-                    ffi::lua_xmove(ref_thread, state, 1);
+                    ffi::lua_xpush(ref_thread, state, index);
                     ffi::lua_pushnil(ref_thread);
                     ffi::lua_replace(ref_thread, index);
                     (*extra).ref_free.push(index);
@@ -87,24 +86,13 @@ where
             let ref_thread = (*extra).ref_thread;
             match self {
                 PreallocatedFailure::New(_) => {
-                    if (*extra).wrapped_failure_pool.len() < WRAPPED_FAILURE_POOL_SIZE {
-                        ffi::lua_rotate(state, 1, -1);
-                        ffi::lua_xmove(state, ref_thread, 1);
-                        let index = ref_stack_pop(extra);
-                        (*extra).wrapped_failure_pool.push(index);
-                    } else {
-                        ffi::lua_remove(state, 1);
-                    }
-                }
-                PreallocatedFailure::Existing(index) => {
-                    if (*extra).wrapped_failure_pool.len() < WRAPPED_FAILURE_POOL_SIZE {
-                        (*extra).wrapped_failure_pool.push(index);
-                    } else {
-                        ffi::lua_pushnil(ref_thread);
-                        ffi::lua_replace(ref_thread, index);
-                        (*extra).ref_free.push(index);
-                    }
+                    ffi::lua_rotate(state, 1, -1);
+                    ffi::lua_xmove(state, ref_thread, 1);
+                    let index = ref_stack_pop(extra);
+                    (*extra).wrapped_failure_pool.push(index);
+                    (*extra).wrapped_failure_top += 1;
                 }
+                PreallocatedFailure::Reserved => (*extra).wrapped_failure_top += 1,
             }
         }
     }