Rename Lua::init_from_ptr to Lua::get_or_init_from_ptr.

Return reference to `&Lua` with unbound lifetime (chosen by user).

Author: khvzak

src/state.rs

@@ -261,16 +261,23 @@ impl Lua {
         lua
     }
 
-    /// Constructs a new Lua instance from an existing raw state.
+    /// Returns or constructs Lua instance from a raw state.
     ///
-    /// Once called, a returned Lua state is cached in the registry and can be retrieved
+    /// Once initialized, the returned Lua instance is cached in the registry and can be retrieved
     /// by calling this function again.
-    #[allow(clippy::missing_safety_doc)]
+    ///
+    /// # Safety
+    /// The `Lua` must outlive the chosen lifetime `'a`.
     #[inline]
-    pub unsafe fn init_from_ptr(state: *mut ffi::lua_State) -> Lua {
-        Lua {
-            raw: RawLua::init_from_ptr(state, false),
-            collect_garbage: true,
+    pub unsafe fn get_or_init_from_ptr<'a>(state: *mut ffi::lua_State) -> &'a Lua {
+        debug_assert!(!state.is_null(), "Lua state is null");
+        match ExtraData::get(state) {
+            extra if !extra.is_null() => (*extra).lua(),
+            _ => {
+                // The `owned` flag is set to `false` as we don't own the Lua state.
+                RawLua::init_from_ptr(state, false);
+                (*ExtraData::get(state)).lua()
+            }
         }
     }
 
@@ -410,11 +417,7 @@ impl Lua {
         R: IntoLua,
     {
         // Make sure that Lua is initialized
-        let mut lua = Self::init_from_ptr(state);
-        lua.collect_garbage = false;
-        // `Lua` is no longer needed and must be dropped at this point to avoid memory leak
-        // in case of possible longjmp (lua_error) below
-        drop(lua);
+        let _ = Self::get_or_init_from_ptr(state);
 
         callback_error_ext(state, ptr::null_mut(), move |extra, nargs| {
             let rawlua = (*extra).raw_lua();

src/state/extra.rs

@@ -195,9 +195,6 @@ impl ExtraData {
             raw: XRc::clone(raw),
             collect_garbage: false,
         });
-        if self.owned {
-            XRc::decrement_strong_count(XRc::as_ptr(raw));
-        }
         self.weak.write(WeakLua(XRc::downgrade(raw)));
     }
 

src/state/raw.rs

@@ -58,12 +58,13 @@ pub struct RawLua {
     pub(super) state: Cell<*mut ffi::lua_State>,
     pub(super) main_state: Option<NonNull<ffi::lua_State>>,
     pub(super) extra: XRc<UnsafeCell<ExtraData>>,
+    owned: bool,
 }
 
 impl Drop for RawLua {
     fn drop(&mut self) {
         unsafe {
-            if !(*self.extra.get()).owned {
+            if !self.owned {
                 return;
             }
 
@@ -233,8 +234,17 @@ impl RawLua {
             // Make sure that we don't store current state as main state (if it's not available)
             main_state: get_main_state(state).and_then(NonNull::new),
             extra: XRc::clone(&extra),
+            owned,
         }));
         (*extra.get()).set_lua(&rawlua);
+        if owned {
+            // If Lua state is managed by us, then make internal `RawLua` reference "weak"
+            XRc::decrement_strong_count(XRc::as_ptr(&rawlua));
+        } else {
+            // If Lua state is not managed by us, then keep internal `RawLua` reference "strong"
+            // but `Extra` reference weak (it will be collected from registry at lua_close time)
+            XRc::decrement_strong_count(XRc::as_ptr(&extra));
+        }
 
         rawlua
     }

tests/tests.rs

@@ -8,8 +8,8 @@ use std::sync::Arc;
 use std::{error, f32, f64, fmt};
 
 use mlua::{
-    ChunkMode, Error, ExternalError, Function, Lua, LuaOptions, Nil, Result, StdLib, String, Table, UserData,
-    Value, Variadic,
+    ffi, ChunkMode, Error, ExternalError, Function, Lua, LuaOptions, Nil, Result, StdLib, String, Table,
+    UserData, Value, Variadic,
 };
 
 #[cfg(not(feature = "luau"))]
@@ -1425,3 +1425,23 @@ fn test_gc_drop_ref_thread() -> Result<()> {
 
     Ok(())
 }
+
+#[cfg(not(feature = "luau"))]
+#[test]
+fn test_get_or_init_from_ptr() -> Result<()> {
+    // This would not work with Luau, the state must be init by mlua internally
+    let state = unsafe { ffi::luaL_newstate() };
+
+    let mut lua = unsafe { Lua::get_or_init_from_ptr(state) };
+    lua.globals().set("hello", "world678")?;
+
+    // The same Lua instance must be returned
+    lua = unsafe { Lua::get_or_init_from_ptr(state) };
+    assert_eq!(lua.globals().get::<String>("hello")?, "world678");
+
+    unsafe { ffi::lua_close(state) };
+
+    // Lua must not be accessed after closing
+
+    Ok(())
+}