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
10pub 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 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 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 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 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 pub fn resource<T: Resource>(self) -> SystemBuilder<'w, ($($curr,)* Res<'static, T>,)> {
105 self.param::<Res<T>>()
106 }
107
108 pub fn resource_mut<T: Resource>(self) -> SystemBuilder<'w, ($($curr,)* ResMut<'static, T>,)> {
110 self.param::<ResMut<T>>()
111 }
112
113 pub fn local<T: Send + FromWorld>(self) -> SystemBuilder<'w, ($($curr,)* Local<'static, T>,)> {
115 self.param::<Local<T>>()
116 }
117
118 pub fn query<D: QueryData>(self) -> SystemBuilder<'w, ($($curr,)* Query<'static, 'static, D, ()>,)> {
120 self.query_filtered::<D, ()>()
121 }
122
123 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 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}