Add support for lua yields and basic continuation support

mlua-sys/src/lua52/lua.rs

@@ -272,6 +272,11 @@ pub unsafe fn lua_yield(L: *mut lua_State, n: c_int) -> c_int {
     lua_yieldk(L, n, 0, None)
 }
 
+#[inline(always)]
+pub unsafe fn lua_yieldc(L: *mut lua_State, n: c_int, k: lua_CFunction) -> c_int {
+    lua_yieldk(L, n, 0, Some(k))
+}
+
 //
 // Garbage-collection function and options
 //

mlua-sys/src/lua53/lua.rs

@@ -286,6 +286,11 @@ pub unsafe fn lua_yield(L: *mut lua_State, n: c_int) -> c_int {
     lua_yieldk(L, n, 0, None)
 }
 
+#[inline(always)]
+pub unsafe fn lua_yieldc(L: *mut lua_State, n: c_int, k: lua_KFunction) -> c_int {
+    lua_yieldk(L, n, 0, Some(k))
+}
+
 //
 // Garbage-collection function and options
 //

mlua-sys/src/lua54/lua.rs

@@ -299,6 +299,11 @@ pub unsafe fn lua_yield(L: *mut lua_State, n: c_int) -> c_int {
     lua_yieldk(L, n, 0, None)
 }
 
+#[inline(always)]
+pub unsafe fn lua_yieldc(L: *mut lua_State, n: c_int, k: lua_KFunction) -> c_int {
+    lua_yieldk(L, n, 0, Some(k))
+}
+
 //
 // Warning-related functions
 //

mlua-sys/src/luau/lua.rs

@@ -426,6 +426,11 @@ pub unsafe fn lua_pushcclosure(L: *mut lua_State, f: lua_CFunction, nup: c_int)
     lua_pushcclosurek(L, f, ptr::null(), nup, None)
 }
 
+#[inline(always)]
+pub unsafe fn lua_pushcclosurec(L: *mut lua_State, f: lua_CFunction, cont: lua_Continuation, nup: c_int) {
+    lua_pushcclosurek(L, f, ptr::null(), nup, Some(cont))
+}
+
 #[inline(always)]
 pub unsafe fn lua_pushcclosured(L: *mut lua_State, f: lua_CFunction, debugname: *const c_char, nup: c_int) {
     lua_pushcclosurek(L, f, debugname, nup, None)

src/error.rs

@@ -321,7 +321,7 @@ impl fmt::Display for Error {
             Error::WithContext { context, cause } => {
                 writeln!(fmt, "{context}")?;
                 write!(fmt, "{cause}")
-            }
+            },
         }
     }
 }

src/lib.rs

@@ -110,7 +110,7 @@ pub use crate::state::{GCMode, Lua, LuaOptions, WeakLua};
 pub use crate::stdlib::StdLib;
 pub use crate::string::{BorrowedBytes, BorrowedStr, String};
 pub use crate::table::{Table, TablePairs, TableSequence};
-pub use crate::thread::{Thread, ThreadStatus};
+pub use crate::thread::{ContinuationStatus, Thread, ThreadStatus};
 pub use crate::traits::{
     FromLua, FromLuaMulti, IntoLua, IntoLuaMulti, LuaNativeFn, LuaNativeFnMut, ObjectLike,
 };

src/prelude.rs

