bevy_ecs/query/
builder.rs

1use std::marker::PhantomData;
2
3use crate::{component::ComponentId, prelude::*};
4
5use super::{FilteredAccess, QueryData, QueryFilter};
6
7/// Builder struct to create [`QueryState`] instances at runtime.
8///
9/// ```
10/// # use bevy_ecs::prelude::*;
11/// #
12/// # #[derive(Component)]
13/// # struct A;
14/// #
15/// # #[derive(Component)]
16/// # struct B;
17/// #
18/// # #[derive(Component)]
19/// # struct C;
20/// #
21/// let mut world = World::new();
22/// let entity_a = world.spawn((A, B)).id();
23/// let entity_b = world.spawn((A, C)).id();
24///
25/// // Instantiate the builder using the type signature of the iterator you will consume
26/// let mut query = QueryBuilder::<(Entity, &B)>::new(&mut world)
27/// // Add additional terms through builder methods
28///     .with::<A>()
29///     .without::<C>()
30///     .build();
31///
32/// // Consume the QueryState
33/// let (entity, b) = query.single(&world);
34///```
35pub struct QueryBuilder<'w, D: QueryData = (), F: QueryFilter = ()> {
36    access: FilteredAccess<ComponentId>,
37    world: &'w mut World,
38    or: bool,
39    first: bool,
40    _marker: PhantomData<(D, F)>,
41}
42
43impl<'w, D: QueryData, F: QueryFilter> QueryBuilder<'w, D, F> {
44    /// Creates a new builder with the accesses required for `Q` and `F`
45    pub fn new(world: &'w mut World) -> Self {
46        let fetch_state = D::init_state(world);
47        let filter_state = F::init_state(world);
48
49        let mut access = FilteredAccess::default();
50        D::update_component_access(&fetch_state, &mut access);
51
52        // Use a temporary empty FilteredAccess for filters. This prevents them from conflicting with the
53        // main Query's `fetch_state` access. Filters are allowed to conflict with the main query fetch
54        // because they are evaluated *before* a specific reference is constructed.
55        let mut filter_access = FilteredAccess::default();
56        F::update_component_access(&filter_state, &mut filter_access);
57
58        // Merge the temporary filter access with the main access. This ensures that filter access is
59        // properly considered in a global "cross-query" context (both within systems and across systems).
60        access.extend(&filter_access);
61
62        Self {
63            access,
64            world,
65            or: false,
66            first: false,
67            _marker: PhantomData,
68        }
69    }
70
71    /// Returns a reference to the world passed to [`Self::new`].
72    pub fn world(&self) -> &World {
73        self.world
74    }
75
76    /// Returns a mutable reference to the world passed to [`Self::new`].
77    pub fn world_mut(&mut self) -> &mut World {
78        self.world
79    }
80
81    /// Adds access to self's underlying [`FilteredAccess`] respecting [`Self::or`] and [`Self::and`]
82    pub fn extend_access(&mut self, mut access: FilteredAccess<ComponentId>) {
83        if self.or {
84            if self.first {
85                access.required.clear();
86                self.access.extend(&access);
87                self.first = false;
88            } else {
89                self.access.append_or(&access);
90            }
91        } else {
92            self.access.extend(&access);
93        }
94    }
95
96    /// Adds accesses required for `T` to self.
97    pub fn data<T: QueryData>(&mut self) -> &mut Self {
98        let state = T::init_state(self.world);
99        let mut access = FilteredAccess::default();
100        T::update_component_access(&state, &mut access);
101        self.extend_access(access);
102        self
103    }
104
105    /// Adds filter from `T` to self.
106    pub fn filter<T: QueryFilter>(&mut self) -> &mut Self {
107        let state = T::init_state(self.world);
108        let mut access = FilteredAccess::default();
109        T::update_component_access(&state, &mut access);
110        self.extend_access(access);
111        self
112    }
113
114    /// Adds [`With<T>`] to the [`FilteredAccess`] of self.
115    pub fn with<T: Component>(&mut self) -> &mut Self {
116        self.filter::<With<T>>();
117        self
118    }
119
120    /// Adds [`With<T>`] to the [`FilteredAccess`] of self from a runtime [`ComponentId`].
121    pub fn with_id(&mut self, id: ComponentId) -> &mut Self {
122        let mut access = FilteredAccess::default();
123        access.and_with(id);
124        self.extend_access(access);
125        self
126    }
127
128    /// Adds [`Without<T>`] to the [`FilteredAccess`] of self.
129    pub fn without<T: Component>(&mut self) -> &mut Self {
130        self.filter::<Without<T>>();
131        self
132    }
133
134    /// Adds [`Without<T>`] to the [`FilteredAccess`] of self from a runtime [`ComponentId`].
135    pub fn without_id(&mut self, id: ComponentId) -> &mut Self {
136        let mut access = FilteredAccess::default();
137        access.and_without(id);
138        self.extend_access(access);
139        self
140    }
141
142    /// Adds `&T` to the [`FilteredAccess`] of self.
143    pub fn ref_id(&mut self, id: ComponentId) -> &mut Self {
144        self.with_id(id);
145        self.access.add_read(id);
146        self
147    }
148
149    /// Adds `&mut T` to the [`FilteredAccess`] of self.
150    pub fn mut_id(&mut self, id: ComponentId) -> &mut Self {
151        self.with_id(id);
152        self.access.add_write(id);
153        self
154    }
155
156    /// Takes a function over mutable access to a [`QueryBuilder`], calls that function
157    /// on an empty builder and then adds all accesses from that builder to self as optional.
158    pub fn optional(&mut self, f: impl Fn(&mut QueryBuilder)) -> &mut Self {
159        let mut builder = QueryBuilder::new(self.world);
160        f(&mut builder);
161        self.access.extend_access(builder.access());
162        self
163    }
164
165    /// Takes a function over mutable access to a [`QueryBuilder`], calls that function
166    /// on an empty builder and then adds all accesses from that builder to self.
167    ///
168    /// Primarily used when inside a [`Self::or`] closure to group several terms.
169    pub fn and(&mut self, f: impl Fn(&mut QueryBuilder)) -> &mut Self {
170        let mut builder = QueryBuilder::new(self.world);
171        f(&mut builder);
172        let access = builder.access().clone();
173        self.extend_access(access);
174        self
175    }
176
177    /// Takes a function over mutable access to a [`QueryBuilder`], calls that function
178    /// on an empty builder, all accesses added to that builder will become terms in an or expression.
179    ///
180    /// ```
181    /// # use bevy_ecs::prelude::*;
182    /// #
183    /// # #[derive(Component)]
184    /// # struct A;
185    /// #
186    /// # #[derive(Component)]
187    /// # struct B;
188    /// #
189    /// # let mut world = World::new();
190    /// #
191    /// QueryBuilder::<Entity>::new(&mut world).or(|builder| {
192    ///     builder.with::<A>();
193    ///     builder.with::<B>();
194    /// });
195    /// // is equivalent to
196    /// QueryBuilder::<Entity>::new(&mut world).filter::<Or<(With<A>, With<B>)>>();
197    /// ```
198    pub fn or(&mut self, f: impl Fn(&mut QueryBuilder)) -> &mut Self {
199        let mut builder = QueryBuilder::new(self.world);
200        builder.or = true;
201        builder.first = true;
202        f(&mut builder);
203        self.access.extend(builder.access());
204        self
205    }
206
207    /// Returns a reference to the [`FilteredAccess`] that will be provided to the built [`Query`].
208    pub fn access(&self) -> &FilteredAccess<ComponentId> {
209        &self.access
210    }
211
212    /// Transmute the existing builder adding required accesses.
213    /// This will maintain all existing accesses.
214    ///
215    /// If including a filter type see [`Self::transmute_filtered`]
216    pub fn transmute<NewD: QueryData>(&mut self) -> &mut QueryBuilder<'w, NewD> {
217        self.transmute_filtered::<NewD, ()>()
218    }
219
220    /// Transmute the existing builder adding required accesses.
221    /// This will maintain all existing accesses.
222    pub fn transmute_filtered<NewD: QueryData, NewF: QueryFilter>(
223        &mut self,
224    ) -> &mut QueryBuilder<'w, NewD, NewF> {
225        let mut fetch_state = NewD::init_state(self.world);
226        let filter_state = NewF::init_state(self.world);
227
228        NewD::set_access(&mut fetch_state, &self.access);
229
230        let mut access = FilteredAccess::default();
231        NewD::update_component_access(&fetch_state, &mut access);
232        NewF::update_component_access(&filter_state, &mut access);
233
234        self.extend_access(access);
235        // SAFETY:
236        // - We have included all required accesses for NewQ and NewF
237        // - The layout of all QueryBuilder instances is the same
238        unsafe { std::mem::transmute(self) }
239    }
240
241    /// Create a [`QueryState`] with the accesses of the builder.
242    ///
243    /// Takes `&mut self` to access the innner world reference while initializing
244    /// state for the new [`QueryState`]
245    pub fn build(&mut self) -> QueryState<D, F> {
246        QueryState::<D, F>::from_builder(self)
247    }
248}
249
250#[cfg(test)]
251mod tests {
252    use crate as bevy_ecs;
253    use crate::prelude::*;
254    use crate::world::FilteredEntityRef;
255
256    #[derive(Component, PartialEq, Debug)]
257    struct A(usize);
258
259    #[derive(Component, PartialEq, Debug)]
260    struct B(usize);
261
262    #[derive(Component, PartialEq, Debug)]
263    struct C(usize);
264
265    #[test]
266    fn builder_with_without_static() {
267        let mut world = World::new();
268        let entity_a = world.spawn((A(0), B(0))).id();
269        let entity_b = world.spawn((A(0), C(0))).id();
270
271        let mut query_a = QueryBuilder::<Entity>::new(&mut world)
272            .with::<A>()
273            .without::<C>()
274            .build();
275        assert_eq!(entity_a, query_a.single(&world));
276
277        let mut query_b = QueryBuilder::<Entity>::new(&mut world)
278            .with::<A>()
279            .without::<B>()
280            .build();
281        assert_eq!(entity_b, query_b.single(&world));
282    }
283
284    #[test]
285    fn builder_with_without_dynamic() {
286        let mut world = World::new();
287        let entity_a = world.spawn((A(0), B(0))).id();
288        let entity_b = world.spawn((A(0), C(0))).id();
289        let component_id_a = world.init_component::<A>();
290        let component_id_b = world.init_component::<B>();
291        let component_id_c = world.init_component::<C>();
292
293        let mut query_a = QueryBuilder::<Entity>::new(&mut world)
294            .with_id(component_id_a)
295            .without_id(component_id_c)
296            .build();
297        assert_eq!(entity_a, query_a.single(&world));
298
299        let mut query_b = QueryBuilder::<Entity>::new(&mut world)
300            .with_id(component_id_a)
301            .without_id(component_id_b)
302            .build();
303        assert_eq!(entity_b, query_b.single(&world));
304    }
305
306    #[test]
307    fn builder_or() {
308        let mut world = World::new();
309        world.spawn((A(0), B(0)));
310        world.spawn(B(0));
311        world.spawn(C(0));
312
313        let mut query_a = QueryBuilder::<Entity>::new(&mut world)
314            .or(|builder| {
315                builder.with::<A>();
316                builder.with::<B>();
317            })
318            .build();
319        assert_eq!(2, query_a.iter(&world).count());
320
321        let mut query_b = QueryBuilder::<Entity>::new(&mut world)
322            .or(|builder| {
323                builder.with::<A>();
324                builder.without::<B>();
325            })
326            .build();
327        dbg!(&query_b.component_access);
328        assert_eq!(2, query_b.iter(&world).count());
329
330        let mut query_c = QueryBuilder::<Entity>::new(&mut world)
331            .or(|builder| {
332                builder.with::<A>();
333                builder.with::<B>();
334                builder.with::<C>();
335            })
336            .build();
337        assert_eq!(3, query_c.iter(&world).count());
338    }
339
340    #[test]
341    fn builder_transmute() {
342        let mut world = World::new();
343        world.spawn(A(0));
344        world.spawn((A(1), B(0)));
345        let mut query = QueryBuilder::<()>::new(&mut world)
346            .with::<B>()
347            .transmute::<&A>()
348            .build();
349
350        query.iter(&world).for_each(|a| assert_eq!(a.0, 1));
351    }
352
353    #[test]
354    fn builder_static_components() {
355        let mut world = World::new();
356        let entity = world.spawn((A(0), B(1))).id();
357
358        let mut query = QueryBuilder::<FilteredEntityRef>::new(&mut world)
359            .data::<&A>()
360            .data::<&B>()
361            .build();
362
363        let entity_ref = query.single(&world);
364
365        assert_eq!(entity, entity_ref.id());
366
367        let a = entity_ref.get::<A>().unwrap();
368        let b = entity_ref.get::<B>().unwrap();
369
370        assert_eq!(0, a.0);
371        assert_eq!(1, b.0);
372    }
373
374    #[test]
375    fn builder_dynamic_components() {
376        let mut world = World::new();
377        let entity = world.spawn((A(0), B(1))).id();
378        let component_id_a = world.init_component::<A>();
379        let component_id_b = world.init_component::<B>();
380
381        let mut query = QueryBuilder::<FilteredEntityRef>::new(&mut world)
382            .ref_id(component_id_a)
383            .ref_id(component_id_b)
384            .build();
385
386        let entity_ref = query.single(&world);
387
388        assert_eq!(entity, entity_ref.id());
389
390        let a = entity_ref.get_by_id(component_id_a).unwrap();
391        let b = entity_ref.get_by_id(component_id_b).unwrap();
392
393        // SAFETY: We set these pointers to point to these components
394        unsafe {
395            assert_eq!(0, a.deref::<A>().0);
396            assert_eq!(1, b.deref::<B>().0);
397        }
398    }
399}
bevy_ecs/query/builder.rs

bevy_ecs/query/
builder.rs
