From e8de2a458ad8dea6847f17b3c7d3fb564109a3f3 Mon Sep 17 00:00:00 2001
From: Alex Orlenko <zxteam@protonmail.com>
Date: Tue, 18 May 2021 20:07:34 +0100
Subject: [PATCH] Allow multiple entrypoints in a single module share the same
 Lua state. Previously it would initialize different Lua instances. Fixes #49.

---
 examples/module/src/lib.rs | 17 +++++++++++++++++
 src/hook.rs                |  2 +-
 src/lua.rs                 | 18 ++++++++++++------
 tests/module/tests/load.rs | 13 +++++++++++++
 4 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/examples/module/src/lib.rs b/examples/module/src/lib.rs
index 38cc0fd..dcd487a 100644
--- a/examples/module/src/lib.rs
+++ b/examples/module/src/lib.rs
@@ -8,10 +8,27 @@ fn used_memory(lua: &Lua, _: ()) -> LuaResult<usize> {
     Ok(lua.used_memory())
 }
 
+fn check_userdata(_: &Lua, ud: MyUserData) -> LuaResult<i32> {
+    Ok(ud.0)
+}
+
 #[mlua::lua_module]
 fn rust_module(lua: &Lua) -> LuaResult<LuaTable> {
     let exports = lua.create_table()?;
     exports.set("sum", lua.create_function(sum)?)?;
     exports.set("used_memory", lua.create_function(used_memory)?)?;
+    exports.set("check_userdata", lua.create_function(check_userdata)?)?;
+    Ok(exports)
+}
+
+#[derive(Clone, Copy)]
+struct MyUserData(i32);
+
+impl LuaUserData for MyUserData {}
+
+#[mlua::lua_module]
+fn rust_module_second(lua: &Lua) -> LuaResult<LuaTable> {
+    let exports = lua.create_table()?;
+    exports.set("userdata", lua.create_userdata(MyUserData(123))?)?;
     Ok(exports)
 }
diff --git a/src/hook.rs b/src/hook.rs
index 2946fc2..d534602 100644
--- a/src/hook.rs
+++ b/src/hook.rs
@@ -174,7 +174,7 @@ pub unsafe extern "C" fn mlua_hook_proc(state: *mut lua_State, ar: *mut lua_Debu
             _phantom: PhantomData,
         };
 
-        let lua = Lua::make_from_ptr(state);
+        let lua = mlua_expect!(Lua::make_from_ptr(state), "cannot make Lua instance");
         let hook_cb = mlua_expect!(lua.hook_callback(), "no hook callback set in hook_proc");
 
         #[allow(clippy::match_wild_err_arm)]
diff --git a/src/lua.rs b/src/lua.rs
index 4db5f93..6fec30b 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -369,6 +369,10 @@ impl Lua {
         let main_state = maybe_main_state.unwrap_or(state);
         let main_state_top = ffi::lua_gettop(main_state);
 
+        if let Some(lua) = Lua::make_from_ptr(state) {
+            return lua;
+        }
+
         let ref_thread = mlua_expect!(
             (|state| {
                 // Before initializing the error registry, we must set Error/Panic size.
@@ -437,8 +441,8 @@ impl Lua {
         );
         let extra_key = &EXTRA_REGISTRY_KEY as *const u8 as *const c_void;
         mlua_expect!(
-            ffi::safe::lua_rawsetp(main_state, ffi::LUA_REGISTRYINDEX, extra_key,),
-            "Error while storing extra data"
+            ffi::safe::lua_rawsetp(main_state, ffi::LUA_REGISTRYINDEX, extra_key),
+            "Error while storing extra data",
         );
 
         mlua_debug_assert!(
@@ -1972,12 +1976,14 @@ impl Lua {
         Ok(())
     }
 
-    pub(crate) unsafe fn make_from_ptr(state: *mut ffi::lua_State) -> Self {
+    pub(crate) unsafe fn make_from_ptr(state: *mut ffi::lua_State) -> Option<Self> {
         let _sg = StackGuard::new(state);
         assert_stack(state, 1);
 
         let extra_key = &EXTRA_REGISTRY_KEY as *const u8 as *const c_void;
-        ffi::lua_rawgetp(state, ffi::LUA_REGISTRYINDEX, extra_key);
+        if ffi::lua_rawgetp(state, ffi::LUA_REGISTRYINDEX, extra_key) != ffi::LUA_TUSERDATA {
+            return None;
+        }
         let extra = mlua_expect!(
             (*get_gc_userdata::<Weak<Mutex<ExtraData>>>(state, -1)).upgrade(),
             "extra is destroyed"
@@ -1986,14 +1992,14 @@ impl Lua {
 
         let safe = mlua_expect!(extra.lock(), "extra is poisoned").safe;
 
-        Lua {
+        Some(Lua {
             state,
             main_state: get_main_state(state),
             extra,
             ephemeral: true,
             safe,
             _no_ref_unwind_safe: PhantomData,
-        }
+        })
     }
 
     pub(crate) unsafe fn hook_callback(&self) -> Option<HookCallback> {
diff --git a/tests/module/tests/load.rs b/tests/module/tests/load.rs
index dedbe1f..5acad79 100644
--- a/tests/module/tests/load.rs
+++ b/tests/module/tests/load.rs
@@ -15,6 +15,19 @@ fn test_module() -> Result<()> {
     .exec()
 }
 
+#[test]
+fn test_module_multi() -> Result<()> {
+    let lua = make_lua()?;
+    lua.load(
+        r#"
+        local mod = require("rust_module")
+        local mod2 = require("rust_module.second")
+        assert(mod.check_userdata(mod2.userdata) == 123)
+    "#,
+    )
+    .exec()
+}
+
 #[cfg(any(
     feature = "lua54",
     feature = "lua53",