@@ -3,15 +3,15 @@
 #[doc(no_inline)]
 pub use crate::{
     AnyUserData as LuaAnyUserData, BorrowedBytes as LuaBorrowedBytes, BorrowedStr as LuaBorrowedStr,
-    Chunk as LuaChunk, Either as LuaEither, Error as LuaError, ErrorContext as LuaErrorContext,
-    ExternalError as LuaExternalError, ExternalResult as LuaExternalResult, FromLua, FromLuaMulti,
-    Function as LuaFunction, FunctionInfo as LuaFunctionInfo, GCMode as LuaGCMode, Integer as LuaInteger,
-    IntoLua, IntoLuaMulti, LightUserData as LuaLightUserData, Lua, LuaNativeFn, LuaNativeFnMut, LuaOptions,
-    MetaMethod as LuaMetaMethod, MultiValue as LuaMultiValue, Nil as LuaNil, Number as LuaNumber,
-    ObjectLike as LuaObjectLike, RegistryKey as LuaRegistryKey, Result as LuaResult, StdLib as LuaStdLib,
-    String as LuaString, Table as LuaTable, TablePairs as LuaTablePairs, TableSequence as LuaTableSequence,
-    Thread as LuaThread, ThreadStatus as LuaThreadStatus, UserData as LuaUserData,
-    UserDataFields as LuaUserDataFields, UserDataMetatable as LuaUserDataMetatable,
+    Chunk as LuaChunk, ContinuationStatus as LuaContinuationStatus, Either as LuaEither, Error as LuaError,
+    ErrorContext as LuaErrorContext, ExternalError as LuaExternalError, ExternalResult as LuaExternalResult,
+    FromLua, FromLuaMulti, Function as LuaFunction, FunctionInfo as LuaFunctionInfo, GCMode as LuaGCMode,
+    Integer as LuaInteger, IntoLua, IntoLuaMulti, LightUserData as LuaLightUserData, Lua, LuaNativeFn,
+    LuaNativeFnMut, LuaOptions, MetaMethod as LuaMetaMethod, MultiValue as LuaMultiValue, Nil as LuaNil,
+    Number as LuaNumber, ObjectLike as LuaObjectLike, RegistryKey as LuaRegistryKey, Result as LuaResult,
+    StdLib as LuaStdLib, String as LuaString, Table as LuaTable, TablePairs as LuaTablePairs,
+    TableSequence as LuaTableSequence, Thread as LuaThread, ThreadStatus as LuaThreadStatus,
+    UserData as LuaUserData, UserDataFields as LuaUserDataFields, UserDataMetatable as LuaUserDataMetatable,
     UserDataMethods as LuaUserDataMethods, UserDataRef as LuaUserDataRef,
     UserDataRefMut as LuaUserDataRefMut, UserDataRegistry as LuaUserDataRegistry, Value as LuaValue,
     Variadic as LuaVariadic, VmState as LuaVmState, WeakLua,

src/state.rs

@@ -18,6 +18,10 @@ use crate::stdlib::StdLib;
 use crate::string::String;
 use crate::table::Table;
 use crate::thread::Thread;
+
+#[cfg(all(not(feature = "lua51"), not(feature = "luajit")))]
+use crate::thread::ContinuationStatus;
+
 use crate::traits::{FromLua, FromLuaMulti, IntoLua, IntoLuaMulti};
 use crate::types::{
     AppDataRef, AppDataRefMut, ArcReentrantMutexGuard, Integer, LuaType, MaybeSend, Number, ReentrantMutex,
@@ -1265,6 +1269,44 @@ impl Lua {
         }))
     }
 
