[PoC] Add support for async "coroutine" systems

[CAD97]

Objective

• Allow the use of async to write systems that span multiple frames
• Similar to that discussed in https://github.com/bevyengine/bevy/discussions/2677
• Related to https://github.com/bevyengine/bevy/pull/2650

Solution

This PR implements a proof-of-concept for "coroutine systems" defined using async. Because it is not sound to hold a reference to ECS resources over an await point, access to the system parameters is limited to during a sync closure through a co_with! macro.

This abuses the fact that macro names are not hygienic to define the co_with! macro used to access system parameters. This scheme allows minimizing the annotation overhead required of the user, importantly allowing full elision of lifetimes just as with normal function systems.

I did attempt to skip the need to re-name the system parameters when entering co_with! (essentially getting the "magic mutation" syntax for yield closures), but unfortunately for this use case, naming the parameters within the closure always refers to the outer names, not the newly bound names due to the new semitransparent macro_rules! layer.

The macro techniques to do so

diff --git a/crates/bevy_async/src/lib.rs b/crates/bevy_async/src/lib.rs
index 5fbeae7f..166a3a76 100644
--- a/crates/bevy_async/src/lib.rs
+++ b/crates/bevy_async/src/lib.rs
@@ -53,14 +53,14 @@
 //!     co!(async move |state: Local<usize>| loop {
 //!         let name = String::from("bevy_async");
 //!         // Access the parameters by using co_with!
-//!         let _ = co_with!(|state| {
+//!         let _ = co_with! {
 //!             // Within this closure, the parameters are available.
 //!             println!("The local's current value is {}", *state);
 //!             // The closure borrows normally from the containing scope.
 //!             println!("The name is {}", name);
 //!             // The closure can pass state back to the async context.
 //!             (&*name, *state)
-//!         });
+//!         };
 //!         // Outside co_with! is an async context where you can use await.
 //!         // It's best practice to spawn async work onto a task pool
 //!         // and await the task handle to minimize redundant polling.
@@ -370,8 +370,12 @@ macro_rules! co {
                     let ($($arg,)*) = unsafe { co.fetch_params_unchecked() };
                     f($($arg,)*)
                 }
-                macro_rules! co_with {($with:expr) => {
-                    co_with(&mut co, $with)
+                $crate::__item_with_dollar! {($_:tt) => {
+                    macro_rules! co_with {($_($_ with:tt)*) => {
+                        co_with(&mut co, |$($arg),*| {
+                            $_($_ with)*
+                        })
+                    }}
                 }}
                 loop {
                     $body;
@@ -383,3 +387,10 @@ macro_rules! co {
         )
     };
 }
+
+#[doc(hidden)]
+#[macro_export]
+macro_rules! __item_with_dollar {(($($in:tt)*) => {$($out:tt)*}) => {
+    macro_rules! __emit {($($in)*) => {$($out)*}}
+    __emit! {$}
+}}

error[E0381]: used binding `state` isn't initialized
  --> src\lib.rs:56:17
   |
9  |  /     co!(async move |state: Local<usize>| loop {
10 |  |         let name = String::from("bevy_async");
11 |  |         // Access the parameters by using co_with!
12 |  |         let _ = co_with! {
   |  |_________________^
13 | ||             // Within this closure, the parameters are available.
14 | ||             println!("The local's current value is {}", *state);
   | ||                                                          ----- borrow occurs due to use in closure
15 | ||             // The closure borrows normally from the containing scope.
...  ||
18 | ||             (&*name, *state)
19 | ||         };
   | ||_________^ `state` used here but it isn't initialized
...   |
26 |  |         }).await;
27 |  |     })
   |  |______- binding declared here but left uninitialized
   |
   = note: this error originates in the macro `co_with` (in Nightly builds, run with -Z macro-backtrace for more info)

The reason this doesn't work seems to be that the co_with! definition introduces a fresh scope which the expanded idents are attached to, rather than the intent in this case of keeping their originally captured context. I've tried a number of permutations but was unable to get the desired behavior here.

---

This can mostly be implemented out-of-tree, but requires making a few more things accessible from bevy_ecs to be able to implement System properly.

[CAD97]

cc @hanabi1224, if you're interested.

[jakobhellermann]

I also tried using the nightly generators feature, but unfortunately I don't think that can work.

    || {
        let (asset_server, images) = yield;
        // do stuff
        let (_asset_server, _images) = yield;
        dbg!(&images); // unfortunately still live here
    }

I thought that the system parameters could be resumed into the generator through yield points, but unfortunately old yielded variables are still live after the next yield.

[CAD97]

feature(generators)

Yeah; using FnPin to illustrate, nightly generators are FnPin(&'exists Args) -> Output, whereas we need for<'all> FnPin(&'all Args) -> Output.

It would be theoretically possible to implement an async state machine desugar in a proc macro, but that's a significant engineering effort that I don't think is worth the effort. Rather, the general use case needs to wait on language-level yield closures that support short-lived arguments. (Streaming is a big use case for them though, so I expect it'll probably happen eventually.)

why require the loop?

Mostly, this is in order to be explicit and avoid making a choice. There's a split on FnPin semantics for what happens after a return; I and the referenced MCP prefer resuming back at the top (this matches the semantics of straightline closures), but it's also a reasonable option to poison the closure (matching futures) making further resumption a panic.

(The former is imho a better option for -> Output, the latter for -> GeneratorState, as the latter differentiates yield and return, but the former doesn't. See the lang team design notes for more context on FnPin.)

Easy to write code that doesn't work

I don't think poll_fn(|_| poll) is that much of a footgun comparatively; anyone who knows what poll_fn does enough to use it would (hopefully) know that this wouldn't recalculate the readiness.

If this async model is adopted more, I suspect we'd add something along the lines of server.poll_ready(&handle)... though needing to write it in a loop still isn't optimal.

loop {
    co_with!(|server, _| server.poll_ready(&handle)).await
}

If bevy were to adopt async even more (though I probably wouldn't recommend it for bevy), we could provide some sort of async channel to notify readiness with a single await.

IDE

Yeah, it's quite nice that the transformation is simple enough that simple IDE functionality works. It's probably worth throwing together the #[coroutine] attribute to see if that improves anything... though being a proc macro it's inherently more difficult to analyze...

3rd party crate

Yeah, although the additional public API surface to bevy_ecs (or similar) in this PR is necessary (as far as I can tell) to be able to implement System for this externally.