Add resume_error to Thread for luau

src/thread.rs

@@ -8,6 +8,7 @@ use crate::traits::{FromLuaMulti, IntoLuaMulti};
 use crate::types::{LuaType, ValueRef};
 use crate::util::{check_stack, error_traceback_thread, pop_error, StackGuard};
 
+use crate::IntoLua;
 #[cfg(not(feature = "luau"))]
 use crate::{
     hook::{Debug, HookTriggers},
@@ -173,6 +174,38 @@ impl Thread {
         }
     }
 
+    #[cfg(feature = "luau")]
+    /// Resumes a thread with an error. This is useful when developing
+    /// custom async schedulers' etc. with mlua as it allows bubbling up
+    /// errors from the yielded thread upwards with working pcall out of
+    /// the box.
+    pub fn resume_error<R>(&self, args: impl IntoLua) -> Result<R>
+    where
+        R: FromLuaMulti,
+    {
+        let lua = self.0.lua.lock();
+        match self.status_inner(&lua) {
+            ThreadStatusInner::New(_) | ThreadStatusInner::Yielded(_) => {}
+            _ => return Err(Error::CoroutineUnresumable),
+        };
+
+        let state = lua.state();
+        let thread_state = self.state();
+        unsafe {
+            let _sg = StackGuard::new(state);
+            let _thread_sg = StackGuard::with_top(thread_state, 0);
+
+            check_stack(state, 1)?;
+            args.push_into_stack(&lua)?;
+            ffi::lua_xmove(state, thread_state, 1);
+
+            let (_, nresults) = self.resumeerror_inner(&lua)?;
+            check_stack(state, nresults + 1)?;
+            ffi::lua_xmove(thread_state, state, nresults);
+            R::from_stack_multi(nresults, &lua)
+        }
+    }
+
     /// Resumes execution of this thread.
     ///
     /// It's similar to `resume()` but leaves `nresults` values on the thread stack.
@@ -196,6 +229,35 @@ impl Thread {
         }
     }
 
+    #[cfg(feature = "luau")]
+    /// Resumes execution of this thread with an error.
+    ///
+    /// It's similar to `resume_error()` but leaves `nresults` values on the thread stack.
+    unsafe fn resumeerror_inner(&self, lua: &RawLua) -> Result<(ThreadStatusInner, c_int)> {
+        let state = lua.state();
+        let thread_state = self.state();
+        let mut nresults = 0;
+        let ret = ffi::luau::lua_resumeerror(thread_state, state);
+
+        if ret == ffi::LUA_OK || ret == ffi::LUA_YIELD {
+            nresults = ffi::lua_gettop(thread_state);
+        }
+
+        match ret {
+            ffi::LUA_OK => Ok((ThreadStatusInner::Finished, nresults)),
+            ffi::LUA_YIELD => Ok((ThreadStatusInner::Yielded(0), nresults)),
+            ffi::LUA_ERRMEM => {
+                // Don't call error handler for memory errors
+                Err(pop_error(thread_state, ret))
+            }
+            _ => {
+                check_stack(state, 3)?;
+                protect_lua!(state, 0, 1, |state| error_traceback_thread(state, thread_state))?;
+                Err(pop_error(state, ret))
+            }
+        }
+    }
+
     /// Gets the status of the thread.
     pub fn status(&self) -> ThreadStatus {
         match self.status_inner(&self.0.lua.lock()) {

tests/thread.rs

@@ -227,3 +227,75 @@ fn test_thread_pointer() -> Result<()> {
 
     Ok(())
 }
+
+#[cfg(feature = "luau")]
+#[cfg(test)]
+mod resumeerror_test {
+    #[test]
+    fn test_resumeerror() {
+        // Create tokio runtime and use spawn_local [this is required for the test to work]
+        let rt = tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .worker_threads(10)
+            .build()
+            .unwrap();
+
+        let local = tokio::task::LocalSet::new();
+
+        local.block_on(&rt, async {
+            use mlua::{Function, Lua, Thread};
+
+            let lua = Lua::new();
+
+            let thread: Function = lua
+                .load(
+                    r#"
+local luacall = ...
+local function callback(...)
+    luacall(coroutine.running(), ...)
+    return coroutine.yield()
+end
+
+return callback
+                "#,
+                )
+                .call(
+                    lua.create_function(|_lua, th: Thread| {
+                        tokio::task::spawn_local(async move {
+                            th.resume_error::<()>("ErrorABC".to_string()).unwrap();
+                            tokio::task::yield_now().await;
+                        });
+                        Ok(())
+                    })
+                    .unwrap(),
+                )
+                .unwrap();
+
+            let thread_b: Thread = lua
+                .load(
+                    r#"
+                    local a = ...
+                    return coroutine.create(function (...)
+                        local b = ...
+                        assert(b == 1)
+                        -- Test 1: Working pcall
+                        local ok, result = pcall(a)
+                        assert(not ok, "Should not be ok")
+                        assert(result == "ErrorABC", "Should be ErrorABC")
+
+                        -- Repeat
+                        local ok, result = pcall(a)
+                        assert(not ok)
+                        assert(result == "ErrorABC")
+                    end)
+                "#,
+                )
+                .call(thread.clone())
+                .unwrap();
+
+            // Actually using this system in practice needs a few extra pieces (such as a way to track return
+            // values etc.)
+            thread_b.resume::<()>(1).expect("Error in thread_b");
+        });
+    }
+}