SystemRunner param and macro syntax

crates/bevy_ecs/macros/Cargo.toml

@@ -11,7 +11,7 @@ proc-macro = true
 [dependencies]
 bevy_macro_utils = { path = "../../bevy_macro_utils", version = "0.18.0-dev" }
 
-syn = { version = "2.0.108", features = ["full", "extra-traits"] }
+syn = { version = "2.0.108", features = ["full", "extra-traits", "visit-mut"] }
 quote = "1.0"
 proc-macro2 = "1.0"
 [lints]

crates/bevy_ecs/macros/src/lib.rs

@@ -9,6 +9,7 @@ mod event;
 mod message;
 mod query_data;
 mod query_filter;
+mod system_composition;
 mod world_query;
 
 use crate::{
@@ -719,3 +720,85 @@ pub fn derive_from_world(input: TokenStream) -> TokenStream {
             }
     })
 }
+
+/// Automatically compose systems together with function syntax.
+///
+/// This macro provides some nice syntax on top of the `SystemRunner` `SystemParam`
+/// to allow running systems inside other systems. Overall, the macro accepts normal
+/// closure syntax:
+///
+/// ```ignore
+/// let system_a = |world: &mut World| { 10 };
+/// let system_b = |a: In<u32>, world: &mut World| { println!("{}", *a + 12) };
+/// compose! {
+///     || -> Result<(), RunSystemError> {
+///         let a = run!(system-a)?;
+///         run!(system_b);
+///     }
+/// }
+/// ```
+///
+/// What's special is that the macro will expand any invocations of `run!()` into
+/// calls to `SystemRunner::run` or `SystemRunner::run_with`. The `run!()` accepts
+/// two parameters: first, a system identifier (or a path to one), and second, an
+/// optional input to invoke the system with.
+///
+/// Notes:
+/// 1. All system runners are passed through a `ParamSet`, so invoked systems will
+///    not conflict with each other. However, invoked systems may still conflict
+///    with system params in the outer closure.
+///
+/// 2. `run!` will not accept expressions that evaluate to systems, only direct
+///    identifiers or paths. So, if you want to call something like:
+///
+///    ```ignore
+///    run!(|query: Query<(&A, &B, &mut C)>| { ... })`
+///    ```
+///
+///    Assign the expression to a variable first.
+#[proc_macro]
+pub fn compose(input: TokenStream) -> TokenStream {
+    system_composition::compose(input, false)
+}
+
+/// Automatically compose systems together with function syntax.
+///
+/// Unlike [`compose`], this macro allows generating systems that take input.
+///
+/// This macro provides some nice syntax on top of the `SystemRunner` `SystemParam`
+/// to allow running systems inside other systems. Overall, the macro accepts normal
+/// closure syntax:
+///
+/// ```ignore
+/// let system_a = |input: In<u32>, world: &mut World| { *input + 10 };
+/// let system_b = |a: In<u32>, world: &mut World| { println!("{}", *a + 12) };
+/// compose_with! {
+///     |input: In<u32>| -> Result<(), RunSystemError> {
+///         let a = run!(system_a, input)?;
+///         run!(system_b);
+///     }
+/// }
+/// ```
+///
+/// What's special is that the macro will expand any invocations of `run!()` into
+/// calls to `SystemRunner::run` or `SystemRunner::run_with`. The `run!()` accepts
+/// two parameters: first, a system identifier (or a path to one), and second, an
+/// optional input to invoke the system with.
+///
+/// Notes:
+/// 1. All system runners are passed through a `ParamSet`, so invoked systems will
+///    not conflict with each other. However, invoked systems may still conflict
+///    with system params in the outer closure.
+///
+/// 2. `run!` will not accept expressions that evaluate to systems, only direct
+///    identifiers or paths. So, if you want to call something like:
+///
+///    ```ignore
+///    run!(|query: Query<(&A, &B, &mut C)>| { ... })`
+///    ```
+///
+///    Assign the expression to a variable first.
+#[proc_macro]
+pub fn compose_with(input: TokenStream) -> TokenStream {
+    system_composition::compose(input, true)
+}

crates/bevy_ecs/macros/src/system_composition.rs

