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<ComponentId>,
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<ComponentId>) {
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<ComponentId> {
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 mut fetch_state = NewD::init_state(self.world);
252        let filter_state = NewF::init_state(self.world);
253
254        NewD::set_access(&mut fetch_state, &self.access);
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::{prelude::*, world::FilteredEntityRef};
279    use std::dbg;
280
281    #[derive(Component, PartialEq, Debug)]
282    struct A(usize);
283
284    #[derive(Component, PartialEq, Debug)]
285    struct B(usize);
286
287    #[derive(Component, PartialEq, Debug)]
288    struct C(usize);
289
290    #[test]
291    fn builder_with_without_static() {
292        let mut world = World::new();
293        let entity_a = world.spawn((A(0), B(0))).id();
294        let entity_b = world.spawn((A(0), C(0))).id();
295
296        let mut query_a = QueryBuilder::<Entity>::new(&mut world)
297            .with::<A>()
298            .without::<C>()
299            .build();
300        assert_eq!(entity_a, query_a.single(&world).unwrap());
301
302        let mut query_b = QueryBuilder::<Entity>::new(&mut world)
303            .with::<A>()
304            .without::<B>()
305            .build();
306        assert_eq!(entity_b, query_b.single(&world).unwrap());
307    }
308
309    #[test]
310    fn builder_with_without_dynamic() {
311        let mut world = World::new();
312        let entity_a = world.spawn((A(0), B(0))).id();
313        let entity_b = world.spawn((A(0), C(0))).id();
314        let component_id_a = world.register_component::<A>();
315        let component_id_b = world.register_component::<B>();
316        let component_id_c = world.register_component::<C>();
317
318        let mut query_a = QueryBuilder::<Entity>::new(&mut world)
319            .with_id(component_id_a)
320            .without_id(component_id_c)
321            .build();
322        assert_eq!(entity_a, query_a.single(&world).unwrap());
323
324        let mut query_b = QueryBuilder::<Entity>::new(&mut world)
325            .with_id(component_id_a)
326            .without_id(component_id_b)
327            .build();
328        assert_eq!(entity_b, query_b.single(&world).unwrap());
329    }
330
331    #[test]
332    fn builder_or() {
333        let mut world = World::new();
334        world.spawn((A(0), B(0)));
335        world.spawn(B(0));
336        world.spawn(C(0));
337
338        let mut query_a = QueryBuilder::<Entity>::new(&mut world)
339            .or(|builder| {
340                builder.with::<A>();
341                builder.with::<B>();
342            })
343            .build();
344        assert_eq!(2, query_a.iter(&world).count());
345
346        let mut query_b = QueryBuilder::<Entity>::new(&mut world)
347            .or(|builder| {
348                builder.with::<A>();
349                builder.without::<B>();
350            })
351            .build();
352        dbg!(&query_b.component_access);
353        assert_eq!(2, query_b.iter(&world).count());
354
355        let mut query_c = QueryBuilder::<Entity>::new(&mut world)
356            .or(|builder| {
357                builder.with::<A>();
358                builder.with::<B>();
359                builder.with::<C>();
360            })
361            .build();
362        assert_eq!(3, query_c.iter(&world).count());
363    }
364
365    #[test]
366    fn builder_transmute() {
367        let mut world = World::new();
368        world.spawn(A(0));
369        world.spawn((A(1), B(0)));
370        let mut query = QueryBuilder::<()>::new(&mut world)
371            .with::<B>()
372            .transmute::<&A>()
373            .build();
374
375        query.iter(&world).for_each(|a| assert_eq!(a.0, 1));
376    }
377
378    #[test]
379    fn builder_static_components() {
380        let mut world = World::new();
381        let entity = world.spawn((A(0), B(1))).id();
382
383        let mut query = QueryBuilder::<FilteredEntityRef>::new(&mut world)
384            .data::<&A>()
385            .data::<&B>()
386            .build();
387
388        let entity_ref = query.single(&world).unwrap();
389
390        assert_eq!(entity, entity_ref.id());
391
392        let a = entity_ref.get::<A>().unwrap();
393        let b = entity_ref.get::<B>().unwrap();
394
395        assert_eq!(0, a.0);
396        assert_eq!(1, b.0);
397    }
398
399    #[test]
400    fn builder_dynamic_components() {
401        let mut world = World::new();
402        let entity = world.spawn((A(0), B(1))).id();
403        let component_id_a = world.register_component::<A>();
404        let component_id_b = world.register_component::<B>();
405
406        let mut query = QueryBuilder::<FilteredEntityRef>::new(&mut world)
407            .ref_id(component_id_a)
408            .ref_id(component_id_b)
409            .build();
410
411        let entity_ref = query.single(&world).unwrap();
412
413        assert_eq!(entity, entity_ref.id());
414
415        let a = entity_ref.get_by_id(component_id_a).unwrap();
416        let b = entity_ref.get_by_id(component_id_b).unwrap();
417
418        // SAFETY: We set these pointers to point to these components
419        unsafe {
420            assert_eq!(0, a.deref::<A>().0);
421            assert_eq!(1, b.deref::<B>().0);
422        }
423    }
424
425    /// Regression test for issue #14348
426    #[test]
427    fn builder_static_dense_dynamic_sparse() {
428        #[derive(Component)]
429        struct Dense;
430
431        #[derive(Component)]
432        #[component(storage = "SparseSet")]
433        struct Sparse;
434
435        let mut world = World::new();
436
437        world.spawn(Dense);
438        world.spawn((Dense, Sparse));
439
440        let mut query = QueryBuilder::<&Dense>::new(&mut world)
441            .with::<Sparse>()
442            .build();
443
444        let matched = query.iter(&world).count();
445        assert_eq!(matched, 1);
446    }
447}
bevy_ecs/query/builder.rs

bevy_ecs/query/
builder.rs
