bevy_ecs/query/iter.rs
1use crate::{
2 archetype::{Archetype, ArchetypeEntity, Archetypes},
3 component::Tick,
4 entity::{Entities, Entity},
5 query::{ArchetypeFilter, DebugCheckedUnwrap, QueryState, StorageId},
6 storage::{Table, TableRow, Tables},
7 world::unsafe_world_cell::UnsafeWorldCell,
8};
9use std::{
10 borrow::Borrow,
11 cmp::Ordering,
12 fmt::{self, Debug, Formatter},
13 iter::FusedIterator,
14 mem::MaybeUninit,
15 ops::Range,
16};
17
18use super::{QueryData, QueryFilter, ReadOnlyQueryData};
19
20/// An [`Iterator`] over query results of a [`Query`](crate::system::Query).
21///
22/// This struct is created by the [`Query::iter`](crate::system::Query::iter) and
23/// [`Query::iter_mut`](crate::system::Query::iter_mut) methods.
24pub struct QueryIter<'w, 's, D: QueryData, F: QueryFilter> {
25 world: UnsafeWorldCell<'w>,
26 tables: &'w Tables,
27 archetypes: &'w Archetypes,
28 query_state: &'s QueryState<D, F>,
29 cursor: QueryIterationCursor<'w, 's, D, F>,
30}
31
32impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
33 /// # Safety
34 /// - `world` must have permission to access any of the components registered in `query_state`.
35 /// - `world` must be the same one used to initialize `query_state`.
36 pub(crate) unsafe fn new(
37 world: UnsafeWorldCell<'w>,
38 query_state: &'s QueryState<D, F>,
39 last_run: Tick,
40 this_run: Tick,
41 ) -> Self {
42 QueryIter {
43 world,
44 query_state,
45 // SAFETY: We only access table data that has been registered in `query_state`.
46 tables: unsafe { &world.storages().tables },
47 archetypes: world.archetypes(),
48 // SAFETY: The invariants are uphold by the caller.
49 cursor: unsafe { QueryIterationCursor::init(world, query_state, last_run, this_run) },
50 }
51 }
52
53 /// Executes the equivalent of [`Iterator::fold`] over a contiguous segment
54 /// from an table.
55 ///
56 /// # Safety
57 /// - all `rows` must be in `[0, table.entity_count)`.
58 /// - `table` must match D and F
59 /// - Both `D::IS_DENSE` and `F::IS_DENSE` must be true.
60 #[inline]
61 pub(super) unsafe fn fold_over_table_range<B, Func>(
62 &mut self,
63 mut accum: B,
64 func: &mut Func,
65 table: &'w Table,
66 rows: Range<usize>,
67 ) -> B
68 where
69 Func: FnMut(B, D::Item<'w>) -> B,
70 {
71 assert!(
72 rows.end <= u32::MAX as usize,
73 "TableRow is only valid up to u32::MAX"
74 );
75
76 D::set_table(&mut self.cursor.fetch, &self.query_state.fetch_state, table);
77 F::set_table(
78 &mut self.cursor.filter,
79 &self.query_state.filter_state,
80 table,
81 );
82
83 let entities = table.entities();
84 for row in rows {
85 // SAFETY: Caller assures `row` in range of the current archetype.
86 let entity = unsafe { entities.get_unchecked(row) };
87 let row = TableRow::from_usize(row);
88
89 // SAFETY: set_table was called prior.
90 // Caller assures `row` in range of the current archetype.
91 let fetched = unsafe { !F::filter_fetch(&mut self.cursor.filter, *entity, row) };
92 if fetched {
93 continue;
94 }
95
96 // SAFETY: set_table was called prior.
97 // Caller assures `row` in range of the current archetype.
98 let item = D::fetch(&mut self.cursor.fetch, *entity, row);
99
100 accum = func(accum, item);
101 }
102 accum
103 }
104
105 /// Executes the equivalent of [`Iterator::fold`] over a contiguous segment
106 /// from an archetype.
107 ///
108 /// # Safety
109 /// - all `indices` must be in `[0, archetype.len())`.
110 /// - `archetype` must match D and F
111 /// - Either `D::IS_DENSE` or `F::IS_DENSE` must be false.
112 #[inline]
113 pub(super) unsafe fn fold_over_archetype_range<B, Func>(
114 &mut self,
115 mut accum: B,
116 func: &mut Func,
117 archetype: &'w Archetype,
118 indices: Range<usize>,
119 ) -> B
120 where
121 Func: FnMut(B, D::Item<'w>) -> B,
122 {
123 let table = self.tables.get(archetype.table_id()).debug_checked_unwrap();
124 D::set_archetype(
125 &mut self.cursor.fetch,
126 &self.query_state.fetch_state,
127 archetype,
128 table,
129 );
130 F::set_archetype(
131 &mut self.cursor.filter,
132 &self.query_state.filter_state,
133 archetype,
134 table,
135 );
136
137 let entities = archetype.entities();
138 for index in indices {
139 // SAFETY: Caller assures `index` in range of the current archetype.
140 let archetype_entity = unsafe { entities.get_unchecked(index) };
141
142 // SAFETY: set_archetype was called prior.
143 // Caller assures `index` in range of the current archetype.
144 let fetched = unsafe {
145 !F::filter_fetch(
146 &mut self.cursor.filter,
147 archetype_entity.id(),
148 archetype_entity.table_row(),
149 )
150 };
151 if fetched {
152 continue;
153 }
154
155 // SAFETY: set_archetype was called prior, `index` is an archetype index in range of the current archetype
156 // Caller assures `index` in range of the current archetype.
157 let item = unsafe {
158 D::fetch(
159 &mut self.cursor.fetch,
160 archetype_entity.id(),
161 archetype_entity.table_row(),
162 )
163 };
164
165 accum = func(accum, item);
166 }
167 accum
168 }
169
170 /// Sorts all query items into a new iterator, using the query lens as a key.
171 ///
172 /// This sort is stable (i.e., does not reorder equal elements).
173 ///
174 /// This uses [`slice::sort`] internally.
175 ///
176 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
177 /// This includes the allowed parameter type changes listed under [allowed transmutes].
178 /// However, the lens uses the filter of the original query when present.
179 ///
180 /// The sort is not cached across system runs.
181 ///
182 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
183 ///
184 /// # Panics
185 ///
186 /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
187 ///
188 /// # Examples
189 /// ```rust
190 /// # use bevy_ecs::prelude::*;
191 /// # use std::{ops::{Deref, DerefMut}, iter::Sum};
192 /// #
193 /// # #[derive(Component)]
194 /// # struct PartMarker;
195 /// #
196 /// # #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
197 /// # struct PartIndex(usize);
198 /// #
199 /// # #[derive(Component, Clone, Copy)]
200 /// # struct PartValue(f32);
201 /// #
202 /// # impl Deref for PartValue {
203 /// # type Target = f32;
204 /// #
205 /// # fn deref(&self) -> &Self::Target {
206 /// # &self.0
207 /// # }
208 /// # }
209 /// #
210 /// # #[derive(Component)]
211 /// # struct ParentValue(f32);
212 /// #
213 /// # impl Deref for ParentValue {
214 /// # type Target = f32;
215 /// #
216 /// # fn deref(&self) -> &Self::Target {
217 /// # &self.0
218 /// # }
219 /// # }
220 /// #
221 /// # impl DerefMut for ParentValue {
222 /// # fn deref_mut(&mut self) -> &mut Self::Target {
223 /// # &mut self.0
224 /// # }
225 /// # }
226 /// #
227 /// # #[derive(Component, Debug, PartialEq, Eq, PartialOrd, Ord)]
228 /// # struct Length(usize);
229 /// #
230 /// # #[derive(Component, Debug, PartialEq, Eq, PartialOrd, Ord)]
231 /// # struct Width(usize);
232 /// #
233 /// # #[derive(Component, Debug, PartialEq, Eq, PartialOrd, Ord)]
234 /// # struct Height(usize);
235 /// #
236 /// # #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
237 /// # struct ParentEntity(Entity);
238 /// #
239 /// # #[derive(Component, Clone, Copy)]
240 /// # struct ChildPartCount(usize);
241 /// #
242 /// # impl Deref for ChildPartCount {
243 /// # type Target = usize;
244 /// #
245 /// # fn deref(&self) -> &Self::Target {
246 /// # &self.0
247 /// # }
248 /// # }
249 /// # let mut world = World::new();
250 /// // We can ensure that a query always returns in the same order.
251 /// fn system_1(query: Query<(Entity, &PartIndex)>) {
252 /// let parts: Vec<(Entity, &PartIndex)> = query.iter().sort::<&PartIndex>().collect();
253 /// }
254 ///
255 /// // We can freely rearrange query components in the key.
256 /// fn system_2(query: Query<(&Length, &Width, &Height), With<PartMarker>>) {
257 /// for (length, width, height) in query.iter().sort::<(&Height, &Length, &Width)>() {
258 /// println!("height: {height:?}, width: {width:?}, length: {length:?}")
259 /// }
260 /// }
261 ///
262 /// // We can sort by Entity without including it in the original Query.
263 /// // Here, we match iteration orders between query iterators.
264 /// fn system_3(
265 /// part_query: Query<(&PartValue, &ParentEntity)>,
266 /// mut parent_query: Query<(&ChildPartCount, &mut ParentValue)>,
267 /// ) {
268 /// let part_values = &mut part_query
269 /// .into_iter()
270 /// .sort::<&ParentEntity>()
271 /// .map(|(&value, parent_entity)| *value);
272 ///
273 /// for (&child_count, mut parent_value) in parent_query.iter_mut().sort::<Entity>() {
274 /// **parent_value = part_values.take(*child_count).sum();
275 /// }
276 /// }
277 /// #
278 /// # let mut schedule = Schedule::default();
279 /// # schedule.add_systems((system_1, system_2, system_3));
280 /// # schedule.run(&mut world);
281 /// ```
282 pub fn sort<L: ReadOnlyQueryData + 'w>(
283 self,
284 ) -> QuerySortedIter<
285 'w,
286 's,
287 D,
288 F,
289 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
290 >
291 where
292 L::Item<'w>: Ord,
293 {
294 // On the first successful iteration of `QueryIterationCursor`, `archetype_entities` or `table_entities`
295 // will be set to a non-zero value. The correctness of this method relies on this.
296 // I.e. this sort method will execute if and only if `next` on `QueryIterationCursor` of a
297 // non-empty `QueryIter` has not yet been called. When empty, this sort method will not panic.
298 if !self.cursor.archetype_entities.is_empty() || !self.cursor.table_entities.is_empty() {
299 panic!("it is not valid to call sort() after next()")
300 }
301
302 let world = self.world;
303
304 let query_lens_state = self
305 .query_state
306 .transmute_filtered::<(L, Entity), F>(world.components());
307
308 // SAFETY:
309 // `self.world` has permission to access the required components.
310 // The original query iter has not been iterated on, so no items are aliased from it.
311 let query_lens = unsafe {
312 query_lens_state.iter_unchecked_manual(
313 world,
314 world.last_change_tick(),
315 world.change_tick(),
316 )
317 };
318 let mut keyed_query: Vec<_> = query_lens
319 .map(|(key, entity)| (key, NeutralOrd(entity)))
320 .collect();
321 keyed_query.sort();
322 let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity.0);
323 // SAFETY:
324 // `self.world` has permission to access the required components.
325 // Each lens query item is dropped before the respective actual query item is accessed.
326 unsafe {
327 QuerySortedIter::new(
328 world,
329 self.query_state,
330 entity_iter,
331 world.last_change_tick(),
332 world.change_tick(),
333 )
334 }
335 }
336
337 /// Sorts all query items into a new iterator, using the query lens as a key.
338 ///
339 /// This sort is unstable (i.e., may reorder equal elements).
340 ///
341 /// This uses [`slice::sort_unstable`] internally.
342 ///
343 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
344 /// This includes the allowed parameter type changes listed under [allowed transmutes]..
345 /// However, the lens uses the filter of the original query when present.
346 ///
347 /// The sort is not cached across system runs.
348 ///
349 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
350 ///
351 /// # Panics
352 ///
353 /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
354 ///
355 /// # Example
356 /// ```
357 /// # use bevy_ecs::prelude::*;
358 /// #
359 /// # let mut world = World::new();
360 /// #
361 /// # #[derive(Component)]
362 /// # struct PartMarker;
363 /// #
364 /// #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
365 /// enum Flying {
366 /// Enabled,
367 /// Disabled
368 /// };
369 ///
370 /// // We perform an unstable sort by a Component with few values.
371 /// fn system_1(query: Query<&Flying, With<PartMarker>>) {
372 /// let part_values: Vec<&Flying> = query.iter().sort_unstable::<&Flying>().collect();
373 /// }
374 /// #
375 /// # let mut schedule = Schedule::default();
376 /// # schedule.add_systems((system_1));
377 /// # schedule.run(&mut world);
378 /// ```
379 pub fn sort_unstable<L: ReadOnlyQueryData + 'w>(
380 self,
381 ) -> QuerySortedIter<
382 'w,
383 's,
384 D,
385 F,
386 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
387 >
388 where
389 L::Item<'w>: Ord,
390 {
391 // On the first successful iteration of `QueryIterationCursor`, `archetype_entities` or `table_entities`
392 // will be set to a non-zero value. The correctness of this method relies on this.
393 // I.e. this sort method will execute if and only if `next` on `QueryIterationCursor` of a
394 // non-empty `QueryIter` has not yet been called. When empty, this sort method will not panic.
395 if !self.cursor.archetype_entities.is_empty() || !self.cursor.table_entities.is_empty() {
396 panic!("it is not valid to call sort() after next()")
397 }
398
399 let world = self.world;
400
401 let query_lens_state = self
402 .query_state
403 .transmute_filtered::<(L, Entity), F>(world.components());
404
405 // SAFETY:
406 // `self.world` has permission to access the required components.
407 // The original query iter has not been iterated on, so no items are aliased from it.
408 let query_lens = unsafe {
409 query_lens_state.iter_unchecked_manual(
410 world,
411 world.last_change_tick(),
412 world.change_tick(),
413 )
414 };
415 let mut keyed_query: Vec<_> = query_lens
416 .map(|(key, entity)| (key, NeutralOrd(entity)))
417 .collect();
418 keyed_query.sort_unstable();
419 let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity.0);
420 // SAFETY:
421 // `self.world` has permission to access the required components.
422 // Each lens query item is dropped before the respective actual query item is accessed.
423 unsafe {
424 QuerySortedIter::new(
425 world,
426 self.query_state,
427 entity_iter,
428 world.last_change_tick(),
429 world.change_tick(),
430 )
431 }
432 }
433
434 /// Sorts all query items into a new iterator with a comparator function over the query lens.
435 ///
436 /// This sort is stable (i.e., does not reorder equal elements).
437 ///
438 /// This uses [`slice::sort_by`] internally.
439 ///
440 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
441 /// This includes the allowed parameter type changes listed under [allowed transmutes].
442 /// However, the lens uses the filter of the original query when present.
443 ///
444 /// The sort is not cached across system runs.
445 ///
446 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
447 ///
448 /// # Panics
449 ///
450 /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
451 ///
452 /// # Example
453 /// ```
454 /// # use bevy_ecs::prelude::*;
455 /// # use std::ops::Deref;
456 /// #
457 /// # impl Deref for PartValue {
458 /// # type Target = f32;
459 /// #
460 /// # fn deref(&self) -> &Self::Target {
461 /// # &self.0
462 /// # }
463 /// # }
464 /// #
465 /// # let mut world = World::new();
466 /// #
467 /// #[derive(Component)]
468 /// struct PartValue(f32);
469 ///
470 /// // We can use a cmp function on components do not implement Ord.
471 /// fn system_1(query: Query<&PartValue>) {
472 /// // Sort part values according to `f32::total_comp`.
473 /// let part_values: Vec<&PartValue> = query
474 /// .iter()
475 /// .sort_by::<&PartValue>(|value_1, value_2| value_1.total_cmp(*value_2))
476 /// .collect();
477 /// }
478 /// #
479 /// # let mut schedule = Schedule::default();
480 /// # schedule.add_systems((system_1));
481 /// # schedule.run(&mut world);
482 /// ```
483 pub fn sort_by<L: ReadOnlyQueryData + 'w>(
484 self,
485 mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
486 ) -> QuerySortedIter<
487 'w,
488 's,
489 D,
490 F,
491 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
492 > {
493 // On the first successful iteration of `QueryIterationCursor`, `archetype_entities` or `table_entities`
494 // will be set to a non-zero value. The correctness of this method relies on this.
495 // I.e. this sort method will execute if and only if `next` on `QueryIterationCursor` of a
496 // non-empty `QueryIter` has not yet been called. When empty, this sort method will not panic.
497 if !self.cursor.archetype_entities.is_empty() || !self.cursor.table_entities.is_empty() {
498 panic!("it is not valid to call sort() after next()")
499 }
500
501 let world = self.world;
502
503 let query_lens_state = self
504 .query_state
505 .transmute_filtered::<(L, Entity), F>(world.components());
506
507 // SAFETY:
508 // `self.world` has permission to access the required components.
509 // The original query iter has not been iterated on, so no items are aliased from it.
510 let query_lens = unsafe {
511 query_lens_state.iter_unchecked_manual(
512 world,
513 world.last_change_tick(),
514 world.change_tick(),
515 )
516 };
517 let mut keyed_query: Vec<_> = query_lens.collect();
518 keyed_query.sort_by(|(key_1, _), (key_2, _)| compare(key_1, key_2));
519 let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity);
520 // SAFETY:
521 // `self.world` has permission to access the required components.
522 // Each lens query item is dropped before the respective actual query item is accessed.
523 unsafe {
524 QuerySortedIter::new(
525 world,
526 self.query_state,
527 entity_iter,
528 world.last_change_tick(),
529 world.change_tick(),
530 )
531 }
532 }
533
534 /// Sorts all query items into a new iterator with a comparator function over the query lens.
535 ///
536 /// This sort is unstable (i.e., may reorder equal elements).
537 ///
538 /// This uses [`slice::sort_unstable_by`] internally.
539 ///
540 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
541 /// This includes the allowed parameter type changes listed under [allowed transmutes].
542 /// However, the lens uses the filter of the original query when present.
543 ///
544 /// The sort is not cached across system runs.
545 ///
546 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
547 ///
548 /// # Panics
549 ///
550 /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
551 pub fn sort_unstable_by<L: ReadOnlyQueryData + 'w>(
552 self,
553 mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
554 ) -> QuerySortedIter<
555 'w,
556 's,
557 D,
558 F,
559 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
560 > {
561 // On the first successful iteration of `QueryIterationCursor`, `archetype_entities` or `table_entities`
562 // will be set to a non-zero value. The correctness of this method relies on this.
563 // I.e. this sort method will execute if and only if `next` on `QueryIterationCursor` of a
564 // non-empty `QueryIter` has not yet been called. When empty, this sort method will not panic.
565 if !self.cursor.archetype_entities.is_empty() || !self.cursor.table_entities.is_empty() {
566 panic!("it is not valid to call sort() after next()")
567 }
568
569 let world = self.world;
570
571 let query_lens_state = self
572 .query_state
573 .transmute_filtered::<(L, Entity), F>(world.components());
574
575 // SAFETY:
576 // `self.world` has permission to access the required components.
577 // The original query iter has not been iterated on, so no items are aliased from it.
578 let query_lens = unsafe {
579 query_lens_state.iter_unchecked_manual(
580 world,
581 world.last_change_tick(),
582 world.change_tick(),
583 )
584 };
585 let mut keyed_query: Vec<_> = query_lens.collect();
586 keyed_query.sort_by(|(key_1, _), (key_2, _)| compare(key_1, key_2));
587 let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity);
588 // SAFETY:
589 // `self.world` has permission to access the required components.
590 // Each lens query item is dropped before the respective actual query item is accessed.
591 unsafe {
592 QuerySortedIter::new(
593 world,
594 self.query_state,
595 entity_iter,
596 world.last_change_tick(),
597 world.change_tick(),
598 )
599 }
600 }
601
602 /// Sorts all query items into a new iterator with a key extraction function over the query lens.
603 ///
604 /// This sort is stable (i.e., does not reorder equal elements).
605 ///
606 /// This uses [`slice::sort_by_key`] internally.
607 ///
608 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
609 /// This includes the allowed parameter type changes listed under [allowed transmutes].
610 /// However, the lens uses the filter of the original query when present.
611 ///
612 /// The sort is not cached across system runs.
613 ///
614 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
615 ///
616 /// # Panics
617 ///
618 /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
619 ///
620 /// # Example
621 /// ```
622 /// # use bevy_ecs::prelude::*;
623 /// # use std::ops::Deref;
624 /// #
625 /// # #[derive(Component)]
626 /// # struct PartMarker;
627 /// #
628 /// # impl Deref for PartValue {
629 /// # type Target = f32;
630 /// #
631 /// # fn deref(&self) -> &Self::Target {
632 /// # &self.0
633 /// # }
634 /// # }
635 /// #
636 /// # let mut world = World::new();
637 /// #
638 /// #[derive(Component)]
639 /// struct AvailableMarker;
640 ///
641 /// #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
642 /// enum Rarity {
643 /// Common,
644 /// Rare,
645 /// Epic,
646 /// Legendary
647 /// };
648 ///
649 /// #[derive(Component)]
650 /// struct PartValue(f32);
651 ///
652 /// // We can sort with the internals of components that do not implement Ord.
653 /// fn system_1(query: Query<(Entity, &PartValue)>) {
654 /// // Sort by the sines of the part values.
655 /// let parts: Vec<(Entity, &PartValue)> = query
656 /// .iter()
657 /// .sort_by_key::<&PartValue, _>(|value| value.sin() as usize)
658 /// .collect();
659 /// }
660 ///
661 /// // We can define our own custom comparison functions over an EntityRef.
662 /// fn system_2(query: Query<EntityRef, With<PartMarker>>) {
663 /// // Sort by whether parts are available and their rarity.
664 /// // We want the available legendaries to come first, so we reverse the iterator.
665 /// let parts: Vec<EntityRef> = query.iter()
666 /// .sort_by_key::<EntityRef, _>(|entity_ref| {
667 /// (
668 /// entity_ref.contains::<AvailableMarker>(),
669 /// entity_ref.get::<Rarity>()
670 /// )
671 /// })
672 /// .rev()
673 /// .collect();
674 /// }
675 /// # let mut schedule = Schedule::default();
676 /// # schedule.add_systems((system_1, system_2));
677 /// # schedule.run(&mut world);
678 /// ```
679 pub fn sort_by_key<L: ReadOnlyQueryData + 'w, K>(
680 self,
681 mut f: impl FnMut(&L::Item<'w>) -> K,
682 ) -> QuerySortedIter<
683 'w,
684 's,
685 D,
686 F,
687 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
688 >
689 where
690 K: Ord,
691 {
692 // On the first successful iteration of `QueryIterationCursor`, `archetype_entities` or `table_entities`
693 // will be set to a non-zero value. The correctness of this method relies on this.
694 // I.e. this sort method will execute if and only if `next` on `QueryIterationCursor` of a
695 // non-empty `QueryIter` has not yet been called. When empty, this sort method will not panic.
696 if !self.cursor.archetype_entities.is_empty() || !self.cursor.table_entities.is_empty() {
697 panic!("it is not valid to call sort() after next()")
698 }
699
700 let world = self.world;
701
702 let query_lens_state = self
703 .query_state
704 .transmute_filtered::<(L, Entity), F>(world.components());
705
706 // SAFETY:
707 // `self.world` has permission to access the required components.
708 // The original query iter has not been iterated on, so no items are aliased from it.
709 let query_lens = unsafe {
710 query_lens_state.iter_unchecked_manual(
711 world,
712 world.last_change_tick(),
713 world.change_tick(),
714 )
715 };
716 let mut keyed_query: Vec<_> = query_lens.collect();
717 keyed_query.sort_by_key(|(lens, _)| f(lens));
718 let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity);
719 // SAFETY:
720 // `self.world` has permission to access the required components.
721 // Each lens query item is dropped before the respective actual query item is accessed.
722 unsafe {
723 QuerySortedIter::new(
724 world,
725 self.query_state,
726 entity_iter,
727 world.last_change_tick(),
728 world.change_tick(),
729 )
730 }
731 }
732
733 /// Sorts all query items into a new iterator with a key extraction function over the query lens.
734 ///
735 /// This sort is unstable (i.e., may reorder equal elements).
736 ///
737 /// This uses [`slice::sort_unstable_by_key`] internally.
738 ///
739 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
740 /// This includes the allowed parameter type changes listed under [allowed transmutes].
741 /// However, the lens uses the filter of the original query when present.
742 ///
743 /// The sort is not cached across system runs.
744 ///
745 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
746 ///
747 /// # Panics
748 ///
749 /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
750 pub fn sort_unstable_by_key<L: ReadOnlyQueryData + 'w, K>(
751 self,
752 mut f: impl FnMut(&L::Item<'w>) -> K,
753 ) -> QuerySortedIter<
754 'w,
755 's,
756 D,
757 F,
758 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
759 >
760 where
761 K: Ord,
762 {
763 // On the first successful iteration of `QueryIterationCursor`, `archetype_entities` or `table_entities`
764 // will be set to a non-zero value. The correctness of this method relies on this.
765 // I.e. this sort method will execute if and only if `next` on `QueryIterationCursor` of a
766 // non-empty `QueryIter` has not yet been called. When empty, this sort method will not panic.
767 if !self.cursor.archetype_entities.is_empty() || !self.cursor.table_entities.is_empty() {
768 panic!("it is not valid to call sort() after next()")
769 }
770
771 let world = self.world;
772
773 let query_lens_state = self
774 .query_state
775 .transmute_filtered::<(L, Entity), F>(world.components());
776
777 // SAFETY:
778 // `self.world` has permission to access the required components.
779 // The original query iter has not been iterated on, so no items are aliased from it.
780 let query_lens = unsafe {
781 query_lens_state.iter_unchecked_manual(
782 world,
783 world.last_change_tick(),
784 world.change_tick(),
785 )
786 };
787 let mut keyed_query: Vec<_> = query_lens.collect();
788 keyed_query.sort_unstable_by_key(|(lens, _)| f(lens));
789 let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity);
790 // SAFETY:
791 // `self.world` has permission to access the required components.
792 // Each lens query item is dropped before the respective actual query item is accessed.
793 unsafe {
794 QuerySortedIter::new(
795 world,
796 self.query_state,
797 entity_iter,
798 world.last_change_tick(),
799 world.change_tick(),
800 )
801 }
802 }
803
804 /// Sort all query items into a new iterator with a key extraction function over the query lens.
805 ///
806 /// This sort is stable (i.e., does not reorder equal elements).
807 ///
808 /// This uses [`slice::sort_by_cached_key`] internally.
809 ///
810 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
811 /// This includes the allowed parameter type changes listed under [allowed transmutes].
812 /// However, the lens uses the filter of the original query when present.
813 ///
814 /// The sort is not cached across system runs.
815 ///
816 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
817 ///
818 /// # Panics
819 ///
820 /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
821 ///
822 pub fn sort_by_cached_key<L: ReadOnlyQueryData + 'w, K>(
823 self,
824 mut f: impl FnMut(&L::Item<'w>) -> K,
825 ) -> QuerySortedIter<
826 'w,
827 's,
828 D,
829 F,
830 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
831 >
832 where
833 K: Ord,
834 {
835 // On the first successful iteration of `QueryIterationCursor`, `archetype_entities` or `table_entities`
836 // will be set to a non-zero value. The correctness of this method relies on this.
837 // I.e. this sort method will execute if and only if `next` on `QueryIterationCursor` of a
838 // non-empty `QueryIter` has not yet been called. When empty, this sort method will not panic.
839 if !self.cursor.archetype_entities.is_empty() || !self.cursor.table_entities.is_empty() {
840 panic!("it is not valid to call sort() after next()")
841 }
842
843 let world = self.world;
844
845 let query_lens_state = self
846 .query_state
847 .transmute_filtered::<(L, Entity), F>(world.components());
848
849 // SAFETY:
850 // `self.world` has permission to access the required components.
851 // The original query iter has not been iterated on, so no items are aliased from it.
852 let query_lens = unsafe {
853 query_lens_state.iter_unchecked_manual(
854 world,
855 world.last_change_tick(),
856 world.change_tick(),
857 )
858 };
859 let mut keyed_query: Vec<_> = query_lens.collect();
860 keyed_query.sort_by_cached_key(|(lens, _)| f(lens));
861 let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity);
862 // SAFETY:
863 // `self.world` has permission to access the required components.
864 // Each lens query item is dropped before the respective actual query item is accessed.
865 unsafe {
866 QuerySortedIter::new(
867 world,
868 self.query_state,
869 entity_iter,
870 world.last_change_tick(),
871 world.change_tick(),
872 )
873 }
874 }
875}
876
877impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for QueryIter<'w, 's, D, F> {
878 type Item = D::Item<'w>;
879
880 #[inline(always)]
881 fn next(&mut self) -> Option<Self::Item> {
882 // SAFETY:
883 // `tables` and `archetypes` belong to the same world that the cursor was initialized for.
884 // `query_state` is the state that was passed to `QueryIterationCursor::init`.
885 unsafe {
886 self.cursor
887 .next(self.tables, self.archetypes, self.query_state)
888 }
889 }
890
891 fn size_hint(&self) -> (usize, Option<usize>) {
892 let max_size = self.cursor.max_remaining(self.tables, self.archetypes);
893 let archetype_query = F::IS_ARCHETYPAL;
894 let min_size = if archetype_query { max_size } else { 0 };
895 (min_size, Some(max_size))
896 }
897
898 #[inline]
899 fn fold<B, Func>(mut self, init: B, mut func: Func) -> B
900 where
901 Func: FnMut(B, Self::Item) -> B,
902 {
903 let mut accum = init;
904 // Empty any remaining uniterated values from the current table/archetype
905 while self.cursor.current_row != self.cursor.current_len {
906 let Some(item) = self.next() else { break };
907 accum = func(accum, item);
908 }
909 for id in self.cursor.storage_id_iter.clone() {
910 if D::IS_DENSE && F::IS_DENSE {
911 // SAFETY: Matched table IDs are guaranteed to still exist.
912 let table = unsafe { self.tables.get(id.table_id).debug_checked_unwrap() };
913 accum =
914 // SAFETY:
915 // - The fetched table matches both D and F
916 // - The provided range is equivalent to [0, table.entity_count)
917 // - The if block ensures that D::IS_DENSE and F::IS_DENSE are both true
918 unsafe { self.fold_over_table_range(accum, &mut func, table, 0..table.entity_count()) };
919 } else {
920 let archetype =
921 // SAFETY: Matched archetype IDs are guaranteed to still exist.
922 unsafe { self.archetypes.get(id.archetype_id).debug_checked_unwrap() };
923 accum =
924 // SAFETY:
925 // - The fetched archetype matches both D and F
926 // - The provided range is equivalent to [0, archetype.len)
927 // - The if block ensures that ether D::IS_DENSE or F::IS_DENSE are false
928 unsafe { self.fold_over_archetype_range(accum, &mut func, archetype, 0..archetype.len()) };
929 }
930 }
931 accum
932 }
933}
934
935// This is correct as [`QueryIter`] always returns `None` once exhausted.
936impl<'w, 's, D: QueryData, F: QueryFilter> FusedIterator for QueryIter<'w, 's, D, F> {}
937
938impl<'w, 's, D: QueryData, F: QueryFilter> Debug for QueryIter<'w, 's, D, F> {
939 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
940 f.debug_struct("QueryIter").finish()
941 }
942}
943
944/// An [`Iterator`] over sorted query results of a [`Query`](crate::system::Query).
945///
946/// This struct is created by the [`QueryIter::sort`], [`QueryIter::sort_unstable`],
947/// [`QueryIter::sort_by`], [`QueryIter::sort_unstable_by`], [`QueryIter::sort_by_key`],
948/// [`QueryIter::sort_unstable_by_key`], and [`QueryIter::sort_by_cached_key`] methods.
949pub struct QuerySortedIter<'w, 's, D: QueryData, F: QueryFilter, I>
950where
951 I: Iterator<Item = Entity>,
952{
953 entity_iter: I,
954 entities: &'w Entities,
955 tables: &'w Tables,
956 archetypes: &'w Archetypes,
957 fetch: D::Fetch<'w>,
958 query_state: &'s QueryState<D, F>,
959}
960
961impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> QuerySortedIter<'w, 's, D, F, I>
962where
963 I: Iterator<Item = Entity>,
964{
965 /// # Safety
966 /// - `world` must have permission to access any of the components registered in `query_state`.
967 /// - `world` must be the same one used to initialize `query_state`.
968 /// - `entity_list` must only contain unique entities or be empty.
969 pub(crate) unsafe fn new<EntityList: IntoIterator<IntoIter = I>>(
970 world: UnsafeWorldCell<'w>,
971 query_state: &'s QueryState<D, F>,
972 entity_list: EntityList,
973 last_run: Tick,
974 this_run: Tick,
975 ) -> QuerySortedIter<'w, 's, D, F, I> {
976 let fetch = D::init_fetch(world, &query_state.fetch_state, last_run, this_run);
977 QuerySortedIter {
978 query_state,
979 entities: world.entities(),
980 archetypes: world.archetypes(),
981 // SAFETY: We only access table data that has been registered in `query_state`.
982 // This means `world` has permission to access the data we use.
983 tables: &world.storages().tables,
984 fetch,
985 entity_iter: entity_list.into_iter(),
986 }
987 }
988
989 /// # Safety
990 /// `entity` must stem from `self.entity_iter`, and not have been passed before.
991 #[inline(always)]
992 unsafe fn fetch_next(&mut self, entity: Entity) -> D::Item<'w> {
993 let (location, archetype, table);
994 // SAFETY:
995 // `tables` and `archetypes` belong to the same world that the [`QueryIter`]
996 // was initialized for.
997 unsafe {
998 location = self.entities.get(entity).debug_checked_unwrap();
999 archetype = self
1000 .archetypes
1001 .get(location.archetype_id)
1002 .debug_checked_unwrap();
1003 table = self.tables.get(location.table_id).debug_checked_unwrap();
1004 }
1005
1006 // SAFETY: `archetype` is from the world that `fetch` was created for,
1007 // `fetch_state` is the state that `fetch` was initialized with
1008 unsafe {
1009 D::set_archetype(
1010 &mut self.fetch,
1011 &self.query_state.fetch_state,
1012 archetype,
1013 table,
1014 );
1015 }
1016
1017 // The entity list has already been filtered by the query lens, so we forego filtering here.
1018 // SAFETY:
1019 // - set_archetype was called prior, `location.archetype_row` is an archetype index in range of the current archetype
1020 // - fetch is only called once for each entity.
1021 unsafe { D::fetch(&mut self.fetch, entity, location.table_row) }
1022 }
1023}
1024
1025impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Iterator
1026 for QuerySortedIter<'w, 's, D, F, I>
1027where
1028 I: Iterator<Item = Entity>,
1029{
1030 type Item = D::Item<'w>;
1031
1032 #[inline(always)]
1033 fn next(&mut self) -> Option<Self::Item> {
1034 let entity = self.entity_iter.next()?;
1035 // SAFETY: `entity` is passed from `entity_iter` the first time.
1036 unsafe { self.fetch_next(entity).into() }
1037 }
1038
1039 fn size_hint(&self) -> (usize, Option<usize>) {
1040 self.entity_iter.size_hint()
1041 }
1042}
1043
1044impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> DoubleEndedIterator
1045 for QuerySortedIter<'w, 's, D, F, I>
1046where
1047 I: DoubleEndedIterator<Item = Entity>,
1048{
1049 #[inline(always)]
1050 fn next_back(&mut self) -> Option<Self::Item> {
1051 let entity = self.entity_iter.next_back()?;
1052 // SAFETY: `entity` is passed from `entity_iter` the first time.
1053 unsafe { self.fetch_next(entity).into() }
1054 }
1055}
1056
1057impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> ExactSizeIterator
1058 for QuerySortedIter<'w, 's, D, F, I>
1059where
1060 I: ExactSizeIterator<Item = Entity>,
1061{
1062}
1063
1064// This is correct as [`QuerySortedIter`] returns `None` once exhausted if `entity_iter` does.
1065impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> FusedIterator
1066 for QuerySortedIter<'w, 's, D, F, I>
1067where
1068 I: FusedIterator<Item = Entity>,
1069{
1070}
1071
1072impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item = Entity>> Debug
1073 for QuerySortedIter<'w, 's, D, F, I>
1074{
1075 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
1076 f.debug_struct("QuerySortedIter").finish()
1077 }
1078}
1079
1080/// An [`Iterator`] over the query items generated from an iterator of [`Entity`]s.
1081///
1082/// Items are returned in the order of the provided iterator.
1083/// Entities that don't match the query are skipped.
1084///
1085/// This struct is created by the [`Query::iter_many`](crate::system::Query::iter_many) and [`Query::iter_many_mut`](crate::system::Query::iter_many_mut) methods.
1086pub struct QueryManyIter<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>
1087where
1088 I::Item: Borrow<Entity>,
1089{
1090 entity_iter: I,
1091 entities: &'w Entities,
1092 tables: &'w Tables,
1093 archetypes: &'w Archetypes,
1094 fetch: D::Fetch<'w>,
1095 filter: F::Fetch<'w>,
1096 query_state: &'s QueryState<D, F>,
1097}
1098
1099impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> QueryManyIter<'w, 's, D, F, I>
1100where
1101 I::Item: Borrow<Entity>,
1102{
1103 /// # Safety
1104 /// - `world` must have permission to access any of the components registered in `query_state`.
1105 /// - `world` must be the same one used to initialize `query_state`.
1106 pub(crate) unsafe fn new<EntityList: IntoIterator<IntoIter = I>>(
1107 world: UnsafeWorldCell<'w>,
1108 query_state: &'s QueryState<D, F>,
1109 entity_list: EntityList,
1110 last_run: Tick,
1111 this_run: Tick,
1112 ) -> QueryManyIter<'w, 's, D, F, I> {
1113 let fetch = D::init_fetch(world, &query_state.fetch_state, last_run, this_run);
1114 let filter = F::init_fetch(world, &query_state.filter_state, last_run, this_run);
1115 QueryManyIter {
1116 query_state,
1117 entities: world.entities(),
1118 archetypes: world.archetypes(),
1119 // SAFETY: We only access table data that has been registered in `query_state`.
1120 // This means `world` has permission to access the data we use.
1121 tables: &world.storages().tables,
1122 fetch,
1123 filter,
1124 entity_iter: entity_list.into_iter(),
1125 }
1126 }
1127
1128 /// Safety:
1129 /// The lifetime here is not restrictive enough for Fetch with &mut access,
1130 /// as calling `fetch_next_aliased_unchecked` multiple times can produce multiple
1131 /// references to the same component, leading to unique reference aliasing.
1132 ///
1133 /// It is always safe for shared access.
1134 #[inline(always)]
1135 unsafe fn fetch_next_aliased_unchecked(&mut self) -> Option<D::Item<'w>> {
1136 for entity in self.entity_iter.by_ref() {
1137 let entity = *entity.borrow();
1138 let Some(location) = self.entities.get(entity) else {
1139 continue;
1140 };
1141
1142 if !self
1143 .query_state
1144 .matched_archetypes
1145 .contains(location.archetype_id.index())
1146 {
1147 continue;
1148 }
1149
1150 let archetype = self
1151 .archetypes
1152 .get(location.archetype_id)
1153 .debug_checked_unwrap();
1154 let table = self.tables.get(location.table_id).debug_checked_unwrap();
1155
1156 // SAFETY: `archetype` is from the world that `fetch/filter` were created for,
1157 // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
1158 unsafe {
1159 D::set_archetype(
1160 &mut self.fetch,
1161 &self.query_state.fetch_state,
1162 archetype,
1163 table,
1164 );
1165 }
1166 // SAFETY: `table` is from the world that `fetch/filter` were created for,
1167 // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
1168 unsafe {
1169 F::set_archetype(
1170 &mut self.filter,
1171 &self.query_state.filter_state,
1172 archetype,
1173 table,
1174 );
1175 }
1176
1177 // SAFETY: set_archetype was called prior.
1178 // `location.archetype_row` is an archetype index row in range of the current archetype, because if it was not, the match above would have `continue`d
1179 if unsafe { F::filter_fetch(&mut self.filter, entity, location.table_row) } {
1180 // SAFETY:
1181 // - set_archetype was called prior, `location.archetype_row` is an archetype index in range of the current archetype
1182 // - fetch is only called once for each entity.
1183 return Some(unsafe { D::fetch(&mut self.fetch, entity, location.table_row) });
1184 }
1185 }
1186 None
1187 }
1188
1189 /// Get next result from the query
1190 #[inline(always)]
1191 pub fn fetch_next(&mut self) -> Option<D::Item<'_>> {
1192 // SAFETY: we are limiting the returned reference to self,
1193 // making sure this method cannot be called multiple times without getting rid
1194 // of any previously returned unique references first, thus preventing aliasing.
1195 unsafe { self.fetch_next_aliased_unchecked().map(D::shrink) }
1196 }
1197}
1198
1199impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator> Iterator
1200 for QueryManyIter<'w, 's, D, F, I>
1201where
1202 I::Item: Borrow<Entity>,
1203{
1204 type Item = D::Item<'w>;
1205
1206 #[inline(always)]
1207 fn next(&mut self) -> Option<Self::Item> {
1208 // SAFETY: It is safe to alias for ReadOnlyWorldQuery.
1209 unsafe { self.fetch_next_aliased_unchecked() }
1210 }
1211
1212 fn size_hint(&self) -> (usize, Option<usize>) {
1213 let (_, max_size) = self.entity_iter.size_hint();
1214 (0, max_size)
1215 }
1216}
1217
1218// This is correct as [`QueryManyIter`] always returns `None` once exhausted.
1219impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator> FusedIterator
1220 for QueryManyIter<'w, 's, D, F, I>
1221where
1222 I::Item: Borrow<Entity>,
1223{
1224}
1225
1226impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Debug for QueryManyIter<'w, 's, D, F, I>
1227where
1228 I::Item: Borrow<Entity>,
1229{
1230 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
1231 f.debug_struct("QueryManyIter").finish()
1232 }
1233}
1234
1235/// An iterator over `K`-sized combinations of query items without repetition.
1236///
1237/// A combination is an arrangement of a collection of items where order does not matter.
1238///
1239/// `K` is the number of items that make up each subset, and the number of items returned by the iterator.
1240/// `N` is the number of total entities output by the query.
1241///
1242/// For example, given the list [1, 2, 3, 4], where `K` is 2, the combinations without repeats are
1243/// [1, 2], [1, 3], [1, 4], [2, 3], [2, 4], [3, 4].
1244/// And in this case, `N` would be defined as 4 since the size of the input list is 4.
1245///
1246/// The number of combinations depend on how `K` relates to the number of entities matching the [`Query`]:
1247/// - if `K = N`, only one combination exists,
1248/// - if `K < N`, there are <sub>N</sub>C<sub>K</sub> combinations (see the [performance section] of `Query`),
1249/// - if `K > N`, there are no combinations.
1250///
1251/// The output combination is not guaranteed to have any order of iteration.
1252///
1253/// # Usage
1254///
1255/// This type is returned by calling [`Query::iter_combinations`] or [`Query::iter_combinations_mut`].
1256///
1257/// It implements [`Iterator`] only if it iterates over read-only query items ([learn more]).
1258///
1259/// In the case of mutable query items, it can be iterated by calling [`fetch_next`] in a `while let` loop.
1260///
1261/// # Examples
1262///
1263/// The following example shows how to traverse the iterator when the query items are read-only.
1264///
1265/// ```
1266/// # use bevy_ecs::prelude::*;
1267/// # #[derive(Component)]
1268/// # struct ComponentA;
1269/// #
1270/// fn some_system(query: Query<&ComponentA>) {
1271/// for [a1, a2] in query.iter_combinations() {
1272/// // ...
1273/// }
1274/// }
1275/// ```
1276///
1277/// The following example shows how `fetch_next` should be called with a `while let` loop to traverse the iterator when the query items are mutable.
1278///
1279/// ```
1280/// # use bevy_ecs::prelude::*;
1281/// # #[derive(Component)]
1282/// # struct ComponentA;
1283/// #
1284/// fn some_system(mut query: Query<&mut ComponentA>) {
1285/// let mut combinations = query.iter_combinations_mut();
1286/// while let Some([a1, a2]) = combinations.fetch_next() {
1287/// // ...
1288/// }
1289/// }
1290/// ```
1291///
1292/// [`fetch_next`]: Self::fetch_next
1293/// [learn more]: Self#impl-Iterator
1294/// [performance section]: crate::system::Query#performance
1295/// [`Query`]: crate::system::Query
1296/// [`Query::iter_combinations`]: crate::system::Query::iter_combinations
1297/// [`Query::iter_combinations_mut`]: crate::system::Query::iter_combinations_mut
1298pub struct QueryCombinationIter<'w, 's, D: QueryData, F: QueryFilter, const K: usize> {
1299 tables: &'w Tables,
1300 archetypes: &'w Archetypes,
1301 query_state: &'s QueryState<D, F>,
1302 cursors: [QueryIterationCursor<'w, 's, D, F>; K],
1303}
1304
1305impl<'w, 's, D: QueryData, F: QueryFilter, const K: usize> QueryCombinationIter<'w, 's, D, F, K> {
1306 /// # Safety
1307 /// - `world` must have permission to access any of the components registered in `query_state`.
1308 /// - `world` must be the same one used to initialize `query_state`.
1309 pub(crate) unsafe fn new(
1310 world: UnsafeWorldCell<'w>,
1311 query_state: &'s QueryState<D, F>,
1312 last_run: Tick,
1313 this_run: Tick,
1314 ) -> Self {
1315 // Initialize array with cursors.
1316 // There is no FromIterator on arrays, so instead initialize it manually with MaybeUninit
1317
1318 let mut array: MaybeUninit<[QueryIterationCursor<'w, 's, D, F>; K]> = MaybeUninit::uninit();
1319 let ptr = array
1320 .as_mut_ptr()
1321 .cast::<QueryIterationCursor<'w, 's, D, F>>();
1322 if K != 0 {
1323 ptr.write(QueryIterationCursor::init(
1324 world,
1325 query_state,
1326 last_run,
1327 this_run,
1328 ));
1329 }
1330 for slot in (1..K).map(|offset| ptr.add(offset)) {
1331 slot.write(QueryIterationCursor::init_empty(
1332 world,
1333 query_state,
1334 last_run,
1335 this_run,
1336 ));
1337 }
1338
1339 QueryCombinationIter {
1340 query_state,
1341 // SAFETY: We only access table data that has been registered in `query_state`.
1342 tables: unsafe { &world.storages().tables },
1343 archetypes: world.archetypes(),
1344 cursors: array.assume_init(),
1345 }
1346 }
1347
1348 /// Safety:
1349 /// The lifetime here is not restrictive enough for Fetch with &mut access,
1350 /// as calling `fetch_next_aliased_unchecked` multiple times can produce multiple
1351 /// references to the same component, leading to unique reference aliasing.
1352 ///.
1353 /// It is always safe for shared access.
1354 unsafe fn fetch_next_aliased_unchecked(&mut self) -> Option<[D::Item<'w>; K]> {
1355 if K == 0 {
1356 return None;
1357 }
1358
1359 // PERF: can speed up the following code using `cursor.remaining()` instead of `next_item.is_none()`
1360 // when D::IS_ARCHETYPAL && F::IS_ARCHETYPAL
1361 //
1362 // let `i` be the index of `c`, the last cursor in `self.cursors` that
1363 // returns `K-i` or more elements.
1364 // Make cursor in index `j` for all `j` in `[i, K)` a copy of `c` advanced `j-i+1` times.
1365 // If no such `c` exists, return `None`
1366 'outer: for i in (0..K).rev() {
1367 match self.cursors[i].next(self.tables, self.archetypes, self.query_state) {
1368 Some(_) => {
1369 for j in (i + 1)..K {
1370 self.cursors[j] = self.cursors[j - 1].clone();
1371 match self.cursors[j].next(self.tables, self.archetypes, self.query_state) {
1372 Some(_) => {}
1373 None if i > 0 => continue 'outer,
1374 None => return None,
1375 }
1376 }
1377 break;
1378 }
1379 None if i > 0 => continue,
1380 None => return None,
1381 }
1382 }
1383
1384 let mut values = MaybeUninit::<[D::Item<'w>; K]>::uninit();
1385
1386 let ptr = values.as_mut_ptr().cast::<D::Item<'w>>();
1387 for (offset, cursor) in self.cursors.iter_mut().enumerate() {
1388 ptr.add(offset).write(cursor.peek_last().unwrap());
1389 }
1390
1391 Some(values.assume_init())
1392 }
1393
1394 /// Get next combination of queried components
1395 #[inline]
1396 pub fn fetch_next(&mut self) -> Option<[D::Item<'_>; K]> {
1397 // SAFETY: we are limiting the returned reference to self,
1398 // making sure this method cannot be called multiple times without getting rid
1399 // of any previously returned unique references first, thus preventing aliasing.
1400 unsafe {
1401 self.fetch_next_aliased_unchecked()
1402 .map(|array| array.map(D::shrink))
1403 }
1404 }
1405}
1406
1407// Iterator type is intentionally implemented only for read-only access.
1408// Doing so for mutable references would be unsound, because calling `next`
1409// multiple times would allow multiple owned references to the same data to exist.
1410impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, const K: usize> Iterator
1411 for QueryCombinationIter<'w, 's, D, F, K>
1412{
1413 type Item = [D::Item<'w>; K];
1414
1415 #[inline]
1416 fn next(&mut self) -> Option<Self::Item> {
1417 // Safety: it is safe to alias for ReadOnlyWorldQuery
1418 unsafe { QueryCombinationIter::fetch_next_aliased_unchecked(self) }
1419 }
1420
1421 fn size_hint(&self) -> (usize, Option<usize>) {
1422 // binomial coefficient: (n ; k) = n! / k!(n-k)! = (n*n-1*...*n-k+1) / k!
1423 // See https://en.wikipedia.org/wiki/Binomial_coefficient
1424 // See https://blog.plover.com/math/choose.html for implementation
1425 // It was chosen to reduce overflow potential.
1426 fn choose(n: usize, k: usize) -> Option<usize> {
1427 if k > n || n == 0 {
1428 return Some(0);
1429 }
1430 let k = k.min(n - k);
1431 let ks = 1..=k;
1432 let ns = (n - k + 1..=n).rev();
1433 ks.zip(ns)
1434 .try_fold(1_usize, |acc, (k, n)| Some(acc.checked_mul(n)? / k))
1435 }
1436 // sum_i=0..k choose(cursors[i].remaining, k-i)
1437 let max_combinations = self
1438 .cursors
1439 .iter()
1440 .enumerate()
1441 .try_fold(0, |acc, (i, cursor)| {
1442 let n = cursor.max_remaining(self.tables, self.archetypes);
1443 Some(acc + choose(n, K - i)?)
1444 });
1445
1446 let archetype_query = F::IS_ARCHETYPAL;
1447 let known_max = max_combinations.unwrap_or(usize::MAX);
1448 let min_combinations = if archetype_query { known_max } else { 0 };
1449 (min_combinations, max_combinations)
1450 }
1451}
1452
1453impl<'w, 's, D: QueryData, F: QueryFilter> ExactSizeIterator for QueryIter<'w, 's, D, F>
1454where
1455 F: ArchetypeFilter,
1456{
1457 fn len(&self) -> usize {
1458 self.size_hint().0
1459 }
1460}
1461
1462// This is correct as [`QueryCombinationIter`] always returns `None` once exhausted.
1463impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, const K: usize> FusedIterator
1464 for QueryCombinationIter<'w, 's, D, F, K>
1465{
1466}
1467
1468impl<'w, 's, D: QueryData, F: QueryFilter, const K: usize> Debug
1469 for QueryCombinationIter<'w, 's, D, F, K>
1470{
1471 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
1472 f.debug_struct("QueryCombinationIter").finish()
1473 }
1474}
1475
1476struct QueryIterationCursor<'w, 's, D: QueryData, F: QueryFilter> {
1477 storage_id_iter: std::slice::Iter<'s, StorageId>,
1478 table_entities: &'w [Entity],
1479 archetype_entities: &'w [ArchetypeEntity],
1480 fetch: D::Fetch<'w>,
1481 filter: F::Fetch<'w>,
1482 // length of the table or length of the archetype, depending on whether both `D`'s and `F`'s fetches are dense
1483 current_len: usize,
1484 // either table row or archetype index, depending on whether both `D`'s and `F`'s fetches are dense
1485 current_row: usize,
1486}
1487
1488impl<D: QueryData, F: QueryFilter> Clone for QueryIterationCursor<'_, '_, D, F> {
1489 fn clone(&self) -> Self {
1490 Self {
1491 storage_id_iter: self.storage_id_iter.clone(),
1492 table_entities: self.table_entities,
1493 archetype_entities: self.archetype_entities,
1494 fetch: self.fetch.clone(),
1495 filter: self.filter.clone(),
1496 current_len: self.current_len,
1497 current_row: self.current_row,
1498 }
1499 }
1500}
1501
1502impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> {
1503 const IS_DENSE: bool = D::IS_DENSE && F::IS_DENSE;
1504
1505 unsafe fn init_empty(
1506 world: UnsafeWorldCell<'w>,
1507 query_state: &'s QueryState<D, F>,
1508 last_run: Tick,
1509 this_run: Tick,
1510 ) -> Self {
1511 QueryIterationCursor {
1512 storage_id_iter: [].iter(),
1513 ..Self::init(world, query_state, last_run, this_run)
1514 }
1515 }
1516
1517 /// # Safety
1518 /// - `world` must have permission to access any of the components registered in `query_state`.
1519 /// - `world` must be the same one used to initialize `query_state`.
1520 unsafe fn init(
1521 world: UnsafeWorldCell<'w>,
1522 query_state: &'s QueryState<D, F>,
1523 last_run: Tick,
1524 this_run: Tick,
1525 ) -> Self {
1526 let fetch = D::init_fetch(world, &query_state.fetch_state, last_run, this_run);
1527 let filter = F::init_fetch(world, &query_state.filter_state, last_run, this_run);
1528 QueryIterationCursor {
1529 fetch,
1530 filter,
1531 table_entities: &[],
1532 archetype_entities: &[],
1533 storage_id_iter: query_state.matched_storage_ids.iter(),
1534 current_len: 0,
1535 current_row: 0,
1536 }
1537 }
1538
1539 /// retrieve item returned from most recent `next` call again.
1540 #[inline]
1541 unsafe fn peek_last(&mut self) -> Option<D::Item<'w>> {
1542 if self.current_row > 0 {
1543 let index = self.current_row - 1;
1544 if Self::IS_DENSE {
1545 let entity = self.table_entities.get_unchecked(index);
1546 Some(D::fetch(
1547 &mut self.fetch,
1548 *entity,
1549 TableRow::from_usize(index),
1550 ))
1551 } else {
1552 let archetype_entity = self.archetype_entities.get_unchecked(index);
1553 Some(D::fetch(
1554 &mut self.fetch,
1555 archetype_entity.id(),
1556 archetype_entity.table_row(),
1557 ))
1558 }
1559 } else {
1560 None
1561 }
1562 }
1563
1564 /// How many values will this cursor return at most?
1565 ///
1566 /// Note that if `F::IS_ARCHETYPAL`, the return value
1567 /// will be **the exact count of remaining values**.
1568 fn max_remaining(&self, tables: &'w Tables, archetypes: &'w Archetypes) -> usize {
1569 let ids = self.storage_id_iter.clone();
1570 let remaining_matched: usize = if Self::IS_DENSE {
1571 // SAFETY: The if check ensures that storage_id_iter stores TableIds
1572 unsafe { ids.map(|id| tables[id.table_id].entity_count()).sum() }
1573 } else {
1574 // SAFETY: The if check ensures that storage_id_iter stores ArchetypeIds
1575 unsafe { ids.map(|id| archetypes[id.archetype_id].len()).sum() }
1576 };
1577 remaining_matched + self.current_len - self.current_row
1578 }
1579
1580 // NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
1581 // QueryIter, QueryIterationCursor, QuerySortedIter, QueryManyIter, QueryCombinationIter, QueryState::par_fold_init_unchecked_manual
1582 /// # Safety
1583 /// `tables` and `archetypes` must belong to the same world that the [`QueryIterationCursor`]
1584 /// was initialized for.
1585 /// `query_state` must be the same [`QueryState`] that was passed to `init` or `init_empty`.
1586 #[inline(always)]
1587 unsafe fn next(
1588 &mut self,
1589 tables: &'w Tables,
1590 archetypes: &'w Archetypes,
1591 query_state: &'s QueryState<D, F>,
1592 ) -> Option<D::Item<'w>> {
1593 if Self::IS_DENSE {
1594 loop {
1595 // we are on the beginning of the query, or finished processing a table, so skip to the next
1596 if self.current_row == self.current_len {
1597 let table_id = self.storage_id_iter.next()?.table_id;
1598 let table = tables.get(table_id).debug_checked_unwrap();
1599 // SAFETY: `table` is from the world that `fetch/filter` were created for,
1600 // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
1601 unsafe {
1602 D::set_table(&mut self.fetch, &query_state.fetch_state, table);
1603 F::set_table(&mut self.filter, &query_state.filter_state, table);
1604 }
1605 self.table_entities = table.entities();
1606 self.current_len = table.entity_count();
1607 self.current_row = 0;
1608 continue;
1609 }
1610
1611 // SAFETY: set_table was called prior.
1612 // `current_row` is a table row in range of the current table, because if it was not, then the above would have been executed.
1613 let entity = unsafe { self.table_entities.get_unchecked(self.current_row) };
1614 let row = TableRow::from_usize(self.current_row);
1615 if !F::filter_fetch(&mut self.filter, *entity, row) {
1616 self.current_row += 1;
1617 continue;
1618 }
1619
1620 // SAFETY:
1621 // - set_table was called prior.
1622 // - `current_row` must be a table row in range of the current table,
1623 // because if it was not, then the above would have been executed.
1624 // - fetch is only called once for each `entity`.
1625 let item = unsafe { D::fetch(&mut self.fetch, *entity, row) };
1626
1627 self.current_row += 1;
1628 return Some(item);
1629 }
1630 } else {
1631 loop {
1632 if self.current_row == self.current_len {
1633 let archetype_id = self.storage_id_iter.next()?.archetype_id;
1634 let archetype = archetypes.get(archetype_id).debug_checked_unwrap();
1635 let table = tables.get(archetype.table_id()).debug_checked_unwrap();
1636 // SAFETY: `archetype` and `tables` are from the world that `fetch/filter` were created for,
1637 // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
1638 unsafe {
1639 D::set_archetype(
1640 &mut self.fetch,
1641 &query_state.fetch_state,
1642 archetype,
1643 table,
1644 );
1645 F::set_archetype(
1646 &mut self.filter,
1647 &query_state.filter_state,
1648 archetype,
1649 table,
1650 );
1651 }
1652 self.archetype_entities = archetype.entities();
1653 self.current_len = archetype.len();
1654 self.current_row = 0;
1655 continue;
1656 }
1657
1658 // SAFETY: set_archetype was called prior.
1659 // `current_row` is an archetype index row in range of the current archetype, because if it was not, then the if above would have been executed.
1660 let archetype_entity =
1661 unsafe { self.archetype_entities.get_unchecked(self.current_row) };
1662 if !F::filter_fetch(
1663 &mut self.filter,
1664 archetype_entity.id(),
1665 archetype_entity.table_row(),
1666 ) {
1667 self.current_row += 1;
1668 continue;
1669 }
1670
1671 // SAFETY:
1672 // - set_archetype was called prior.
1673 // - `current_row` must be an archetype index row in range of the current archetype,
1674 // because if it was not, then the if above would have been executed.
1675 // - fetch is only called once for each `archetype_entity`.
1676 let item = unsafe {
1677 D::fetch(
1678 &mut self.fetch,
1679 archetype_entity.id(),
1680 archetype_entity.table_row(),
1681 )
1682 };
1683 self.current_row += 1;
1684 return Some(item);
1685 }
1686 }
1687 }
1688}
1689
1690// A wrapper struct that gives its data a neutral ordering.
1691#[derive(Copy, Clone)]
1692struct NeutralOrd<T>(T);
1693
1694impl<T> PartialEq for NeutralOrd<T> {
1695 fn eq(&self, _other: &Self) -> bool {
1696 true
1697 }
1698}
1699
1700impl<T> Eq for NeutralOrd<T> {}
1701
1702#[allow(clippy::non_canonical_partial_ord_impl)]
1703impl<T> PartialOrd for NeutralOrd<T> {
1704 fn partial_cmp(&self, _other: &Self) -> Option<Ordering> {
1705 Some(Ordering::Equal)
1706 }
1707}
1708
1709impl<T> Ord for NeutralOrd<T> {
1710 fn cmp(&self, _other: &Self) -> Ordering {
1711 Ordering::Equal
1712 }
1713}
1714
1715#[cfg(test)]
1716mod tests {
1717 #[allow(unused_imports)]
1718 use crate::{self as bevy_ecs, component::Component, entity::Entity, prelude::World};
1719
1720 #[derive(Component, Debug, PartialEq, PartialOrd, Clone, Copy)]
1721 struct A(f32);
1722 #[derive(Component, Debug, Eq, PartialEq, Clone, Copy)]
1723 #[component(storage = "SparseSet")]
1724 struct Sparse(usize);
1725
1726 #[allow(clippy::unnecessary_sort_by)]
1727 #[test]
1728 fn query_sorts() {
1729 let mut world = World::new();
1730
1731 let mut query = world.query::<Entity>();
1732
1733 let sort = query.iter(&world).sort::<Entity>().collect::<Vec<_>>();
1734
1735 let sort_unstable = query
1736 .iter(&world)
1737 .sort_unstable::<Entity>()
1738 .collect::<Vec<_>>();
1739
1740 let sort_by = query
1741 .iter(&world)
1742 .sort_by::<Entity>(|e1, e2| e1.cmp(e2))
1743 .collect::<Vec<_>>();
1744
1745 let sort_unstable_by = query
1746 .iter(&world)
1747 .sort_unstable_by::<Entity>(|e1, e2| e1.cmp(e2))
1748 .collect::<Vec<_>>();
1749
1750 let sort_by_key = query
1751 .iter(&world)
1752 .sort_by_key::<Entity, _>(|&e| e)
1753 .collect::<Vec<_>>();
1754
1755 let sort_unstable_by_key = query
1756 .iter(&world)
1757 .sort_unstable_by_key::<Entity, _>(|&e| e)
1758 .collect::<Vec<_>>();
1759
1760 let sort_by_cached_key = query
1761 .iter(&world)
1762 .sort_by_cached_key::<Entity, _>(|&e| e)
1763 .collect::<Vec<_>>();
1764
1765 let mut sort_v2 = query.iter(&world).collect::<Vec<_>>();
1766 sort_v2.sort();
1767
1768 let mut sort_unstable_v2 = query.iter(&world).collect::<Vec<_>>();
1769 sort_unstable_v2.sort_unstable();
1770
1771 let mut sort_by_v2 = query.iter(&world).collect::<Vec<_>>();
1772 sort_by_v2.sort_by(|e1, e2| e1.cmp(e2));
1773
1774 let mut sort_unstable_by_v2 = query.iter(&world).collect::<Vec<_>>();
1775 sort_unstable_by_v2.sort_unstable_by(|e1, e2| e1.cmp(e2));
1776
1777 let mut sort_by_key_v2 = query.iter(&world).collect::<Vec<_>>();
1778 sort_by_key_v2.sort_by_key(|&e| e);
1779
1780 let mut sort_unstable_by_key_v2 = query.iter(&world).collect::<Vec<_>>();
1781 sort_unstable_by_key_v2.sort_unstable_by_key(|&e| e);
1782
1783 let mut sort_by_cached_key_v2 = query.iter(&world).collect::<Vec<_>>();
1784 sort_by_cached_key_v2.sort_by_cached_key(|&e| e);
1785
1786 assert_eq!(sort, sort_v2);
1787 assert_eq!(sort_unstable, sort_unstable_v2);
1788 assert_eq!(sort_by, sort_by_v2);
1789 assert_eq!(sort_unstable_by, sort_unstable_by_v2);
1790 assert_eq!(sort_by_key, sort_by_key_v2);
1791 assert_eq!(sort_unstable_by_key, sort_unstable_by_key_v2);
1792 assert_eq!(sort_by_cached_key, sort_by_cached_key_v2);
1793 }
1794
1795 #[test]
1796 #[should_panic]
1797 fn query_sort_after_next() {
1798 let mut world = World::new();
1799 world.spawn((A(0.),));
1800 world.spawn((A(1.1),));
1801 world.spawn((A(2.22),));
1802
1803 {
1804 let mut query = world.query::<&A>();
1805 let mut iter = query.iter(&world);
1806 println!(
1807 "archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
1808 iter.cursor.archetype_entities.len(),
1809 iter.cursor.table_entities.len(),
1810 iter.cursor.current_len,
1811 iter.cursor.current_row
1812 );
1813 _ = iter.next();
1814 println!(
1815 "archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
1816 iter.cursor.archetype_entities.len(),
1817 iter.cursor.table_entities.len(),
1818 iter.cursor.current_len,
1819 iter.cursor.current_row
1820 );
1821 println!("{}", iter.sort::<Entity>().len());
1822 }
1823 }
1824
1825 #[test]
1826 #[should_panic]
1827 fn query_sort_after_next_dense() {
1828 let mut world = World::new();
1829 world.spawn((Sparse(11),));
1830 world.spawn((Sparse(22),));
1831 world.spawn((Sparse(33),));
1832
1833 {
1834 let mut query = world.query::<&Sparse>();
1835 let mut iter = query.iter(&world);
1836 println!(
1837 "before_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
1838 iter.cursor.archetype_entities.len(),
1839 iter.cursor.table_entities.len(),
1840 iter.cursor.current_len,
1841 iter.cursor.current_row
1842 );
1843 _ = iter.next();
1844 println!(
1845 "after_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
1846 iter.cursor.archetype_entities.len(),
1847 iter.cursor.table_entities.len(),
1848 iter.cursor.current_len,
1849 iter.cursor.current_row
1850 );
1851 println!("{}", iter.sort::<Entity>().len());
1852 }
1853 }
1854
1855 #[test]
1856 fn empty_query_sort_after_next_does_not_panic() {
1857 let mut world = World::new();
1858 {
1859 let mut query = world.query::<(&A, &Sparse)>();
1860 let mut iter = query.iter(&world);
1861 println!(
1862 "before_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
1863 iter.cursor.archetype_entities.len(),
1864 iter.cursor.table_entities.len(),
1865 iter.cursor.current_len,
1866 iter.cursor.current_row
1867 );
1868 _ = iter.next();
1869 println!(
1870 "after_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
1871 iter.cursor.archetype_entities.len(),
1872 iter.cursor.table_entities.len(),
1873 iter.cursor.current_len,
1874 iter.cursor.current_row
1875 );
1876 println!("{}", iter.sort::<Entity>().len());
1877 }
1878 }
1879
1880 #[test]
1881 fn query_iter_cursor_state_non_empty_after_next() {
1882 let mut world = World::new();
1883 world.spawn((A(0.), Sparse(11)));
1884 world.spawn((A(1.1), Sparse(22)));
1885 world.spawn((A(2.22), Sparse(33)));
1886 {
1887 let mut query = world.query::<(&A, &Sparse)>();
1888 let mut iter = query.iter(&world);
1889 println!(
1890 "before_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
1891 iter.cursor.archetype_entities.len(),
1892 iter.cursor.table_entities.len(),
1893 iter.cursor.current_len,
1894 iter.cursor.current_row
1895 );
1896 assert!(iter.cursor.table_entities.len() | iter.cursor.archetype_entities.len() == 0);
1897 _ = iter.next();
1898 println!(
1899 "after_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
1900 iter.cursor.archetype_entities.len(),
1901 iter.cursor.table_entities.len(),
1902 iter.cursor.current_len,
1903 iter.cursor.current_row
1904 );
1905 assert!(
1906 (
1907 iter.cursor.table_entities.len(),
1908 iter.cursor.archetype_entities.len()
1909 ) != (0, 0)
1910 );
1911 }
1912 }
1913}