bevy_ecs/system/
builder.rs

1use bevy_utils::all_tuples;
2
3use super::{
4    BuildableSystemParam, FunctionSystem, Local, Res, ResMut, Resource, SystemMeta, SystemParam,
5    SystemParamFunction, SystemState,
6};
7use crate::prelude::{FromWorld, Query, World};
8use crate::query::{QueryData, QueryFilter};
9
10/// Builder struct used to construct state for [`SystemParam`] passed to a system.
11///
12/// ```
13/// # use bevy_ecs::prelude::*;
14/// # use bevy_ecs_macros::SystemParam;
15/// # use bevy_ecs::system::RunSystemOnce;
16/// #
17/// # #[derive(Component)]
18/// # struct A;
19/// #
20/// # #[derive(Component)]
21/// # struct B;
22/// #
23/// # #[derive(Resource)]
24/// # struct R;
25/// #
26/// # #[derive(SystemParam)]
27/// # struct MyParam;
28/// #
29/// # let mut world = World::new();
30/// # world.insert_resource(R);
31///
32/// fn my_system(res: Res<R>, query: Query<&A>, param: MyParam) {
33///     // ...
34/// }
35///
36/// // Create a builder from the world, helper methods exist to add `SystemParam`,
37/// // alternatively use `.param::<T>()` for any other `SystemParam` types.
38/// let system = SystemBuilder::<()>::new(&mut world)
39///     .resource::<R>()
40///     .query::<&A>()
41///     .param::<MyParam>()
42///     .build(my_system);
43///
44/// // Parameters that the builder is initialised with will appear first in the arguments.
45/// let system = SystemBuilder::<(Res<R>, Query<&A>)>::new(&mut world)
46///     .param::<MyParam>()
47///     .build(my_system);
48///
49/// // Parameters that implement `BuildableSystemParam` can use `.builder::<T>()` to build in place.
50/// let system = SystemBuilder::<()>::new(&mut world)
51///     .resource::<R>()
52///     .builder::<Query<&A>>(|builder| { builder.with::<B>(); })
53///     .param::<MyParam>()
54///     .build(my_system);
55///
56/// world.run_system_once(system);
57///```
58pub struct SystemBuilder<'w, T: SystemParam = ()> {
59    pub(crate) meta: SystemMeta,
60    pub(crate) state: T::State,
61    pub(crate) world: &'w mut World,
62}
63
64impl<'w, T: SystemParam> SystemBuilder<'w, T> {
65    /// Construct a new builder with the default state for `T`
66    pub fn new(world: &'w mut World) -> Self {
67        let mut meta = SystemMeta::new::<T>();
68        Self {
69            state: T::init_state(world, &mut meta),
70            meta,
71            world,
72        }
73    }
74
75    /// Construct the a system with the built params
76    pub fn build<F, Marker>(self, func: F) -> FunctionSystem<Marker, F>
77    where
78        F: SystemParamFunction<Marker, Param = T>,
79    {
80        FunctionSystem::from_builder(self, func)
81    }
82
83    /// Return the constructed [`SystemState`]
84    pub fn state(self) -> SystemState<T> {
85        SystemState::from_builder(self)
86    }
87}
88
89macro_rules! impl_system_builder {
90    ($($curr: ident),*) => {
91        impl<'w, $($curr: SystemParam,)*> SystemBuilder<'w, ($($curr,)*)> {
92            /// Add `T` as a parameter built from the world
93            pub fn param<T: SystemParam>(mut self) -> SystemBuilder<'w, ($($curr,)* T,)> {
94                #[allow(non_snake_case)]
95                let ($($curr,)*) = self.state;
96                SystemBuilder {
97                    state: ($($curr,)* T::init_state(self.world, &mut self.meta),),
98                    meta: self.meta,
99                    world: self.world,
100                }
101            }
102
103            /// Helper method for reading a [`Resource`] as a param, equivalent to `.param::<Res<T>>()`
104            pub fn resource<T: Resource>(self) -> SystemBuilder<'w,  ($($curr,)* Res<'static, T>,)> {
105                self.param::<Res<T>>()
106            }
107
108            /// Helper method for mutably accessing a [`Resource`] as a param, equivalent to `.param::<ResMut<T>>()`
109            pub fn resource_mut<T: Resource>(self) -> SystemBuilder<'w,  ($($curr,)* ResMut<'static, T>,)> {
110                self.param::<ResMut<T>>()
111            }
112
113            /// Helper method for adding a [`Local`] as a param, equivalent to `.param::<Local<T>>()`
114            pub fn local<T: Send + FromWorld>(self) -> SystemBuilder<'w,  ($($curr,)* Local<'static, T>,)> {
115                self.param::<Local<T>>()
116            }
117
118            /// Helper method for adding a [`Query`] as a param, equivalent to `.param::<Query<D>>()`
119            pub fn query<D: QueryData>(self) -> SystemBuilder<'w,  ($($curr,)* Query<'static, 'static, D, ()>,)> {
120                self.query_filtered::<D, ()>()
121            }
122
123            /// Helper method for adding a filtered [`Query`] as a param, equivalent to `.param::<Query<D, F>>()`
124            pub fn query_filtered<D: QueryData, F: QueryFilter>(self) -> SystemBuilder<'w,  ($($curr,)* Query<'static, 'static, D, F>,)> {
125                self.param::<Query<D, F>>()
126            }
127
128            /// Add `T` as a parameter built with the given function
129            pub fn builder<T: BuildableSystemParam>(
130                mut self,
131                func: impl FnOnce(&mut T::Builder<'_>),
132            ) -> SystemBuilder<'w, ($($curr,)* T,)> {
133                #[allow(non_snake_case)]
134                let ($($curr,)*) = self.state;
135                SystemBuilder {
136                    state: ($($curr,)* T::build(self.world, &mut self.meta, func),),
137                    meta: self.meta,
138                    world: self.world,
139                }
140            }
141        }
142    };
143}
144
145all_tuples!(impl_system_builder, 0, 15, P);
146
147#[cfg(test)]
148mod tests {
149    use crate as bevy_ecs;
150    use crate::prelude::{Component, Query};
151    use crate::system::{Local, RunSystemOnce};
152
153    use super::*;
154
155    #[derive(Component)]
156    struct A;
157
158    fn local_system(local: Local<u64>) -> u64 {
159        *local
160    }
161
162    fn query_system(query: Query<()>) -> usize {
163        query.iter().count()
164    }
165
166    fn multi_param_system(a: Local<u64>, b: Local<u64>) -> u64 {
167        *a + *b + 1
168    }
169
170    #[test]
171    fn local_builder() {
172        let mut world = World::new();
173
174        let system = SystemBuilder::<()>::new(&mut world)
175            .builder::<Local<u64>>(|x| *x = 10)
176            .build(local_system);
177
178        let result = world.run_system_once(system);
179        assert_eq!(result, 10);
180    }
181
182    #[test]
183    fn query_builder() {
184        let mut world = World::new();
185
186        world.spawn(A);
187        world.spawn_empty();
188
189        let system = SystemBuilder::<()>::new(&mut world)
190            .builder::<Query<()>>(|query| {
191                query.with::<A>();
192            })
193            .build(query_system);
194
195        let result = world.run_system_once(system);
196        assert_eq!(result, 1);
197    }
198
199    #[test]
200    fn multi_param_builder() {
201        let mut world = World::new();
202
203        world.spawn(A);
204        world.spawn_empty();
205
206        let system = SystemBuilder::<()>::new(&mut world)
207            .local::<u64>()
208            .param::<Local<u64>>()
209            .build(multi_param_system);
210
211        let result = world.run_system_once(system);
212        assert_eq!(result, 1);
213    }
214}