Add Lua::yield_with to allow yielding Rust async functions and exchange values between Lua coroutine and Rust.

This functionality is similar to `coroutine.yield` and `coroutine.resume` without C restrictions.

Author: khvzak

src/state.rs

@@ -37,6 +37,7 @@ use crate::{buffer::Buffer, chunk::Compiler};
 use {
     crate::types::LightUserData,
     std::future::{self, Future},
+    std::task::Poll,
 };
 
 #[cfg(feature = "serde")]
@@ -2079,6 +2080,101 @@ impl Lua {
         LightUserData(&ASYNC_POLL_TERMINATE as *const u8 as *mut std::os::raw::c_void)
     }
 
+    #[cfg(feature = "async")]
+    #[inline(always)]
+    pub(crate) fn poll_yield() -> LightUserData {
+        static ASYNC_POLL_YIELD: u8 = 0;
+        LightUserData(&ASYNC_POLL_YIELD as *const u8 as *mut std::os::raw::c_void)
+    }
+
+    /// Suspends the current async function, returning the provided arguments to caller.
+    ///
+    /// This function is similar to [`coroutine.yield`] but allow yeilding Rust functions
+    /// and passing values to the caller.
+    /// Please note that you cannot cross [`Thread`] boundaries (e.g. calling `yield_with` on one
+    /// thread and resuming on another).
+    ///
+    /// # Examples
+    ///
+    /// Async iterator:
+    ///
+    /// ```
+    /// # use mlua::{Lua, Result};
+    ///
+    /// async fn generator(lua: Lua, _: ()) -> Result<()> {
+    ///     for i in 0..10 {
+    ///         lua.yield_with::<()>(i).await?;
+    ///     }
+    ///     Ok(())
+    /// }
+    ///
+    /// fn main() -> Result<()> {
+    ///     let lua = Lua::new();
+    ///     lua.globals().set("generator", lua.create_async_function(generator)?)?;
+    ///
+    ///     lua.load(r#"
+    ///        local n = 0
+    ///        for i in coroutine.wrap(generator) do
+    ///            n = n + i
+    ///        end
+    ///        assert(n == 45)
+    ///     "#)
+    ///     .exec()
+    /// }
+    /// ```
+    ///
+    /// Exchange values on yield:
+    ///
+    /// ```
+    /// # use mlua::{Lua, Result, Value};
+    ///
+    /// async fn pingpong(lua: Lua, mut val: i32) -> Result<()> {
+    ///     loop {
+    ///         val = lua.yield_with::<i32>(val).await? + 1;
+    ///     }
+    ///     Ok(())
+    /// }
+    ///
+    /// # fn main() -> Result<()> {
+    /// let lua = Lua::new();
+    ///
+    /// let co = lua.create_thread(lua.create_async_function(pingpong)?)?;
+    /// assert_eq!(co.resume::<i32>(1)?, 1);
+    /// assert_eq!(co.resume::<i32>(2)?, 3);
+    /// assert_eq!(co.resume::<i32>(3)?, 4);
+    ///
+    /// # Ok(())
+    /// # }
+    /// ```
+    ///
+    /// [`coroutine.yield`]: https://www.lua.org/manual/5.4/manual.html#pdf-coroutine.yield
+    #[cfg(feature = "async")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "async")))]
+    pub async fn yield_with<R: FromLuaMulti>(&self, args: impl IntoLuaMulti) -> Result<R> {
+        let mut args = Some(args.into_lua_multi(self)?);
+        future::poll_fn(move |_cx| match args.take() {
+            Some(args) => unsafe {
+                let lua = self.lock();
+                lua.push(Self::poll_yield())?; // yield marker
+                if args.len() <= 1 {
+                    lua.push(args.front())?;
+                } else {
+                    lua.push(lua.create_sequence_from(&args)?)?;
+                }
+                lua.push(args.len())?;
+                Poll::Pending
+            },
+            None => unsafe {
+                let lua = self.lock();
+                let state = lua.state();
+                let _sg = StackGuard::with_top(state, 0);
+                let nvals = ffi::lua_gettop(state);
+                Poll::Ready(R::from_stack_multi(nvals, &lua))
+            },
+        })
+        .await
+    }
+
     /// Returns a weak reference to the Lua instance.
     ///
     /// This is useful for creating a reference to the Lua instance that does not prevent it from

src/state/raw.rs

