Skip to content

Commit

Permalink
SystemParamBuilder - Allow deriving a SystemParamBuilder struct when …
Browse files Browse the repository at this point in the history
…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());
```
  • Loading branch information
chescock authored Aug 28, 2024
1 parent 4648f7b commit 4be8e49
Show file tree
Hide file tree
Showing 3 changed files with 130 additions and 0 deletions.
54 changes: 54 additions & 0 deletions crates/bevy_ecs/macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
})
}

Expand Down
27 changes: 27 additions & 0 deletions crates/bevy_ecs/src/system/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
49 changes: 49 additions & 0 deletions crates/bevy_ecs/src/system/system_param.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down

0 comments on commit 4be8e49

Please sign in to comment.