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        let Ok(component_accesses) = self.access.access().try_iter_component_access() else {
85            // Access is unbounded, pessimistically assume it's sparse.
86            return false;
87        };
88
89        component_accesses
90            .map(|access| *access.index())
91            .all(is_dense)
92            && !self.access.access().has_read_all_components()
93            && self.access.with_filters().all(is_dense)
94            && self.access.without_filters().all(is_dense)
95    }
96
97    /// Returns a reference to the world passed to [`Self::new`].
98    pub fn world(&self) -> &World {
99        self.world
100    }
101
102    /// Returns a mutable reference to the world passed to [`Self::new`].
103    pub fn world_mut(&mut self) -> &mut World {
104        self.world
105    }
106
107    /// Adds access to self's underlying [`FilteredAccess`] respecting [`Self::or`] and [`Self::and`]
108    pub fn extend_access(&mut self, mut access: FilteredAccess) {
109        if self.or {
110            if self.first {
111                access.required.clear();
112                self.access.extend(&access);
113                self.first = false;
114            } else {
115                self.access.append_or(&access);
116            }
117        } else {
118            self.access.extend(&access);
119        }
120    }
121
122    /// Adds accesses required for `T` to self.
123    pub fn data<T: QueryData>(&mut self) -> &mut Self {
124        let state = T::init_state(self.world);
125        let mut access = FilteredAccess::default();
126        T::update_component_access(&state, &mut access);
127        self.extend_access(access);
128        self
129    }
130
131    /// Adds filter from `T` to self.
132    pub fn filter<T: QueryFilter>(&mut self) -> &mut Self {
133        let state = T::init_state(self.world);
134        let mut access = FilteredAccess::default();
135        T::update_component_access(&state, &mut access);
136        self.extend_access(access);
137        self
138    }
139
140    /// Adds [`With<T>`] to the [`FilteredAccess`] of self.
141    pub fn with<T: Component>(&mut self) -> &mut Self {
142        self.filter::<With<T>>();
143        self
144    }
145
146    /// Adds [`With<T>`] to the [`FilteredAccess`] of self from a runtime [`ComponentId`].
147    pub fn with_id(&mut self, id: ComponentId) -> &mut Self {
148        let mut access = FilteredAccess::default();
149        access.and_with(id);
150        self.extend_access(access);
151        self
152    }
153
154    /// Adds [`Without<T>`] to the [`FilteredAccess`] of self.
155    pub fn without<T: Component>(&mut self) -> &mut Self {
156        self.filter::<Without<T>>();
157        self
158    }
159
160    /// Adds [`Without<T>`] to the [`FilteredAccess`] of self from a runtime [`ComponentId`].
161    pub fn without_id(&mut self, id: ComponentId) -> &mut Self {
162        let mut access = FilteredAccess::default();
163        access.and_without(id);
164        self.extend_access(access);
165        self
166    }
167
168    /// Adds `&T` to the [`FilteredAccess`] of self.
169    pub fn ref_id(&mut self, id: ComponentId) -> &mut Self {
170        self.with_id(id);
171        self.access.add_component_read(id);
172        self
173    }
174
175    /// Adds `&mut T` to the [`FilteredAccess`] of self.
176    pub fn mut_id(&mut self, id: ComponentId) -> &mut Self {
177        self.with_id(id);
178        self.access.add_component_write(id);
179        self
180    }
181
182    /// Takes a function over mutable access to a [`QueryBuilder`], calls that function
183    /// on an empty builder and then adds all accesses from that builder to self as optional.
184    pub fn optional(&mut self, f: impl Fn(&mut QueryBuilder)) -> &mut Self {
185        let mut builder = QueryBuilder::new(self.world);
186        f(&mut builder);
187        self.access.extend_access(builder.access());
188        self
189    }
190
191    /// Takes a function over mutable access to a [`QueryBuilder`], calls that function
192    /// on an empty builder and then adds all accesses from that builder to self.
193    ///
194    /// Primarily used when inside a [`Self::or`] closure to group several terms.
195    pub fn and(&mut self, f: impl Fn(&mut QueryBuilder)) -> &mut Self {
196        let mut builder = QueryBuilder::new(self.world);
197        f(&mut builder);
198        let access = builder.access().clone();
199        self.extend_access(access);
200        self
201    }
202
203    /// Takes a function over mutable access to a [`QueryBuilder`], calls that function
204    /// on an empty builder, all accesses added to that builder will become terms in an or expression.
205    ///
206    /// ```
207    /// # use bevy_ecs::prelude::*;
208    /// #
209    /// # #[derive(Component)]
210    /// # struct A;
211    /// #
212    /// # #[derive(Component)]
213    /// # struct B;
214    /// #
215    /// # let mut world = World::new();
216    /// #
217    /// QueryBuilder::<Entity>::new(&mut world).or(|builder| {
218    ///     builder.with::<A>();
219    ///     builder.with::<B>();
220    /// });
221    /// // is equivalent to
222    /// QueryBuilder::<Entity>::new(&mut world).filter::<Or<(With<A>, With<B>)>>();
223    /// ```
224    pub fn or(&mut self, f: impl Fn(&mut QueryBuilder)) -> &mut Self {
225        let mut builder = QueryBuilder::new(self.world);
226        builder.or = true;
227        builder.first = true;
228        f(&mut builder);
229        self.access.extend(builder.access());
230        self
231    }
232
233    /// Returns a reference to the [`FilteredAccess`] that will be provided to the built [`Query`].
234    pub fn access(&self) -> &FilteredAccess {
235        &self.access
236    }
237
238    /// Transmute the existing builder adding required accesses.
239    /// This will maintain all existing accesses.
240    ///
241    /// If including a filter type see [`Self::transmute_filtered`]
242    pub fn transmute<NewD: QueryData>(&mut self) -> &mut QueryBuilder<'w, NewD> {
243        self.transmute_filtered::<NewD, ()>()
244    }
245
246    /// Transmute the existing builder adding required accesses.
247    /// This will maintain all existing accesses.
248    pub fn transmute_filtered<NewD: QueryData, NewF: QueryFilter>(
249        &mut self,
250    ) -> &mut QueryBuilder<'w, NewD, NewF> {
251        let fetch_state = NewD::init_state(self.world);
252        let filter_state = NewF::init_state(self.world);
253
254        let mut access = FilteredAccess::default();
255        NewD::update_component_access(&fetch_state, &mut access);
256        NewF::update_component_access(&filter_state, &mut access);
257
258        self.extend_access(access);
259        // SAFETY:
260        // - We have included all required accesses for NewQ and NewF
261        // - The layout of all QueryBuilder instances is the same
262        unsafe { core::mem::transmute(self) }
263    }
264
265    /// Create a [`QueryState`] with the accesses of the builder.
266    ///
267    /// Takes `&mut self` to access the inner world reference while initializing
268    /// state for the new [`QueryState`]
269    pub fn build(&mut self) -> QueryState<D, F> {
270        QueryState::<D, F>::from_builder(self)
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use crate::{
277        prelude::*,
278        world::{EntityMutExcept, EntityRefExcept, FilteredEntityMut, FilteredEntityRef},
279    };
280    use std::dbg;
281
282    #[derive(Component, PartialEq, Debug)]
283    struct A(usize);
284
285    #[derive(Component, PartialEq, Debug)]
286    struct B(usize);
287
288    #[derive(Component, PartialEq, Debug)]
289    struct C(usize);
290
291    #[test]
292    fn builder_with_without_static() {
293        let mut world = World::new();
294        let entity_a = world.spawn((A(0), B(0))).id();
295        let entity_b = world.spawn((A(0), C(0))).id();
296
297        let mut query_a = QueryBuilder::<Entity>::new(&mut world)
298            .with::<A>()
299            .without::<C>()
300            .build();
301        assert_eq!(entity_a, query_a.single(&world).unwrap());
302
303        let mut query_b = QueryBuilder::<Entity>::new(&mut world)
304            .with::<A>()
305            .without::<B>()
306            .build();
307        assert_eq!(entity_b, query_b.single(&world).unwrap());
308    }
309
310    #[test]
311    fn builder_with_without_dynamic() {
312        let mut world = World::new();
313        let entity_a = world.spawn((A(0), B(0))).id();
314        let entity_b = world.spawn((A(0), C(0))).id();
315        let component_id_a = world.register_component::<A>();
316        let component_id_b = world.register_component::<B>();
317        let component_id_c = world.register_component::<C>();
318
319        let mut query_a = QueryBuilder::<Entity>::new(&mut world)
320            .with_id(component_id_a)
321            .without_id(component_id_c)
322            .build();
323        assert_eq!(entity_a, query_a.single(&world).unwrap());
324
325        let mut query_b = QueryBuilder::<Entity>::new(&mut world)
326            .with_id(component_id_a)
327            .without_id(component_id_b)
328            .build();
329        assert_eq!(entity_b, query_b.single(&world).unwrap());
330    }
331
332    #[test]
333    fn builder_or() {
334        let mut world = World::new();
335        world.spawn((A(0), B(0)));
336        world.spawn(B(0));
337        world.spawn(C(0));
338
339        let mut query_a = QueryBuilder::<Entity>::new(&mut world)
340            .or(|builder| {
341                builder.with::<A>();
342                builder.with::<B>();
343            })
344            .build();
345        assert_eq!(2, query_a.iter(&world).count());
346
347        let mut query_b = QueryBuilder::<Entity>::new(&mut world)
348            .or(|builder| {
349                builder.with::<A>();
350                builder.without::<B>();
351            })
352            .build();
353        dbg!(&query_b.component_access);
354        assert_eq!(2, query_b.iter(&world).count());
355
356        let mut query_c = QueryBuilder::<Entity>::new(&mut world)
357            .or(|builder| {
358                builder.with::<A>();
359                builder.with::<B>();
360                builder.with::<C>();
361            })
362            .build();
363        assert_eq!(3, query_c.iter(&world).count());
364    }
365
366    #[test]
367    fn builder_transmute() {
368        let mut world = World::new();
369        world.spawn(A(0));
370        world.spawn((A(1), B(0)));
371        let mut query = QueryBuilder::<()>::new(&mut world)
372            .with::<B>()
373            .transmute::<&A>()
374            .build();
375
376        query.iter(&world).for_each(|a| assert_eq!(a.0, 1));
377    }
378
379    #[test]
380    fn builder_static_components() {
381        let mut world = World::new();
382        let entity = world.spawn((A(0), B(1))).id();
383
384        let mut query = QueryBuilder::<FilteredEntityRef>::new(&mut world)
385            .data::<&A>()
386            .data::<&B>()
387            .build();
388
389        let entity_ref = query.single(&world).unwrap();
390
391        assert_eq!(entity, entity_ref.id());
392
393        let a = entity_ref.get::<A>().unwrap();
394        let b = entity_ref.get::<B>().unwrap();
395
396        assert_eq!(0, a.0);
397        assert_eq!(1, b.0);
398    }
399
400    #[test]
401    fn builder_dynamic_components() {
402        let mut world = World::new();
403        let entity = world.spawn((A(0), B(1))).id();
404        let component_id_a = world.register_component::<A>();
405        let component_id_b = world.register_component::<B>();
406
407        let mut query = QueryBuilder::<FilteredEntityRef>::new(&mut world)
408            .ref_id(component_id_a)
409            .ref_id(component_id_b)
410            .build();
411
412        let entity_ref = query.single(&world).unwrap();
413
414        assert_eq!(entity, entity_ref.id());
415
416        let a = entity_ref.get_by_id(component_id_a).unwrap();
417        let b = entity_ref.get_by_id(component_id_b).unwrap();
418
419        // SAFETY: We set these pointers to point to these components
420        unsafe {
421            assert_eq!(0, a.deref::<A>().0);
422            assert_eq!(1, b.deref::<B>().0);
423        }
424    }
425
426    #[test]
427    fn builder_provide_access() {
428        let mut world = World::new();
429        world.spawn((A(0), B(1)));
430
431        let mut query =
432            QueryBuilder::<(Entity, FilteredEntityRef, FilteredEntityMut)>::new(&mut world)
433                .data::<&mut A>()
434                .data::<&B>()
435                .build();
436
437        // The `FilteredEntityRef` only has read access, so the `FilteredEntityMut` can have read access without conflicts
438        let (_entity, entity_ref_1, mut entity_ref_2) = query.single_mut(&mut world).unwrap();
439        assert!(entity_ref_1.get::<A>().is_some());
440        assert!(entity_ref_1.get::<B>().is_some());
441        assert!(entity_ref_2.get::<A>().is_some());
442        assert!(entity_ref_2.get_mut::<A>().is_none());
443        assert!(entity_ref_2.get::<B>().is_some());
444        assert!(entity_ref_2.get_mut::<B>().is_none());
445
446        let mut query =
447            QueryBuilder::<(Entity, FilteredEntityMut, FilteredEntityMut)>::new(&mut world)
448                .data::<&mut A>()
449                .data::<&B>()
450                .build();
451
452        // The first `FilteredEntityMut` has write access to A, so the second one cannot have write access
453        let (_entity, mut entity_ref_1, mut entity_ref_2) = query.single_mut(&mut world).unwrap();
454        assert!(entity_ref_1.get::<A>().is_some());
455        assert!(entity_ref_1.get_mut::<A>().is_some());
456        assert!(entity_ref_1.get::<B>().is_some());
457        assert!(entity_ref_1.get_mut::<B>().is_none());
458        assert!(entity_ref_2.get::<A>().is_none());
459        assert!(entity_ref_2.get_mut::<A>().is_none());
460        assert!(entity_ref_2.get::<B>().is_some());
461        assert!(entity_ref_2.get_mut::<B>().is_none());
462
463        let mut query = QueryBuilder::<(FilteredEntityMut, &mut A, &B)>::new(&mut world)
464            .data::<&mut A>()
465            .data::<&mut B>()
466            .build();
467
468        // Any `A` access would conflict with `&mut A`, and write access to `B` would conflict with `&B`.
469        let (mut entity_ref, _a, _b) = query.single_mut(&mut world).unwrap();
470        assert!(entity_ref.get::<A>().is_none());
471        assert!(entity_ref.get_mut::<A>().is_none());
472        assert!(entity_ref.get::<B>().is_some());
473        assert!(entity_ref.get_mut::<B>().is_none());
474
475        let mut query = QueryBuilder::<(FilteredEntityMut, &mut A, &B)>::new(&mut world)
476            .data::<EntityMut>()
477            .build();
478
479        // Same as above, but starting from "all" access
480        let (mut entity_ref, _a, _b) = query.single_mut(&mut world).unwrap();
481        assert!(entity_ref.get::<A>().is_none());
482        assert!(entity_ref.get_mut::<A>().is_none());
483        assert!(entity_ref.get::<B>().is_some());
484        assert!(entity_ref.get_mut::<B>().is_none());
485
486        let mut query = QueryBuilder::<(FilteredEntityMut, EntityMutExcept<A>)>::new(&mut world)
487            .data::<EntityMut>()
488            .build();
489
490        // Removing `EntityMutExcept<A>` just leaves A
491        let (mut entity_ref_1, _entity_ref_2) = query.single_mut(&mut world).unwrap();
492        assert!(entity_ref_1.get::<A>().is_some());
493        assert!(entity_ref_1.get_mut::<A>().is_some());
494        assert!(entity_ref_1.get::<B>().is_none());
495        assert!(entity_ref_1.get_mut::<B>().is_none());
496
497        let mut query = QueryBuilder::<(FilteredEntityMut, EntityRefExcept<A>)>::new(&mut world)
498            .data::<EntityMut>()
499            .build();
500
501        // Removing `EntityRefExcept<A>` just leaves A, plus read access
502        let (mut entity_ref_1, _entity_ref_2) = query.single_mut(&mut world).unwrap();
503        assert!(entity_ref_1.get::<A>().is_some());
504        assert!(entity_ref_1.get_mut::<A>().is_some());
505        assert!(entity_ref_1.get::<B>().is_some());
506        assert!(entity_ref_1.get_mut::<B>().is_none());
507    }
508
509    /// Regression test for issue #14348
510    #[test]
511    fn builder_static_dense_dynamic_sparse() {
512        #[derive(Component)]
513        struct Dense;
514
515        #[derive(Component)]
516        #[component(storage = "SparseSet")]
517        struct Sparse;
518
519        let mut world = World::new();
520
521        world.spawn(Dense);
522        world.spawn((Dense, Sparse));
523
524        let mut query = QueryBuilder::<&Dense>::new(&mut world)
525            .with::<Sparse>()
526            .build();
527
528        let matched = query.iter(&world).count();
529        assert_eq!(matched, 1);
530    }
531}
bevy_ecs/query/builder.rs

bevy_ecs/query/
builder.rs
