From 4be8e497ca2130561faa0a8e70644e8e36b7adbf Mon Sep 17 00:00:00 2001
From: Chris Russell <8494645+chescock@users.noreply.github.com>
Date: Wed, 28 Aug 2024 14:24:52 -0400
Subject: [PATCH] SystemParamBuilder - Allow deriving a SystemParamBuilder
 struct when deriving SystemParam. (#14818)

# Objective

Allow `SystemParamBuilder` implementations for custom system parameters
created using `#[derive(SystemParam)]`.

## Solution

Extend the derive macro to accept a `#[system_param(builder)]`
attribute. When present, emit a builder type with a field corresponding
to each field of the param.

## Example

```rust
#[derive(SystemParam)]
#[system_param(builder)]
struct CustomParam<'w, 's> {
    query: Query<'w, 's, ()>,
    local: Local<'s, usize>,
}

let system = (CustomParamBuilder {
    local: LocalBuilder(100),
    query: QueryParamBuilder::new(|builder| {
        builder.with::<A>();
    }),
},)
    .build_state(&mut world)
    .build_system(|param: CustomParam| *param.local + param.query.iter().count());
```
---
 crates/bevy_ecs/macros/src/lib.rs          | 54 ++++++++++++++++++++++
 crates/bevy_ecs/src/system/builder.rs      | 27 +++++++++++
 crates/bevy_ecs/src/system/system_param.rs | 49 ++++++++++++++++++++
 3 files changed, 130 insertions(+)

diff --git a/crates/bevy_ecs/macros/src/lib.rs b/crates/bevy_ecs/macros/src/lib.rs
index ed5cc22c24..f33659ca14 100644
--- a/crates/bevy_ecs/macros/src/lib.rs
+++ b/crates/bevy_ecs/macros/src/lib.rs
@@ -423,6 +423,56 @@ pub fn derive_system_param(input: TokenStream) -> TokenStream {
     let state_struct_visibility = &ast.vis;
     let state_struct_name = ensure_no_collision(format_ident!("FetchState"), token_stream);
 
+    let mut builder_name = None;
+    for meta in ast
+        .attrs
+        .iter()
+        .filter(|a| a.path().is_ident("system_param"))
+    {
+        if let Err(e) = meta.parse_nested_meta(|nested| {
+            if nested.path.is_ident("builder") {
+                builder_name = Some(format_ident!("{struct_name}Builder"));
+                Ok(())
+            } else {
+                Err(nested.error("Unsupported attribute"))
+            }
+        }) {
+            return e.into_compile_error().into();
+        }
+    }
+
+    let builder = builder_name.map(|builder_name| {
+        let builder_type_parameters: Vec<_> = (0..fields.len()).map(|i| format_ident!("B{i}")).collect();
+        let builder_doc_comment = format!("A [`SystemParamBuilder`] for a [`{struct_name}`].");
+        let builder_struct = quote! {
+            #[doc = #builder_doc_comment]
+            struct #builder_name<#(#builder_type_parameters,)*> {
+                #(#fields: #builder_type_parameters,)*
+            }
+        };
+        let lifetimes: Vec<_> = generics.lifetimes().collect();
+        let generic_struct = quote!{ #struct_name <#(#lifetimes,)* #punctuated_generic_idents> };
+        let builder_impl = quote!{
+            // SAFETY: This delegates to the `SystemParamBuilder` for tuples.
+            unsafe impl<
+                #(#lifetimes,)*
+                #(#builder_type_parameters: #path::system::SystemParamBuilder<#field_types>,)*
+                #punctuated_generics
+            > #path::system::SystemParamBuilder<#generic_struct> for #builder_name<#(#builder_type_parameters,)*>
+                #where_clause
+            {
+                fn build(self, world: &mut #path::world::World, meta: &mut #path::system::SystemMeta) -> <#generic_struct as #path::system::SystemParam>::State {
+                    let #builder_name { #(#fields: #field_locals,)* } = self;
+                    #state_struct_name {
+                        state: #path::system::SystemParamBuilder::build((#(#tuple_patterns,)*), world, meta)
+                    }
+                }
+            }
+        };
+        (builder_struct, builder_impl)
+    });
+    let (builder_struct, builder_impl) = builder.unzip();
+
     TokenStream::from(quote! {
         // We define the FetchState struct in an anonymous scope to avoid polluting the user namespace.
         // The struct can still be accessed via SystemParam::State, e.g. EventReaderState can be accessed via
@@ -479,7 +529,11 @@ pub fn derive_system_param(input: TokenStream) -> TokenStream {
 
             // Safety: Each field is `ReadOnlySystemParam`, so this can only read from the `World`
             unsafe impl<'w, 's, #punctuated_generics> #path::system::ReadOnlySystemParam for #struct_name #ty_generics #read_only_where_clause {}
+
+            #builder_impl
         };
+
+        #builder_struct
     })
 }
 