@@ -0,0 +1,128 @@
+use bevy_macro_utils::ensure_no_collision;
+use proc_macro::TokenStream;
+use quote::{format_ident, quote};
+use syn::{
+    parse::{Parse, ParseStream},
+    parse_macro_input, parse_quote, parse_quote_spanned,
+    spanned::Spanned,
+    visit_mut::VisitMut,
+    Expr, ExprMacro, Ident, Pat, Path, Token,
+};
+
+use crate::bevy_ecs_path;
+
+struct ExpandSystemCalls {
+    call_ident: Ident,
+    systems_ident: Ident,
+    system_paths: Vec<Path>,
+}
+
+struct SystemCall {
+    path: Path,
+    _comma: Option<Token![,]>,
+    input: Option<Box<Expr>>,
+}
+
+impl Parse for SystemCall {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        let path = input.parse()?;
+        let comma: Option<Token![,]> = input.parse()?;
+        let input = comma.as_ref().and_then(|_| input.parse().ok());
+        Ok(Self {
+            path,
+            _comma: comma,
+            input,
+        })
+    }
+}
+
+impl VisitMut for ExpandSystemCalls {
+    fn visit_expr_mut(&mut self, i: &mut Expr) {
+        if let Expr::Macro(ExprMacro { attrs: _, mac }) = i
+            && mac.path.is_ident(&self.call_ident)
+        {
+            let call = match mac.parse_body::<SystemCall>() {
+                Ok(call) => call,
+                Err(err) => {
+                    *i = Expr::Verbatim(err.into_compile_error());
+                    return;
+                }
+            };
+
+            let call_index = match self.system_paths.iter().position(|p| p == &call.path) {
+                Some(i) => i,
+                None => {
+                    let len = self.system_paths.len();
+                    self.system_paths.push(call.path.clone());
+                    len
+                }
+            };
+
+            let systems_ident = &self.systems_ident;
+            let system_accessor = format_ident!("p{}", call_index);
+            let expr: Expr = match &call.input {
+                Some(input) => {
+                    parse_quote_spanned!(mac.span()=> #systems_ident.#system_accessor().run_with(#input))
+                }
+                None => {
+                    parse_quote_spanned!(mac.span()=> #systems_ident.#system_accessor().run())
+                }
+            };
+            *i = expr;
+        } else {
+            syn::visit_mut::visit_expr_mut(self, i);
+        }
+    }
+}
+
+pub fn compose(input: TokenStream, has_input: bool) -> TokenStream {
+    let bevy_ecs_path = bevy_ecs_path();
+    let call_ident = format_ident!("run");
+    let systems_ident = ensure_no_collision(format_ident!("__systems"), input.clone());
+    let mut expr_closure = parse_macro_input!(input as syn::ExprClosure);
+
+    let mut visitor = ExpandSystemCalls {
+        call_ident,
+        systems_ident: systems_ident.clone(),
+        system_paths: Vec::new(),
+    };
+
+    syn::visit_mut::visit_expr_closure_mut(&mut visitor, &mut expr_closure);
+
+    let runner_types: Vec<syn::Type> = visitor
+        .system_paths
+        .iter()
+        .map(|path| parse_quote_spanned!(path.span()=> #bevy_ecs_path::system::SystemRunner<_, _, _>))
+        .collect();
+
+    let param_count = if has_input {
+        if expr_closure.inputs.is_empty() {
+            return TokenStream::from(
+                syn::Error::new_spanned(
+                    &expr_closure.inputs,
+                    "closure must have at least one parameter",
+                )
+                .into_compile_error(),
+            );
+        }
+        expr_closure.inputs.len() - 1
+    } else {
+        expr_closure.inputs.len()
+    };
+
+    let mut builders: Vec<Expr> =
+        vec![parse_quote!(#bevy_ecs_path::system::ParamBuilder); param_count];
+    let system_builders: Vec<Expr> = visitor.system_paths.iter().map(|path| parse_quote_spanned!(path.span()=> #bevy_ecs_path::system::ParamBuilder::system(#path))).collect();
+    builders.push(parse_quote!(#bevy_ecs_path::system::ParamSetBuilder((#(#system_builders,)*))));
+
+    expr_closure.inputs.push(Pat::Type(
+        parse_quote!(mut #systems_ident: #bevy_ecs_path::system::ParamSet<(#(#runner_types,)*)>),
+    ));
+
+    TokenStream::from(quote! {
+        #bevy_ecs_path::system::SystemParamBuilder::build_system(
+            (#(#builders,)*),
+            #expr_closure
+        )
+    })
+}