Use list of auxillary threads instead of single auxillary thread

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/buffer.rs

@@ -61,7 +61,7 @@ impl Buffer {
     unsafe fn as_raw_parts(&self) -> (*mut u8, usize) {
         let lua = self.0.lua.lock();
         let mut size = 0usize;
-        let buf = ffi::lua_tobuffer(lua.ref_thread(), self.0.index, &mut size);
+        let buf = ffi::lua_tobuffer(lua.ref_thread(self.0.aux_thread), self.0.index, &mut size);
         mlua_assert!(!buf.is_null(), "invalid Luau buffer");
         (buf as *mut u8, size)
     }

src/conversion.rs

@@ -12,6 +12,7 @@ use num_traits::cast;
 
 use crate::error::{Error, Result};
 use crate::function::Function;
+use crate::state::util::get_next_spot;
 use crate::state::{Lua, RawLua};
 use crate::string::{BorrowedBytes, BorrowedStr, String};
 use crate::table::Table;
@@ -83,8 +84,12 @@ impl FromLua for String {
         let state = lua.state();
         let type_id = ffi::lua_type(state, idx);
         if type_id == ffi::LUA_TSTRING {
-            ffi::lua_xpush(state, lua.ref_thread(), idx);
-            return Ok(String(lua.pop_ref_thread()));
+            let (aux_thread, idxs, replace) = get_next_spot(lua.extra());
+            ffi::lua_xpush(state, lua.ref_thread(aux_thread), idx);
+            if replace {
+                ffi::lua_replace(lua.ref_thread(aux_thread), idxs);
+            }
+            return Ok(String(lua.new_value_ref(aux_thread, idxs)));
         }
         // Fallback to default
         Self::from_lua(lua.stack_value(idx, Some(type_id)), lua.lua())

src/error.rs

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

src/function.rs

@@ -3,6 +3,8 @@ use std::os::raw::{c_int, c_void};
 use std::{mem, ptr, slice};
 
 use crate::error::{Error, Result};
+#[cfg(feature = "luau")]
+use crate::state::util::get_next_spot;
 use crate::state::Lua;
 use crate::table::Table;
 use crate::traits::{FromLuaMulti, IntoLua, IntoLuaMulti, LuaNativeFn, LuaNativeFnMut};
@@ -494,14 +496,22 @@ impl Function {
     #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
     pub fn deep_clone(&self) -> Self {
         let lua = self.0.lua.lock();
-        let ref_thread = lua.ref_thread();
+        let ref_thread = lua.ref_thread(self.0.aux_thread);
         unsafe {
             if ffi::lua_iscfunction(ref_thread, self.0.index) != 0 {
                 return self.clone();
             }
 
             ffi::lua_clonefunction(ref_thread, self.0.index);
-            Function(lua.pop_ref_thread())
+
+            // Get the real next spot
+            let (aux_thread, index, replace) = get_next_spot(lua.extra());
+            ffi::lua_xpush(lua.ref_thread(self.0.aux_thread), lua.ref_thread(aux_thread), -1);
+            if replace {
+                ffi::lua_replace(lua.ref_thread(aux_thread), index);
+            }
+
+            Function(lua.new_value_ref(aux_thread, index))
         }
     }
 }

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/scope.rs

@@ -267,7 +267,7 @@ impl<'scope, 'env: 'scope> Scope<'scope, 'env> {
         let f = self.lua.create_callback(f)?;
 
         let destructor: DestructorCallback = Box::new(|rawlua, vref| {
-            let ref_thread = rawlua.ref_thread();
+            let ref_thread = rawlua.ref_thread(vref.aux_thread);
             ffi::lua_getupvalue(ref_thread, vref.index, 1);
             let upvalue = get_userdata::<CallbackUpvalue>(ref_thread, -1);
             let data = (*upvalue).data.take();
@@ -287,13 +287,13 @@ impl<'scope, 'env: 'scope> Scope<'scope, 'env> {
                 Ok(Some(_)) => {}
                 Ok(None) => {
                     // Deregister metatable
-                    let mt_ptr = get_metatable_ptr(rawlua.ref_thread(), vref.index);
+                    let mt_ptr = get_metatable_ptr(rawlua.ref_thread(vref.aux_thread), vref.index);
                     rawlua.deregister_userdata_metatable(mt_ptr);
                 }
                 Err(_) => return vec![],
             }
 
-            let data = take_userdata::<UserDataStorage<T>>(rawlua.ref_thread(), vref.index);
+            let data = take_userdata::<UserDataStorage<T>>(rawlua.ref_thread(vref.aux_thread), vref.index);
             vec![Box::new(move || drop(data))]
         });
         self.destructors.0.borrow_mut().push((ud.0.clone(), destructor));

src/serde/mod.rs

@@ -7,6 +7,7 @@ use serde::ser::Serialize;
 
 use crate::error::Result;
 use crate::private::Sealed;
+use crate::state::util::get_next_spot;
 use crate::state::Lua;
 use crate::table::Table;
 use crate::util::check_stack;
@@ -183,8 +184,14 @@ impl LuaSerdeExt for Lua {
     fn array_metatable(&self) -> Table {
         let lua = self.lock();
         unsafe {
-            push_array_metatable(lua.ref_thread());
-            Table(lua.pop_ref_thread())
+            let (aux_thread, index, replace) = get_next_spot(lua.extra());
+            push_array_metatable(lua.state());
+            ffi::lua_xmove(lua.state(), lua.ref_thread(aux_thread), 1);
+            if replace {
+                ffi::lua_replace(lua.ref_thread(aux_thread), index);
+            }
+
+            Table(lua.new_value_ref(aux_thread, index))
         }
     }
 

src/state.rs

@@ -14,10 +14,15 @@ use crate::hook::Debug;
 use crate::memory::MemoryState;
 use crate::multi::MultiValue;
 use crate::scope::Scope;
+use crate::state::util::get_next_spot;
 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,
@@ -522,7 +527,7 @@ impl Lua {
                         ffi::luaL_sandboxthread(state);
                     } else {
                         // Restore original `LUA_GLOBALSINDEX`
-                        ffi::lua_xpush(lua.ref_thread(), state, ffi::LUA_GLOBALSINDEX);
+                        ffi::lua_xpush(lua.ref_thread_internal(), state, ffi::LUA_GLOBALSINDEX);
                         ffi::lua_replace(state, ffi::LUA_GLOBALSINDEX);
                         ffi::luaL_sandbox(state, 0);
                     }
@@ -762,8 +767,12 @@ impl Lua {
                 return; // Don't allow recursion
             }
             ffi::lua_pushthread(child);
-            ffi::lua_xmove(child, (*extra).ref_thread, 1);
-            let value = Thread((*extra).raw_lua().pop_ref_thread(), child);
+            let (aux_thread, index, replace) = get_next_spot(extra);
+            ffi::lua_xmove(child, (*extra).raw_lua().ref_thread(aux_thread), 1);
+            if replace {
+                ffi::lua_replace((*extra).raw_lua().ref_thread(aux_thread), index);
+            }
+            let value = Thread((*extra).raw_lua().new_value_ref(aux_thread, index), child);
             callback_error_ext(parent, extra, false, move |extra, _| {
                 callback((*extra).lua(), value)
             })
@@ -1265,6 +1274,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.
@@ -1286,8 +1333,13 @@ impl Lua {
     /// This function is unsafe because provides a way to execute unsafe C function.
     pub unsafe fn create_c_function(&self, func: ffi::lua_CFunction) -> Result<Function> {
         let lua = self.lock();
-        ffi::lua_pushcfunction(lua.ref_thread(), func);
-        Ok(Function(lua.pop_ref_thread()))
+        let (aux_thread, idx, replace) = get_next_spot(lua.extra());
+        ffi::lua_pushcfunction(lua.ref_thread(aux_thread), func);
+        if replace {
+            ffi::lua_replace(lua.ref_thread(aux_thread), idx);
+        }
+
+        Ok(Function(lua.new_value_ref(aux_thread, idx)))
     }
 
     /// Wraps a Rust async function or closure, creating a callable Lua function handle to it.
@@ -2080,6 +2132,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 {