@@ -1278,6 +1278,13 @@ impl RawLua {
                 let mut ctx = Context::from_waker(rawlua.waker());
                 match fut.as_mut().map(|fut| fut.as_mut().poll(&mut ctx)) {
                     Some(Poll::Pending) => {
+                        let fut_nvals = ffi::lua_gettop(state);
+                        if fut_nvals >= 3 && ffi::lua_tolightuserdata(state, -3) == Lua::poll_yield().0 {
+                            // We have some values to yield
+                            ffi::lua_pushnil(state);
+                            ffi::lua_replace(state, -4);
+                            return Ok(3);
+                        }
                         ffi::lua_pushnil(state);
                         ffi::lua_pushlightuserdata(state, Lua::poll_pending().0);
                         Ok(2)
@@ -1348,6 +1355,7 @@ impl RawLua {
             local poll = get_poll(...)
             local nres, res, res2 = poll()
             while true do
+                -- Poll::Ready branch, `nres` is the number of results
                 if nres ~= nil then
                     if nres == 0 then
                         return
@@ -1363,10 +1371,20 @@ impl RawLua {
                         return unpack(res, nres)
                     end
                 end
-                -- `res` is a "pending" value
-                -- `yield` can return a signal to drop the future that we should propagate
-                -- to the poller
-                nres, res, res2 = poll(yield(res))
+
+                -- Poll::Pending branch
+                if res2 == nil then
+                    -- `res` is a "pending" value
+                    -- `yield` can return a signal to drop the future that we should propagate
+                    -- to the poller
+                    nres, res, res2 = poll(yield(res))
+                elseif res2 == 0 then
+                    nres, res, res2 = poll(yield())
+                elseif res2 == 1 then
+                    nres, res, res2 = poll(yield(res))
+                else
+                    nres, res, res2 = poll(yield(unpack(res, res2)))
+                end
             end
             "#,
         )
@@ -1378,14 +1396,14 @@ impl RawLua {
 
     #[cfg(feature = "async")]
     #[inline]
-    pub(crate) unsafe fn waker(&self) -> &Waker {
-        (*self.extra.get()).waker.as_ref()
+    pub(crate) fn waker(&self) -> &Waker {
+        unsafe { (*self.extra.get()).waker.as_ref() }
     }
 
     #[cfg(feature = "async")]
     #[inline]
-    pub(crate) unsafe fn set_waker(&self, waker: NonNull<Waker>) -> NonNull<Waker> {
-        mem::replace(&mut (*self.extra.get()).waker, waker)
+    pub(crate) fn set_waker(&self, waker: NonNull<Waker>) -> NonNull<Waker> {
+        unsafe { mem::replace(&mut (*self.extra.get()).waker, waker) }
     }
 }
 

src/thread.rs

@@ -604,7 +604,7 @@ impl<R: FromLuaMulti> Future for AsyncThread<R> {
 
             if status.is_yielded() {
                 if !(nresults == 1 && is_poll_pending(thread_state)) {
-                    // Ignore value returned via yield()
+                    // Ignore values returned via yield()
                     cx.waker().wake_by_ref();
                 }
                 return Poll::Pending;
@@ -635,7 +635,7 @@ struct WakerGuard<'lua, 'a> {
 impl<'lua, 'a> WakerGuard<'lua, 'a> {
     #[inline]
     pub fn new(lua: &'lua RawLua, waker: &'a Waker) -> Result<WakerGuard<'lua, 'a>> {
-        let prev = unsafe { lua.set_waker(NonNull::from(waker)) };
+        let prev = lua.set_waker(NonNull::from(waker));
         Ok(WakerGuard {
             lua,
             prev,
@@ -647,7 +647,7 @@ impl<'lua, 'a> WakerGuard<'lua, 'a> {
 #[cfg(feature = "async")]
 impl Drop for WakerGuard<'_, '_> {
     fn drop(&mut self) {
-        unsafe { self.lua.set_waker(self.prev) };
+        self.lua.set_waker(self.prev);
     }
 }
 

tests/async.rs

@@ -8,7 +8,7 @@ use futures_util::stream::TryStreamExt;
 use tokio::sync::Mutex;
 
 use mlua::{
-    Error, Function, Lua, LuaOptions, MultiValue, ObjectLike, Result, StdLib, Table, UserData,
+    Error, Function, Lua, LuaOptions, MultiValue, ObjectLike, Result, StdLib, Table, ThreadStatus, UserData,
     UserDataMethods, UserDataRef, Value,
 };
 
@@ -667,3 +667,35 @@ async fn test_async_hook() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_async_yield_with() -> Result<()> {
+    let lua = Lua::new();
+
+    let func = lua.create_async_function(|lua, (mut a, mut b): (i32, i32)| async move {
+        let zero = lua.yield_with::<MultiValue>(()).await?;
+        assert!(zero.is_empty());
+        let one = lua.yield_with::<MultiValue>(a + b).await?;
+        assert_eq!(one.len(), 1);
+
+        for _ in 0..3 {
+            (a, b) = lua.yield_with((a + b, a * b)).await?;
+        }
+        Ok((0, 0))
+    })?;
+
+    let thread = lua.create_thread(func)?;
+
+    let zero = thread.resume::<MultiValue>((2, 3))?; // function arguments
+    assert!(zero.is_empty());
+    let one = thread.resume::<i32>(())?; // value of "zero" is passed here
+    assert_eq!(one, 5);
+
+    assert_eq!(thread.resume::<(i32, i32)>(1)?, (5, 6)); // value of "one" is passed here
+    assert_eq!(thread.resume::<(i32, i32)>((10, 11))?, (21, 110));
+    assert_eq!(thread.resume::<(i32, i32)>((11, 12))?, (23, 132));
+    assert_eq!(thread.resume::<(i32, i32)>((12, 13))?, (0, 0));
+    assert_eq!(thread.status(), ThreadStatus::Finished);
+
+    Ok(())
+}