Add Lua::with_raw_state to provide easy low-level access to the Lua state.

Author: khvzak

src/lib.rs

@@ -109,8 +109,8 @@ pub use crate::error::{Error, ErrorContext, ExternalError, ExternalResult, Resul
 pub use crate::function::{Function, FunctionInfo};
 pub use crate::hook::{Debug, DebugEvent, DebugNames, DebugSource, DebugStack};
 pub use crate::multi::Variadic;
+pub use crate::scope::Scope;
 pub use crate::state::{GCMode, Lua, LuaOptions};
-// pub use crate::scope::Scope;
 pub use crate::stdlib::StdLib;
 pub use crate::string::{BorrowedBytes, BorrowedStr, String};
 pub use crate::table::{Table, TablePairs, TableSequence};

src/state.rs

@@ -22,7 +22,9 @@ use crate::types::{
     ReentrantMutex, ReentrantMutexGuard, RegistryKey, XRc, XWeak,
 };
 use crate::userdata::{AnyUserData, UserData, UserDataProxy, UserDataRegistry, UserDataStorage};
-use crate::util::{assert_stack, check_stack, push_string, push_table, rawset_field, StackGuard};
+use crate::util::{
+    assert_stack, check_stack, protect_lua_closure, push_string, push_table, rawset_field, StackGuard,
+};
 use crate::value::{FromLua, FromLuaMulti, IntoLua, IntoLuaMulti, MultiValue, Nil, Value};
 
 #[cfg(not(feature = "luau"))]
@@ -276,6 +278,28 @@ impl Lua {
         }
     }
 
+    /// Calls provided function passing a raw lua state.
+    ///
+    /// The arguments will be pushed onto the stack before calling the function.
+    ///
+    /// This method ensures that the Lua instance is locked while the function is called
+    /// and restores Lua stack after the function returns.
+    pub unsafe fn with_raw_state<R: FromLuaMulti>(
+        &self,
+        args: impl IntoLuaMulti,
+        f: impl FnOnce(*mut ffi::lua_State),
+    ) -> Result<R> {
+        let lua = self.lock();
+        let state = lua.state();
+        let _sg = StackGuard::new(state);
+        let stack_start = ffi::lua_gettop(state);
+        let nargs = args.push_into_stack_multi(&lua)?;
+        check_stack(state, 3)?;
+        protect_lua_closure::<_, ()>(state, nargs, ffi::LUA_MULTRET, f)?;
+        let nresults = ffi::lua_gettop(state) - stack_start;
+        R::from_stack_multi(nresults, &lua)
+    }
+
     /// FIXME: Deprecated load_from_std_lib
 
     /// Loads the specified subset of the standard libraries into an existing Lua state.

src/state/raw.rs

@@ -43,6 +43,7 @@ use {
 };
 
 /// An inner Lua struct which holds a raw Lua state.
+#[doc(hidden)]
 pub struct RawLua {
     // The state is dynamic and depends on context
     pub(super) state: Cell<*mut ffi::lua_State>,
@@ -83,8 +84,11 @@ impl RawLua {
         unsafe { (*self.extra.get()).weak() }
     }
 
+    /// Returns a pointer to the current Lua state.
+    ///
+    /// The pointer refers to the active Lua coroutine and depends on the context.
     #[inline(always)]
-    pub(crate) fn state(&self) -> *mut ffi::lua_State {
+    pub fn state(&self) -> *mut ffi::lua_State {
         self.state.get()
     }
 
@@ -500,10 +504,9 @@ impl RawLua {
 
     /// Pushes a value that implements `IntoLua` onto the Lua stack.
     ///
-    /// Uses 2 stack spaces, does not call checkstack.
-    #[doc(hidden)]
+    /// Uses up to 2 stack spaces to push a single value, does not call `checkstack`.
     #[inline(always)]
-    pub unsafe fn push(&self, value: impl IntoLua) -> Result<()> {
+    pub(crate) unsafe fn push(&self, value: impl IntoLua) -> Result<()> {
         value.push_into_stack(self)
     }
 

src/util/error.rs

@@ -204,24 +204,25 @@ pub(crate) unsafe fn protect_lua_closure<F, R>(
     f: F,
 ) -> Result<R>
 where
-    F: Fn(*mut ffi::lua_State) -> R,
+    F: FnOnce(*mut ffi::lua_State) -> R,
     R: Copy,
 {
     struct Params<F, R: Copy> {
-        function: F,
+        function: Option<F>,
         result: MaybeUninit<R>,
         nresults: c_int,
     }
 
     unsafe extern "C-unwind" fn do_call<F, R>(state: *mut ffi::lua_State) -> c_int
     where
-        F: Fn(*mut ffi::lua_State) -> R,
+        F: FnOnce(*mut ffi::lua_State) -> R,
         R: Copy,
     {
         let params = ffi::lua_touserdata(state, -1) as *mut Params<F, R>;
         ffi::lua_pop(state, 1);
 
-        (*params).result.write(((*params).function)(state));
+        let f = (*params).function.take().unwrap();
+        (*params).result.write(f(state));
 
         if (*params).nresults == ffi::LUA_MULTRET {
             ffi::lua_gettop(state)
@@ -241,7 +242,7 @@ where
     }
 
     let mut params = Params {
-        function: f,
+        function: Some(f),
         result: MaybeUninit::uninit(),
         nresults,
     };

src/util/mod.rs

@@ -69,6 +69,7 @@ impl StackGuard {
 }
 
 impl Drop for StackGuard {
+    #[track_caller]
     fn drop(&mut self) {
         unsafe {
             let top = ffi::lua_gettop(self.state);

tests/tests.rs

@@ -1289,3 +1289,37 @@ fn test_multi_thread() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_with_raw_state() -> Result<()> {
+    let lua = Lua::new();
+
+    let sum = lua.create_function(|_, args: Variadic<i32>| {
+        let mut sum = 0;
+        for i in args {
+            sum += i;
+        }
+        Ok(sum)
+    })?;
+    lua.globals().set("sum", sum)?;
+
+    let n: i32 = unsafe {
+        lua.with_raw_state((), |state| {
+            ffi::lua_getglobal(state, b"sum\0".as_ptr() as _);
+            ffi::lua_pushinteger(state, 1);
+            ffi::lua_pushinteger(state, 7);
+            ffi::lua_call(state, 2, 1);
+        })
+    }?;
+    assert_eq!(n, 8);
+
+    // Test error handling
+    let res: Result<()> = unsafe {
+        lua.with_raw_state("test error", |state| {
+            ffi::lua_error(state);
+        })
+    };
+    assert!(matches!(res, Err(Error::RuntimeError(err)) if err.contains("test error")));
+
+    Ok(())
+}