diff --git a/crates/bevy_ecs/src/system/builder.rs b/crates/bevy_ecs/src/system/builder.rs
index cd532776e5..ea5b070010 100644
--- a/crates/bevy_ecs/src/system/builder.rs
+++ b/crates/bevy_ecs/src/system/builder.rs
@@ -556,4 +556,31 @@ mod tests {
         let result = world.run_system_once(system);
         assert_eq!(result, 4);
     }
+
+    #[derive(SystemParam)]
+    #[system_param(builder)]
+    struct CustomParam<'w, 's> {
+        query: Query<'w, 's, ()>,
+        local: Local<'s, usize>,
+    }
+
+    #[test]
+    fn custom_param_builder() {
+        let mut world = World::new();
+
+        world.spawn(A);
+        world.spawn_empty();
+
+        let system = (CustomParamBuilder {
+            local: LocalBuilder(100),
+            query: QueryParamBuilder::new(|builder| {
+                builder.with::<A>();
+            }),
+        },)
+            .build_state(&mut world)
+            .build_system(|param: CustomParam| *param.local + param.query.iter().count());
+
+        let result = world.run_system_once(system);
+        assert_eq!(result, 101);
+    }
 }
diff --git a/crates/bevy_ecs/src/system/system_param.rs b/crates/bevy_ecs/src/system/system_param.rs
index 10e7fd4ec3..6b604c54c2 100644
--- a/crates/bevy_ecs/src/system/system_param.rs
+++ b/crates/bevy_ecs/src/system/system_param.rs
@@ -121,6 +121,55 @@ use std::{
 /// This will most commonly occur when working with `SystemParam`s generically, as the requirement
 /// has not been proven to the compiler.
 ///
+/// ## Builders
+///
+/// If you want to use a [`SystemParamBuilder`](crate::system::SystemParamBuilder) with a derived [`SystemParam`] implementation,
+/// add a `#[system_param(builder)]` attribute to the struct.
+/// This will generate a builder struct whose name is the param struct suffixed with `Builder`.
+/// The builder will not be `pub`, so you may want to expose a method that returns an `impl SystemParamBuilder<T>`.
+///
+/// ```
+/// mod custom_param {
+/// #     use bevy_ecs::{
+/// #         prelude::*,
+/// #         system::{LocalBuilder, QueryParamBuilder, SystemParam},
+/// #     };
+/// #
+///     #[derive(SystemParam)]
+///     #[system_param(builder)]
+///     pub struct CustomParam<'w, 's> {
+///         query: Query<'w, 's, ()>,
+///         local: Local<'s, usize>,
+///     }
+///
+///     impl<'w, 's> CustomParam<'w, 's> {
+///         pub fn builder(
+///             local: usize,
+///             query: impl FnOnce(&mut QueryBuilder<()>),
+///         ) -> impl SystemParamBuilder<Self> {
+///             CustomParamBuilder {
+///                 local: LocalBuilder(local),
+///                 query: QueryParamBuilder::new(query),
+///             }
+///         }
+///     }
+/// }
+///
+/// use custom_param::CustomParam;
+///
+/// # use bevy_ecs::prelude::*;
+/// # #[derive(Component)]
+/// # struct A;
+/// #
+/// # let mut world = World::new();
+/// #
+/// let system = (CustomParam::builder(100, |builder| {
+///     builder.with::<A>();
+/// }),)
+///     .build_state(&mut world)
+///     .build_system(|param: CustomParam| {});
+/// ```
+///
 /// # Safety
 ///
 /// The implementor must ensure the following is true.