bevy_ecs/query/
builder.rs

1use core::marker::PhantomData;
2
3use crate::{
4    component::{ComponentId, StorageType},
5    prelude::*,
6};
7
8use super::{FilteredAccess, QueryData, QueryFilter};
9
10/// Builder struct to create [`QueryState`] instances at runtime.
11///
12/// ```
13/// # use bevy_ecs::prelude::*;
14/// #
15/// # #[derive(Component)]
16/// # struct A;
17/// #
18/// # #[derive(Component)]
19/// # struct B;
20/// #
21/// # #[derive(Component)]
22/// # struct C;
23/// #
24/// let mut world = World::new();
25/// let entity_a = world.spawn((A, B)).id();
26/// let entity_b = world.spawn((A, C)).id();
27///
28/// // Instantiate the builder using the type signature of the iterator you will consume
29/// let mut query = QueryBuilder::<(Entity, &B)>::new(&mut world)
30/// // Add additional terms through builder methods
31///     .with::<A>()
32///     .without::<C>()
33///     .build();
34///
35/// // Consume the QueryState
36/// let (entity, b) = query.single(&world).unwrap();
37/// ```
38pub struct QueryBuilder<'w, D: QueryData = (), F: QueryFilter = ()> {
39    access: FilteredAccess,
40    world: &'w mut World,
41    or: bool,
42    first: bool,
43    _marker: PhantomData<(D, F)>,
44}
45
46impl<'w, D: QueryData, F: QueryFilter> QueryBuilder<'w, D, F> {
47    /// Creates a new builder with the accesses required for `Q` and `F`
48    pub fn new(world: &'w mut World) -> Self {
49        let fetch_state = D::init_state(world);
50        let filter_state = F::init_state(world);
51
52        let mut access = FilteredAccess::default();
53        D::update_component_access(&fetch_state, &mut access);
54
55        // Use a temporary empty FilteredAccess for filters. This prevents them from conflicting with the
56        // main Query's `fetch_state` access. Filters are allowed to conflict with the main query fetch
57        // because they are evaluated *before* a specific reference is constructed.
58        let mut filter_access = FilteredAccess::default();
59        F::update_component_access(&filter_state, &mut filter_access);
60
61        // Merge the temporary filter access with the main access. This ensures that filter access is
62        // properly considered in a global "cross-query" context (both within systems and across systems).
63        access.extend(&filter_access);
64
65        Self {
66            access,
67            world,
68            or: false,
69            first: false,
70            _marker: PhantomData,
71        }
72    }
73
74    pub(super) fn is_dense(&self) -> bool {
75        // Note: `component_id` comes from the user in safe code, so we cannot trust it to
76        // exist. If it doesn't exist we pessimistically assume it's sparse.
77        let is_dense = |component_id| {
78            self.world()
79                .components()
80                .get_info(component_id)
81                .is_some_and(|info| info.storage_type() == StorageType::Table)
82        };
83
84        // Use dense iteration if possible, but fall back to sparse if we need to.
85        // Both `D` and `F` must allow dense iteration, just as for queries without dynamic filters.
86        // All `with` and `without` filters must be dense to ensure that we match all archetypes in a table.
87        // We also need to ensure that any sparse set components in `access.required` cause sparse iteration,
88        // but anything that adds a `required` component also adds a `with` filter.
89        //
90        // Note that `EntityRef` and `EntityMut` types, including `FilteredEntityRef` and `FilteredEntityMut`, have `D::IS_DENSE = true`.
91        // Calling `builder.data::<&Sparse>()` will add a filter and force sparse iteration,
92        // but calling `builder.data::<Option<&Sparse>>()` will still allow them to use dense iteration!
93        D::IS_DENSE
94            && F::IS_DENSE
95            && self.access.with_filters().all(is_dense)
96            && self.access.without_filters().all(is_dense)
97    }
98
99    /// Returns a reference to the world passed to [`Self::new`].
100    pub fn world(&self) -> &World {
101        self.world
102    }
103
104    /// Returns a mutable reference to the world passed to [`Self::new`].
105    pub fn world_mut(&mut self) -> &mut World {
106        self.world
107    }
108
109    /// Adds access to self's underlying [`FilteredAccess`] respecting [`Self::or`] and [`Self::and`]
110    pub fn extend_access(&mut self, mut access: FilteredAccess) {
111        if self.or {
112            if self.first {
113                access.required.clear();
114                self.access.extend(&access);
115                self.first = false;
116            } else {
117                self.access.append_or(&access);
118            }
119        } else {
120            self.access.extend(&access);
121        }
122    }
123
124    /// Adds accesses required for `T` to self.
125    pub fn data<T: QueryData>(&mut self) -> &mut Self {
126        let state = T::init_state(self.world);
127        let mut access = FilteredAccess::default();
128        T::update_component_access(&state, &mut access);
129        self.extend_access(access);
130        self
131    }
132
133    /// Adds filter from `T` to self.
134    pub fn filter<T: QueryFilter>(&mut self) -> &mut Self {
135        let state = T::init_state(self.world);
136        let mut access = FilteredAccess::default();
137        T::update_component_access(&state, &mut access);
138        self.extend_access(access);
139        self
140    }
141
142    /// Adds [`With<T>`] to the [`FilteredAccess`] of self.
143    pub fn with<T: Component>(&mut self) -> &mut Self {
144        self.filter::<With<T>>();
145        self
146    }
147
148    /// Adds [`With<T>`] to the [`FilteredAccess`] of self from a runtime [`ComponentId`].
149    pub fn with_id(&mut self, id: ComponentId) -> &mut Self {
150        let mut access = FilteredAccess::default();
151        access.and_with(id);
152        self.extend_access(access);
153        self
154    }
155
156    /// Adds [`Without<T>`] to the [`FilteredAccess`] of self.
157    pub fn without<T: Component>(&mut self) -> &mut Self {
158        self.filter::<Without<T>>();
159        self
160    }
161
162    /// Adds [`Without<T>`] to the [`FilteredAccess`] of self from a runtime [`ComponentId`].
163    pub fn without_id(&mut self, id: ComponentId) -> &mut Self {
164        let mut access = FilteredAccess::default();
165        access.and_without(id);
166        self.extend_access(access);
167        self
168    }
169
170    /// Adds `&T` to the [`FilteredAccess`] of self.
171    pub fn ref_id(&mut self, id: ComponentId) -> &mut Self {
172        self.with_id(id);
173        self.access.add_component_read(id);
174        self
175    }
176
177    /// Adds `&mut T` to the [`FilteredAccess`] of self.
178    pub fn mut_id(&mut self, id: ComponentId) -> &mut Self {
179        self.with_id(id);
180        self.access.add_component_write(id);
181        self
182    }
183
184    /// Takes a function over mutable access to a [`QueryBuilder`], calls that function
185    /// on an empty builder and then adds all accesses from that builder to self as optional.
186    pub fn optional(&mut self, f: impl Fn(&mut QueryBuilder)) -> &mut Self {
187        let mut builder = QueryBuilder::new(self.world);
188        f(&mut builder);
189        self.access.extend_access(builder.access());
190        self
191    }
192
193    /// Takes a function over mutable access to a [`QueryBuilder`], calls that function
194    /// on an empty builder and then adds all accesses from that builder to self.
195    ///
196    /// Primarily used when inside a [`Self::or`] closure to group several terms.
197    pub fn and(&mut self, f: impl Fn(&mut QueryBuilder)) -> &mut Self {
198        let mut builder = QueryBuilder::new(self.world);
199        f(&mut builder);
200        let access = builder.access().clone();
201        self.extend_access(access);
202        self
203    }
204
205    /// Takes a function over mutable access to a [`QueryBuilder`], calls that function
206    /// on an empty builder, all accesses added to that builder will become terms in an or expression.
207    ///
208    /// ```
209    /// # use bevy_ecs::prelude::*;
210    /// #
211    /// # #[derive(Component)]
212    /// # struct A;
213    /// #
214    /// # #[derive(Component)]
215    /// # struct B;
216    /// #
217    /// # let mut world = World::new();
218    /// #
219    /// QueryBuilder::<Entity>::new(&mut world).or(|builder| {
220    ///     builder.with::<A>();
221    ///     builder.with::<B>();
222    /// });
223    /// // is equivalent to
224    /// QueryBuilder::<Entity>::new(&mut world).filter::<Or<(With<A>, With<B>)>>();
225    /// ```
226    pub fn or(&mut self, f: impl Fn(&mut QueryBuilder)) -> &mut Self {
227        let mut builder = QueryBuilder::new(self.world);
228        builder.or = true;
229        builder.first = true;
230        f(&mut builder);
231        self.access.extend(builder.access());
232        self
233    }
234
235    /// Returns a reference to the [`FilteredAccess`] that will be provided to the built [`Query`].
236    pub fn access(&self) -> &FilteredAccess {
237        &self.access
238    }
239
240    /// Transmute the existing builder adding required accesses.
241    /// This will maintain all existing accesses.
242    ///
243    /// If including a filter type see [`Self::transmute_filtered`]
244    pub fn transmute<NewD: QueryData>(&mut self) -> &mut QueryBuilder<'w, NewD> {
245        self.transmute_filtered::<NewD, ()>()
246    }
247
248    /// Transmute the existing builder adding required accesses.
249    /// This will maintain all existing accesses.
250    pub fn transmute_filtered<NewD: QueryData, NewF: QueryFilter>(
251        &mut self,
252    ) -> &mut QueryBuilder<'w, NewD, NewF> {
253        let fetch_state = NewD::init_state(self.world);
254        let filter_state = NewF::init_state(self.world);
255
256        let mut access = FilteredAccess::default();
257        NewD::update_component_access(&fetch_state, &mut access);
258        NewF::update_component_access(&filter_state, &mut access);
259
260        self.extend_access(access);
261        // SAFETY:
262        // - We have included all required accesses for NewQ and NewF
263        // - The layout of all QueryBuilder instances is the same
264        unsafe { core::mem::transmute(self) }
265    }
266
267    /// Create a [`QueryState`] with the accesses of the builder.
268    ///
269    /// Takes `&mut self` to access the inner world reference while initializing
270    /// state for the new [`QueryState`]
271    pub fn build(&mut self) -> QueryState<D, F> {
272        QueryState::<D, F>::from_builder(self)
273    }
274}
275
276#[cfg(test)]
277mod tests {
278    use crate::{
279        prelude::*,
280        world::{EntityMutExcept, EntityRefExcept, FilteredEntityMut, FilteredEntityRef},
281    };
282    use std::dbg;
283
284    #[derive(Component, PartialEq, Debug)]
285    struct A(usize);
286
287    #[derive(Component, PartialEq, Debug)]
288    struct B(usize);
289
290    #[derive(Component, PartialEq, Debug)]
291    struct C(usize);
292
293    #[derive(Component)]
294    struct D;
295
296    #[test]
297    fn builder_with_without_static() {
298        let mut world = World::new();
299        let entity_a = world.spawn((A(0), B(0))).id();
300        let entity_b = world.spawn((A(0), C(0))).id();
301
302        let mut query_a = QueryBuilder::<Entity>::new(&mut world)
303            .with::<A>()
304            .without::<C>()
305            .build();
306        assert_eq!(entity_a, query_a.single(&world).unwrap());
307
308        let mut query_b = QueryBuilder::<Entity>::new(&mut world)
309            .with::<A>()
310            .without::<B>()
311            .build();
312        assert_eq!(entity_b, query_b.single(&world).unwrap());
313    }
314
315    #[test]
316    fn builder_with_without_dynamic() {
317        let mut world = World::new();
318        let entity_a = world.spawn((A(0), B(0))).id();
319        let entity_b = world.spawn((A(0), C(0))).id();
320        let component_id_a = world.register_component::<A>();
321        let component_id_b = world.register_component::<B>();
322        let component_id_c = world.register_component::<C>();
323
324        let mut query_a = QueryBuilder::<Entity>::new(&mut world)
325            .with_id(component_id_a)
326            .without_id(component_id_c)
327            .build();
328        assert_eq!(entity_a, query_a.single(&world).unwrap());
329
330        let mut query_b = QueryBuilder::<Entity>::new(&mut world)
331            .with_id(component_id_a)
332            .without_id(component_id_b)
333            .build();
334        assert_eq!(entity_b, query_b.single(&world).unwrap());
335    }
336
337    #[test]
338    fn builder_or() {
339        let mut world = World::new();
340        world.spawn((A(0), B(0), D));
341        world.spawn((B(0), D));
342        world.spawn((C(0), D));
343
344        let mut query_a = QueryBuilder::<&D>::new(&mut world)
345            .or(|builder| {
346                builder.with::<A>();
347                builder.with::<B>();
348            })
349            .build();
350        assert_eq!(2, query_a.iter(&world).count());
351
352        let mut query_b = QueryBuilder::<&D>::new(&mut world)
353            .or(|builder| {
354                builder.with::<A>();
355                builder.without::<B>();
356            })
357            .build();
358        dbg!(&query_b.component_access);
359        assert_eq!(2, query_b.iter(&world).count());
360
361        let mut query_c = QueryBuilder::<&D>::new(&mut world)
362            .or(|builder| {
363                builder.with::<A>();
364                builder.with::<B>();
365                builder.with::<C>();
366            })
367            .build();
368        assert_eq!(3, query_c.iter(&world).count());
369    }
370
371    #[test]
372    fn builder_transmute() {
373        let mut world = World::new();
374        world.spawn(A(0));
375        world.spawn((A(1), B(0)));
376        let mut query = QueryBuilder::<()>::new(&mut world)
377            .with::<B>()
378            .transmute::<&A>()
379            .build();
380
381        query.iter(&world).for_each(|a| assert_eq!(a.0, 1));
382    }
383
384    #[test]
385    fn builder_static_components() {
386        let mut world = World::new();
387        let entity = world.spawn((A(0), B(1))).id();
388
389        let mut query = QueryBuilder::<FilteredEntityRef>::new(&mut world)
390            .data::<&A>()
391            .data::<&B>()
392            .build();
393
394        let entity_ref = query.single(&world).unwrap();
395
396        assert_eq!(entity, entity_ref.id());
397
398        let a = entity_ref.get::<A>().unwrap();
399        let b = entity_ref.get::<B>().unwrap();
400
401        assert_eq!(0, a.0);
402        assert_eq!(1, b.0);
403    }
404
405    #[test]
406    fn builder_dynamic_components() {
407        let mut world = World::new();
408        let entity = world.spawn((A(0), B(1))).id();
409        let component_id_a = world.register_component::<A>();
410        let component_id_b = world.register_component::<B>();
411
412        let mut query = QueryBuilder::<FilteredEntityRef>::new(&mut world)
413            .ref_id(component_id_a)
414            .ref_id(component_id_b)
415            .build();
416
417        let entity_ref = query.single(&world).unwrap();
418
419        assert_eq!(entity, entity_ref.id());
420
421        let a = entity_ref.get_by_id(component_id_a).unwrap();
422        let b = entity_ref.get_by_id(component_id_b).unwrap();
423
424        // SAFETY: We set these pointers to point to these components
425        unsafe {
426            assert_eq!(0, a.deref::<A>().0);
427            assert_eq!(1, b.deref::<B>().0);
428        }
429    }
430
431    #[test]
432    fn builder_provide_access() {
433        let mut world = World::new();
434        world.spawn((A(0), B(1), D));
435
436        let mut query =
437            QueryBuilder::<(Entity, FilteredEntityRef, FilteredEntityMut), With<D>>::new(
438                &mut world,
439            )
440            .data::<&mut A>()
441            .data::<&B>()
442            .build();
443
444        // The `FilteredEntityRef` only has read access, so the `FilteredEntityMut` can have read access without conflicts
445        let (_entity, entity_ref_1, mut entity_ref_2) = query.single_mut(&mut world).unwrap();
446        assert!(entity_ref_1.get::<A>().is_some());
447        assert!(entity_ref_1.get::<B>().is_some());
448        assert!(entity_ref_2.get::<A>().is_some());
449        assert!(entity_ref_2.get_mut::<A>().is_none());
450        assert!(entity_ref_2.get::<B>().is_some());
451        assert!(entity_ref_2.get_mut::<B>().is_none());
452
453        let mut query =
454            QueryBuilder::<(Entity, FilteredEntityMut, FilteredEntityMut), With<D>>::new(
455                &mut world,
456            )
457            .data::<&mut A>()
458            .data::<&B>()
459            .build();
460
461        // The first `FilteredEntityMut` has write access to A, so the second one cannot have write access
462        let (_entity, mut entity_ref_1, mut entity_ref_2) = query.single_mut(&mut world).unwrap();
463        assert!(entity_ref_1.get::<A>().is_some());
464        assert!(entity_ref_1.get_mut::<A>().is_some());
465        assert!(entity_ref_1.get::<B>().is_some());
466        assert!(entity_ref_1.get_mut::<B>().is_none());
467        assert!(entity_ref_2.get::<A>().is_none());
468        assert!(entity_ref_2.get_mut::<A>().is_none());
469        assert!(entity_ref_2.get::<B>().is_some());
470        assert!(entity_ref_2.get_mut::<B>().is_none());
471
472        let mut query = QueryBuilder::<(FilteredEntityMut, &mut A, &B), With<D>>::new(&mut world)
473            .data::<&mut A>()
474            .data::<&mut B>()
475            .build();
476
477        // Any `A` access would conflict with `&mut A`, and write access to `B` would conflict with `&B`.
478        let (mut entity_ref, _a, _b) = query.single_mut(&mut world).unwrap();
479        assert!(entity_ref.get::<A>().is_none());
480        assert!(entity_ref.get_mut::<A>().is_none());
481        assert!(entity_ref.get::<B>().is_some());
482        assert!(entity_ref.get_mut::<B>().is_none());
483
484        let mut query = QueryBuilder::<(FilteredEntityMut, &mut A, &B), With<D>>::new(&mut world)
485            .data::<EntityMut>()
486            .build();
487
488        // Same as above, but starting from "all" access
489        let (mut entity_ref, _a, _b) = query.single_mut(&mut world).unwrap();
490        assert!(entity_ref.get::<A>().is_none());
491        assert!(entity_ref.get_mut::<A>().is_none());
492        assert!(entity_ref.get::<B>().is_some());
493        assert!(entity_ref.get_mut::<B>().is_none());
494
495        let mut query =
496            QueryBuilder::<(FilteredEntityMut, EntityMutExcept<A>), With<D>>::new(&mut world)
497                .data::<EntityMut>()
498                .build();
499
500        // Removing `EntityMutExcept<A>` just leaves A
501        let (mut entity_ref_1, _entity_ref_2) = query.single_mut(&mut world).unwrap();
502        assert!(entity_ref_1.get::<A>().is_some());
503        assert!(entity_ref_1.get_mut::<A>().is_some());
504        assert!(entity_ref_1.get::<B>().is_none());
505        assert!(entity_ref_1.get_mut::<B>().is_none());
506
507        let mut query =
508            QueryBuilder::<(FilteredEntityMut, EntityRefExcept<A>), With<D>>::new(&mut world)
509                .data::<EntityMut>()
510                .build();
511
512        // Removing `EntityRefExcept<A>` just leaves A, plus read access
513        let (mut entity_ref_1, _entity_ref_2) = query.single_mut(&mut world).unwrap();
514        assert!(entity_ref_1.get::<A>().is_some());
515        assert!(entity_ref_1.get_mut::<A>().is_some());
516        assert!(entity_ref_1.get::<B>().is_some());
517        assert!(entity_ref_1.get_mut::<B>().is_none());
518    }
519
520    /// Regression test for issue #14348
521    #[test]
522    fn builder_static_dense_dynamic_sparse() {
523        #[derive(Component)]
524        struct Dense;
525
526        #[derive(Component)]
527        #[component(storage = "SparseSet")]
528        struct Sparse;
529
530        let mut world = World::new();
531
532        world.spawn(Dense);
533        world.spawn((Dense, Sparse));
534
535        let mut query = QueryBuilder::<&Dense>::new(&mut world)
536            .with::<Sparse>()
537            .build();
538
539        let matched = query.iter(&world).count();
540        assert_eq!(matched, 1);
541    }
542
543    #[test]
544    fn builder_dynamic_can_be_dense() {
545        #[derive(Component)]
546        #[component(storage = "SparseSet")]
547        struct Sparse;
548
549        let mut world = World::new();
550
551        // FilteredEntityRef and FilteredEntityMut are dense by default
552        let query = QueryBuilder::<FilteredEntityRef>::new(&mut world).build();
553        assert!(query.is_dense);
554
555        let query = QueryBuilder::<FilteredEntityMut>::new(&mut world).build();
556        assert!(query.is_dense);
557
558        // Adding a required sparse term makes the query sparse
559        let query = QueryBuilder::<FilteredEntityRef>::new(&mut world)
560            .data::<&Sparse>()
561            .build();
562        assert!(!query.is_dense);
563
564        // Adding an optional sparse term lets it remain dense
565        let query = QueryBuilder::<FilteredEntityRef>::new(&mut world)
566            .data::<Option<&Sparse>>()
567            .build();
568        assert!(query.is_dense);
569    }
570}
bevy_ecs/query/builder.rs

bevy_ecs/query/
builder.rs
