From 335f433df4142deb8fc4bb325f3dffaecb695c99 Mon Sep 17 00:00:00 2001
From: Alex Orlenko <zxteam@protonmail.com>
Date: Sun, 21 Feb 2021 23:51:49 +0000
Subject: [PATCH] Optimize callbacks

---
 src/lua.rs  | 47 +++++++++++++++++++++++++----------------------
 src/util.rs |  6 +-----
 2 files changed, 26 insertions(+), 27 deletions(-)

diff --git a/src/lua.rs b/src/lua.rs
index 91e0edd..894b1d0 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -22,10 +22,10 @@ use crate::types::{
 };
 use crate::userdata::{AnyUserData, MetaMethod, UserData, UserDataMethods, UserDataWrapped};
 use crate::util::{
-    assert_stack, callback_error, check_stack, get_gc_userdata, get_main_state,
-    get_meta_gc_userdata, get_wrapped_error, init_error_registry, init_gc_metatable_for,
-    init_userdata_metatable, pop_error, protect_lua, protect_lua_closure, push_gc_userdata,
-    push_meta_gc_userdata, push_string, push_userdata, push_wrapped_error, StackGuard,
+    assert_stack, callback_error, check_stack, get_gc_userdata, get_main_state, get_userdata,
+    get_wrapped_error, init_error_registry, init_gc_metatable_for, init_userdata_metatable,
+    pop_error, protect_lua, protect_lua_closure, push_gc_userdata, push_meta_gc_userdata,
+    push_string, push_userdata, push_wrapped_error, StackGuard,
 };
 use crate::value::{FromLua, FromLuaMulti, MultiValue, Nil, ToLua, ToLuaMulti, Value};
 
@@ -1504,9 +1504,10 @@ impl Lua {
     }
 
     pub(crate) unsafe fn userdata_metatable<T: 'static + UserData>(&self) -> Result<c_int> {
+        let type_id = TypeId::of::<T>();
         if let Some(table_id) = mlua_expect!(self.extra.lock(), "extra is poisoned")
             .registered_userdata
-            .get(&TypeId::of::<T>())
+            .get(&type_id)
         {
             return Ok(*table_id);
         }
@@ -1567,7 +1568,7 @@ impl Lua {
         })?;
 
         let mut extra = mlua_expect!(self.extra.lock(), "extra is poisoned");
-        extra.registered_userdata.insert(TypeId::of::<T>(), id);
+        extra.registered_userdata.insert(type_id, id);
         extra.registered_userdata_mt.insert(ptr);
 
         Ok(id)
@@ -1616,12 +1617,13 @@ impl Lua {
     {
         unsafe extern "C" fn call_callback(state: *mut ffi::lua_State) -> c_int {
             callback_error(state, |nargs| {
-                let func =
-                    get_meta_gc_userdata::<Callback, Callback>(state, ffi::lua_upvalueindex(1));
-                let lua = get_gc_userdata::<Lua>(state, ffi::lua_upvalueindex(2));
-                if func.is_null() || lua.is_null() {
+                if ffi::lua_type(state, ffi::lua_upvalueindex(1)) == ffi::LUA_TNIL
+                    || ffi::lua_type(state, ffi::lua_upvalueindex(2)) == ffi::LUA_TNIL
+                {
                     return Err(Error::CallbackDestructed);
                 }
+                let func = get_userdata::<Callback>(state, ffi::lua_upvalueindex(1));
+                let lua = get_userdata::<Lua>(state, ffi::lua_upvalueindex(2));
 
                 if nargs < ffi::LUA_MINSTACK {
                     check_stack(state, ffi::LUA_MINSTACK - nargs)?;
@@ -1681,14 +1683,13 @@ impl Lua {
 
         unsafe extern "C" fn call_callback(state: *mut ffi::lua_State) -> c_int {
             callback_error(state, |nargs| {
-                let func = get_meta_gc_userdata::<AsyncCallback, AsyncCallback>(
-                    state,
-                    ffi::lua_upvalueindex(1),
-                );
-                let lua = get_gc_userdata::<Lua>(state, ffi::lua_upvalueindex(2));
-                if func.is_null() || lua.is_null() {
+                if ffi::lua_type(state, ffi::lua_upvalueindex(1)) == ffi::LUA_TNIL
+                    || ffi::lua_type(state, ffi::lua_upvalueindex(2)) == ffi::LUA_TNIL
+                {
                     return Err(Error::CallbackDestructed);
                 }
+                let func = get_userdata::<AsyncCallback>(state, ffi::lua_upvalueindex(1));
+                let lua = get_userdata::<Lua>(state, ffi::lua_upvalueindex(2));
 
                 if nargs < ffi::LUA_MINSTACK {
                     check_stack(state, ffi::LUA_MINSTACK - nargs)?;
@@ -1715,14 +1716,16 @@ impl Lua {
 
         unsafe extern "C" fn poll_future(state: *mut ffi::lua_State) -> c_int {
             callback_error(state, |nargs| {
-                let fut = get_gc_userdata::<LocalBoxFuture<Result<MultiValue>>>(
+                if ffi::lua_type(state, ffi::lua_upvalueindex(1)) == ffi::LUA_TNIL
+                    || ffi::lua_type(state, ffi::lua_upvalueindex(2)) == ffi::LUA_TNIL
+                {
+                    return Err(Error::CallbackDestructed);
+                }
+                let fut = get_userdata::<LocalBoxFuture<Result<MultiValue>>>(
                     state,
                     ffi::lua_upvalueindex(1),
                 );
-                let lua = get_gc_userdata::<Lua>(state, ffi::lua_upvalueindex(2));
-                if fut.is_null() || lua.is_null() {
-                    return Err(Error::CallbackDestructed);
-                }
+                let lua = get_userdata::<Lua>(state, ffi::lua_upvalueindex(2));
 
                 if nargs < ffi::LUA_MINSTACK {
                     check_stack(state, ffi::LUA_MINSTACK - nargs)?;
@@ -1791,7 +1794,7 @@ impl Lua {
         self.load(
             r#"
             poll = get_poll(...)
-            local poll = poll
+            local poll, yield, unpack = poll, yield, unpack
             while true do
                 ready, res = poll()
                 if ready then
diff --git a/src/util.rs b/src/util.rs
index c1052d6..ae668fc 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -279,15 +279,11 @@ pub unsafe fn push_meta_gc_userdata<MT: Any, T>(state: *mut ffi::lua_State, t: T
 
 // Uses 2 stack spaces, does not call checkstack
 pub unsafe fn get_gc_userdata<T: Any>(state: *mut ffi::lua_State, index: c_int) -> *mut T {
-    get_meta_gc_userdata::<T, T>(state, index)
-}
-
-pub unsafe fn get_meta_gc_userdata<MT: Any, T>(state: *mut ffi::lua_State, index: c_int) -> *mut T {
     let ud = ffi::lua_touserdata(state, index) as *mut T;
     if ud.is_null() || ffi::lua_getmetatable(state, index) == 0 {
         return ptr::null_mut();
     }
-    get_gc_metatable_for::<MT>(state);
+    get_gc_metatable_for::<T>(state);
     let res = ffi::lua_rawequal(state, -1, -2) != 0;
     ffi::lua_pop(state, 2);
     if !res {