+    /// Same as ``create_function`` but with an added continuation function.
+    ///
+    /// The values passed to the continuation will be the yielded arguments
+    /// from the function for the initial continuation call. If yielding from a
+    /// continuation, the yielded results will be returned to the ``Thread::resume`` caller. The
+    /// arguments passed in the next ``Thread::resume`` call will then be the arguments passed
+    /// to the yielding continuation upon resumption.
+    ///
+    /// Returning a value from a continuation without setting yield
+    /// arguments will then be returned as the final return value of the Lua function call.
+    /// Values returned in a function in which there is also yielding will be ignored
+    #[cfg(all(not(feature = "lua51"), not(feature = "luajit")))]
+    pub fn create_function_with_continuation<F, FC, A, AC, R, RC>(
+        &self,
+        func: F,
+        cont: FC,
+    ) -> Result<Function>
+    where
+        F: Fn(&Lua, A) -> Result<R> + MaybeSend + 'static,
+        FC: Fn(&Lua, ContinuationStatus, AC) -> Result<RC> + MaybeSend + 'static,
+        A: FromLuaMulti,
+        AC: FromLuaMulti,
+        R: IntoLuaMulti,
+        RC: IntoLuaMulti,
+    {
+        (self.lock()).create_callback_with_continuation(
+            Box::new(move |rawlua, nargs| unsafe {
+                let args = A::from_stack_args(nargs, 1, None, rawlua)?;
+                func(rawlua.lua(), args)?.push_into_stack_multi(rawlua)
+            }),
+            Box::new(move |rawlua, nargs, status| unsafe {
+                let args = AC::from_stack_args(nargs, 1, None, rawlua)?;
+                let status = ContinuationStatus::from_status(status);
+                cont(rawlua.lua(), status, args)?.push_into_stack_multi(rawlua)
+            }),
+        )
+    }
+
     /// Wraps a Rust mutable closure, creating a callable Lua function handle to it.
     ///
     /// This is a version of [`Lua::create_function`] that accepts a `FnMut` argument.
@@ -2080,6 +2122,42 @@ impl Lua {
     pub(crate) unsafe fn raw_lua(&self) -> &RawLua {
         &*self.raw.data_ptr()
     }
+
+    /// Set the yield arguments. Note that Lua will not yield until you return from the function
+    ///
+    /// This method is mostly useful with continuations and Rust-Rust yields
+    /// due to the Rust/Lua boundary
+    ///
+    /// Example:
+    ///
+    /// ```rust
+    /// fn test() -> mlua::Result<()> {
+    ///     let lua = mlua::Lua::new();
+    ///     let always_yield = lua.create_function(|lua, ()| lua.yield_with((42, "69420".to_string(), 45.6)))?;
+    ///
+    ///     let thread = lua.create_thread(always_yield)?;
+    ///     assert_eq!(
+    ///         thread.resume::<(i32, String, f32)>(())?,
+    ///         (42, String::from("69420"), 45.6)
+    ///     );
+    ///
+    ///     Ok(())
+    /// }
+    /// ```
+    pub fn yield_with(&self, args: impl IntoLuaMulti) -> Result<()> {
+        let raw = self.lock();
+        unsafe {
+            raw.extra.get().as_mut().unwrap_unchecked().yielded_values = Some(args.into_lua_multi(self)?);
+        }
+        Ok(())
+    }
+
+    /// Checks if Lua is be allowed to yield.
+    #[cfg(not(any(feature = "lua51", feature = "lua52", feature = "luajit")))]
+    #[inline]
+    pub fn is_yieldable(&self) -> bool {
+        self.lock().is_yieldable()
+    }
 }
 
 impl WeakLua {

src/state/extra.rs

@@ -13,11 +13,13 @@ use crate::error::Result;
 use crate::state::RawLua;
 use crate::stdlib::StdLib;
 use crate::types::{AppData, ReentrantMutex, XRc};
+
 use crate::userdata::RawUserDataRegistry;
 use crate::util::{get_internal_metatable, push_internal_userdata, TypeKey, WrappedFailure};
 
 #[cfg(any(feature = "luau", doc))]
 use crate::chunk::Compiler;
+use crate::MultiValue;
 
 #[cfg(feature = "async")]
 use {futures_util::task::noop_waker_ref, std::ptr::NonNull, std::task::Waker};
@@ -94,6 +96,9 @@ pub(crate) struct ExtraData {
     pub(super) compiler: Option<Compiler>,
     #[cfg(feature = "luau-jit")]
     pub(super) enable_jit: bool,
+
+    // Values currently being yielded from Lua.yield()
+    pub(super) yielded_values: Option<MultiValue>,
 }
 
 impl Drop for ExtraData {
@@ -196,6 +201,7 @@ impl ExtraData {
             enable_jit: true,
             #[cfg(feature = "luau")]
             running_gc: false,
+            yielded_values: None,
         }));
 
         // Store it in the registry