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());
```
This commit is contained in:
Chris Russell 2024-08-28 14:24:52 -04:00 committed by GitHub
parent 4648f7bf72
commit 4be8e497ca
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 130 additions and 0 deletions

View file

@ -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
})
}

View file

@ -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);
}
}

View file

@ -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.