bevy_ecs/query/iter.rs
1use super::{QueryData, QueryFilter, ReadOnlyQueryData};
2use crate::{
3 archetype::{Archetype, ArchetypeEntity, Archetypes},
4 bundle::Bundle,
5 change_detection::Tick,
6 entity::{ContainsEntity, Entities, Entity, EntityEquivalent, EntitySet, EntitySetIterator},
7 query::{
8 ArchetypeFilter, ArchetypeQueryData, ContiguousQueryData, DebugCheckedUnwrap,
9 IterQueryData, QueryState, SingleEntityQueryData, StorageId,
10 },
11 storage::{Table, TableRow, Tables},
12 world::{
13 unsafe_world_cell::UnsafeWorldCell, EntityMut, EntityMutExcept, EntityRef, EntityRefExcept,
14 FilteredEntityMut, FilteredEntityRef,
15 },
16};
17use alloc::vec::Vec;
18use core::{
19 cmp::Ordering,
20 fmt::{self, Debug, Formatter},
21 iter::FusedIterator,
22 mem::MaybeUninit,
23 ops::Range,
24};
25use nonmax::NonMaxU32;
26
27/// An [`Iterator`] over query results of a [`Query`](crate::system::Query).
28///
29/// This struct is created by the [`Query::iter`](crate::system::Query::iter) and
30/// [`Query::iter_mut`](crate::system::Query::iter_mut) methods.
31pub struct QueryIter<'w, 's, D: QueryData, F: QueryFilter> {
32 world: UnsafeWorldCell<'w>,
33 tables: &'w Tables,
34 archetypes: &'w Archetypes,
35 query_state: &'s QueryState<D, F>,
36 cursor: QueryIterationCursor<'w, 's, D, F>,
37}
38
39impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
40 /// # Safety
41 /// - `world` must have permission to access any of the components registered in `query_state`.
42 /// - `world` must be the same one used to initialize `query_state`.
43 pub(crate) unsafe fn new(
44 world: UnsafeWorldCell<'w>,
45 query_state: &'s QueryState<D, F>,
46 last_run: Tick,
47 this_run: Tick,
48 ) -> Self {
49 QueryIter {
50 world,
51 query_state,
52 // SAFETY: We only access table data that has been registered in `query_state`.
53 tables: unsafe { &world.storages().tables },
54 archetypes: world.archetypes(),
55 // SAFETY: The invariants are upheld by the caller.
56 cursor: unsafe { QueryIterationCursor::init(world, query_state, last_run, this_run) },
57 }
58 }
59
60 /// Creates a new separate iterator yielding the same remaining items of the current one.
61 /// Advancing the new iterator will not advance the original one, which will resume at the
62 /// point it was left at.
63 ///
64 /// Differently from [`remaining_mut`](QueryIter::remaining_mut) the new iterator does not
65 /// borrow from the original one. However it can only be called from an iterator over read only
66 /// items.
67 ///
68 /// # Example
69 ///
70 /// ```
71 /// # use bevy_ecs::prelude::*;
72 /// #
73 /// # #[derive(Component)]
74 /// # struct ComponentA;
75 ///
76 /// fn combinations(query: Query<&ComponentA>) {
77 /// let mut iter = query.iter();
78 /// while let Some(a) = iter.next() {
79 /// for b in iter.remaining() {
80 /// // Check every combination (a, b)
81 /// }
82 /// }
83 /// }
84 /// ```
85 pub fn remaining(&self) -> QueryIter<'w, 's, D, F>
86 where
87 D: ReadOnlyQueryData,
88 {
89 QueryIter {
90 world: self.world,
91 tables: self.tables,
92 archetypes: self.archetypes,
93 query_state: self.query_state,
94 cursor: self.cursor.clone(),
95 }
96 }
97
98 /// Creates a new separate iterator yielding the same remaining items of the current one.
99 /// Advancing the new iterator will not advance the original one, which will resume at the
100 /// point it was left at.
101 ///
102 /// This method can be called on iterators over mutable items. However the original iterator
103 /// will be borrowed while the new iterator exists and will thus not be usable in that timespan.
104 ///
105 /// # Example
106 ///
107 /// ```
108 /// # use bevy_ecs::prelude::*;
109 /// #
110 /// # #[derive(Component)]
111 /// # struct ComponentA;
112 ///
113 /// fn combinations(mut query: Query<&mut ComponentA>) {
114 /// let mut iter = query.iter_mut();
115 /// while let Some(a) = iter.next() {
116 /// for b in iter.remaining_mut() {
117 /// // Check every combination (a, b)
118 /// }
119 /// }
120 /// }
121 /// ```
122 pub fn remaining_mut(&mut self) -> QueryIter<'_, 's, D, F> {
123 QueryIter {
124 world: self.world,
125 tables: self.tables,
126 archetypes: self.archetypes,
127 query_state: self.query_state,
128 cursor: self.cursor.reborrow(),
129 }
130 }
131
132 /// Get the next result from the query.
133 ///
134 /// If the [`QueryData`] does not implement [`IterQueryData`],
135 /// then it is not sound to yield multiple items concurrently
136 /// and the resulting [`QueryIter`] will not implement [`Iterator`].
137 /// In that case, this method can be used to iterate over the items
138 /// while ensuring only one is alive at a time.
139 ///
140 /// Most queries do implement [`IterQueryData`],
141 /// and can use the ordinary [`Iterator::next`]
142 /// method or a `for` loop.
143 ///
144 /// # Example
145 ///
146 /// ```
147 /// # use bevy_ecs::prelude::*;
148 /// # #[derive(Component)]
149 /// # struct C;
150 /// fn system(mut query: Query<&mut C>) {
151 /// let mut iter = query.iter_mut();
152 /// while let Some(mut c) = iter.fetch_next() {
153 /// //
154 /// }
155 /// }
156 /// # bevy_ecs::system::assert_is_system(system);
157 /// ```
158 pub fn fetch_next(&mut self) -> Option<D::Item<'_, 's>> {
159 // SAFETY:
160 // - `tables` and `archetypes` belong to the same world that the cursor was initialized for.
161 // - `query_state` is the state that was passed to `QueryIterationCursor::init`.
162 // - `self` is mutably borrowed, so there are no other items alive for any entity.
163 unsafe {
164 self.cursor
165 .next(self.tables, self.archetypes, self.query_state)
166 .map(D::shrink)
167 }
168 }
169}
170
171impl<'w, 's, D: IterQueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
172 /// Executes the equivalent of [`Iterator::fold`] over a contiguous segment
173 /// from a storage.
174 ///
175 /// # Safety
176 /// - `range` must be in `[0, storage::entity_count)` or None.
177 #[inline]
178 pub(super) unsafe fn fold_over_storage_range<B, Func>(
179 &mut self,
180 mut accum: B,
181 func: &mut Func,
182 storage: StorageId,
183 range: Option<Range<u32>>,
184 ) -> B
185 where
186 Func: FnMut(B, D::Item<'w, 's>) -> B,
187 {
188 if self.cursor.is_dense {
189 // SAFETY: `self.cursor.is_dense` is true, so storage ids are guaranteed to be table ids.
190 let table_id = unsafe { storage.table_id };
191 // SAFETY: Matched table IDs are guaranteed to still exist.
192 let table = unsafe { self.tables.get(table_id).debug_checked_unwrap() };
193
194 let range = range.unwrap_or(0..table.entity_count());
195 accum =
196 // SAFETY:
197 // - The fetched table matches both D and F
198 // - caller ensures `range` is within `[0, table.entity_count)`
199 // - The if block ensures that the query iteration is dense
200 unsafe { self.fold_over_table_range(accum, func, table, range) };
201 } else {
202 // SAFETY: `self.cursor.is_dense` is false, so storage ids are guaranteed to be archetype ids.
203 let archetype_id = unsafe { storage.archetype_id };
204 // SAFETY: Matched archetype IDs are guaranteed to still exist.
205 let archetype = unsafe { self.archetypes.get(archetype_id).debug_checked_unwrap() };
206 // SAFETY: Matched table IDs are guaranteed to still exist.
207 let table = unsafe { self.tables.get(archetype.table_id()).debug_checked_unwrap() };
208
209 let range = range.unwrap_or(0..archetype.len());
210
211 // When an archetype and its table have equal entity counts, dense iteration can be safely used.
212 // this leverages cache locality to optimize performance.
213 if table.entity_count() == archetype.len() {
214 accum =
215 // SAFETY:
216 // - The fetched archetype matches both D and F
217 // - The provided archetype and its' table have the same length.
218 // - caller ensures `range` is within `[0, archetype.len)`
219 // - The if block ensures that the query iteration is not dense.
220 unsafe { self.fold_over_dense_archetype_range(accum, func, archetype,range) };
221 } else {
222 accum =
223 // SAFETY:
224 // - The fetched archetype matches both D and F
225 // - caller ensures `range` is within `[0, archetype.len)`
226 // - The if block ensures that the query iteration is not dense.
227 unsafe { self.fold_over_archetype_range(accum, func, archetype,range) };
228 }
229 }
230 accum
231 }
232
233 /// Executes the equivalent of [`Iterator::fold`] over a contiguous segment
234 /// from a table.
235 ///
236 /// # Safety
237 /// - all `rows` must be in `[0, table.entity_count)`.
238 /// - `table` must match D and F
239 /// - The query iteration must be dense (i.e. `self.query_state.is_dense` must be true).
240 #[inline]
241 pub(super) unsafe fn fold_over_table_range<B, Func>(
242 &mut self,
243 mut accum: B,
244 func: &mut Func,
245 table: &'w Table,
246 rows: Range<u32>,
247 ) -> B
248 where
249 Func: FnMut(B, D::Item<'w, 's>) -> B,
250 {
251 if table.is_empty() {
252 return accum;
253 }
254
255 D::set_table(&mut self.cursor.fetch, &self.query_state.fetch_state, table);
256 F::set_table(
257 &mut self.cursor.filter,
258 &self.query_state.filter_state,
259 table,
260 );
261
262 let entities = table.entities();
263 for row in rows {
264 // SAFETY: Caller assures `row` in range of the current archetype.
265 let entity = unsafe { entities.get_unchecked(row as usize) };
266 // SAFETY: This is from an exclusive range, so it can't be max.
267 let row = unsafe { TableRow::new(NonMaxU32::new_unchecked(row)) };
268
269 // SAFETY: set_table was called prior.
270 // Caller assures `row` in range of the current archetype.
271 let fetched = unsafe {
272 !F::filter_fetch(
273 &self.query_state.filter_state,
274 &mut self.cursor.filter,
275 *entity,
276 row,
277 )
278 };
279 if fetched {
280 continue;
281 }
282
283 // SAFETY:
284 // - set_table was called prior.
285 // - Caller assures `row` in range of the current archetype.
286 // - Each row is unique, so each entity is only alive once
287 // - `D: IterQueryData`
288 if let Some(item) = D::fetch(
289 &self.query_state.fetch_state,
290 &mut self.cursor.fetch,
291 *entity,
292 row,
293 ) {
294 accum = func(accum, item);
295 }
296 }
297 accum
298 }
299
300 /// Executes the equivalent of [`Iterator::fold`] over a contiguous segment
301 /// from an archetype.
302 ///
303 /// # Safety
304 /// - all `indices` must be in `[0, archetype.len())`.
305 /// - `archetype` must match D and F
306 /// - The query iteration must not be dense (i.e. `self.query_state.is_dense` must be false).
307 #[inline]
308 pub(super) unsafe fn fold_over_archetype_range<B, Func>(
309 &mut self,
310 mut accum: B,
311 func: &mut Func,
312 archetype: &'w Archetype,
313 indices: Range<u32>,
314 ) -> B
315 where
316 Func: FnMut(B, D::Item<'w, 's>) -> B,
317 {
318 if archetype.is_empty() {
319 return accum;
320 }
321 let table = self.tables.get(archetype.table_id()).debug_checked_unwrap();
322 D::set_archetype(
323 &mut self.cursor.fetch,
324 &self.query_state.fetch_state,
325 archetype,
326 table,
327 );
328 F::set_archetype(
329 &mut self.cursor.filter,
330 &self.query_state.filter_state,
331 archetype,
332 table,
333 );
334
335 let entities = archetype.entities();
336 for index in indices {
337 // SAFETY: Caller assures `index` in range of the current archetype.
338 let archetype_entity = unsafe { entities.get_unchecked(index as usize) };
339
340 // SAFETY: set_archetype was called prior.
341 // Caller assures `index` in range of the current archetype.
342 let fetched = unsafe {
343 !F::filter_fetch(
344 &self.query_state.filter_state,
345 &mut self.cursor.filter,
346 archetype_entity.id(),
347 archetype_entity.table_row(),
348 )
349 };
350 if fetched {
351 continue;
352 }
353
354 // SAFETY:
355 // - set_archetype was called prior, `index` is an archetype index in range of the current archetype
356 // - Caller assures `index` in range of the current archetype.
357 // - Each row is unique, so each entity is only alive once
358 // - `D: IterQueryData`
359 if let Some(item) = unsafe {
360 D::fetch(
361 &self.query_state.fetch_state,
362 &mut self.cursor.fetch,
363 archetype_entity.id(),
364 archetype_entity.table_row(),
365 )
366 } {
367 accum = func(accum, item);
368 }
369 }
370 accum
371 }
372
373 /// Executes the equivalent of [`Iterator::fold`] over a contiguous segment
374 /// from an archetype which has the same entity count as its table.
375 ///
376 /// # Safety
377 /// - all `indices` must be in `[0, archetype.len())`.
378 /// - `archetype` must match D and F
379 /// - `archetype` must have the same length as its table.
380 /// - The query iteration must not be dense (i.e. `self.query_state.is_dense` must be false).
381 #[inline]
382 pub(super) unsafe fn fold_over_dense_archetype_range<B, Func>(
383 &mut self,
384 mut accum: B,
385 func: &mut Func,
386 archetype: &'w Archetype,
387 rows: Range<u32>,
388 ) -> B
389 where
390 Func: FnMut(B, D::Item<'w, 's>) -> B,
391 {
392 if archetype.is_empty() {
393 return accum;
394 }
395 let table = self.tables.get(archetype.table_id()).debug_checked_unwrap();
396 debug_assert!(
397 archetype.len() == table.entity_count(),
398 "archetype and its table must have the same length. "
399 );
400
401 D::set_archetype(
402 &mut self.cursor.fetch,
403 &self.query_state.fetch_state,
404 archetype,
405 table,
406 );
407 F::set_archetype(
408 &mut self.cursor.filter,
409 &self.query_state.filter_state,
410 archetype,
411 table,
412 );
413 let entities = table.entities();
414 for row in rows {
415 // SAFETY: Caller assures `row` in range of the current archetype.
416 let entity = unsafe { *entities.get_unchecked(row as usize) };
417 // SAFETY: This is from an exclusive range, so it can't be max.
418 let row = unsafe { TableRow::new(NonMaxU32::new_unchecked(row)) };
419
420 // SAFETY: set_table was called prior.
421 // Caller assures `row` in range of the current archetype.
422 let filter_matched = unsafe {
423 F::filter_fetch(
424 &self.query_state.filter_state,
425 &mut self.cursor.filter,
426 entity,
427 row,
428 )
429 };
430 if !filter_matched {
431 continue;
432 }
433
434 // SAFETY:
435 // - set_table was called prior.
436 // - Caller assures `row` in range of the current archetype.
437 // - Each row is unique, so each entity is only alive once
438 // - `D: IterQueryData`
439 if let Some(item) = D::fetch(
440 &self.query_state.fetch_state,
441 &mut self.cursor.fetch,
442 entity,
443 row,
444 ) {
445 accum = func(accum, item);
446 }
447 }
448 accum
449 }
450}
451
452impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
453 /// Sorts all query items into a new iterator, using the query lens as a key.
454 ///
455 /// This sort is stable (i.e., does not reorder equal elements).
456 ///
457 /// This uses [`slice::sort`] internally.
458 ///
459 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
460 /// This includes the allowed parameter type changes listed under [allowed transmutes].
461 /// However, the lens uses the filter of the original query when present.
462 ///
463 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
464 /// of query transmutes does not support nested queries.
465 /// This restriction may be lifted in the future.
466 ///
467 /// The sort is not cached across system runs.
468 ///
469 /// If the [`QueryData`] does not implement [`IterQueryData`],
470 /// then it is not sound to yield multiple items concurrently
471 /// and the resulting [`QuerySortedIter`] will not implement [`Iterator`].
472 /// To iterate over the items in that case,
473 /// use the [`QuerySortedIter::fetch_next()`](crate::query::QuerySortedIter::fetch_next) method,
474 /// which ensures only one item is alive at a time.
475 ///
476 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
477 ///
478 /// # Panics
479 ///
480 /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
481 ///
482 /// # Examples
483 /// ```rust
484 /// # use bevy_ecs::prelude::*;
485 /// # use std::{ops::{Deref, DerefMut}, iter::Sum};
486 /// #
487 /// # #[derive(Component)]
488 /// # struct PartMarker;
489 /// #
490 /// # #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
491 /// # struct PartIndex(usize);
492 /// #
493 /// # #[derive(Component, Clone, Copy)]
494 /// # struct PartValue(f32);
495 /// #
496 /// # impl Deref for PartValue {
497 /// # type Target = f32;
498 /// #
499 /// # fn deref(&self) -> &Self::Target {
500 /// # &self.0
501 /// # }
502 /// # }
503 /// #
504 /// # #[derive(Component)]
505 /// # struct ParentValue(f32);
506 /// #
507 /// # impl Deref for ParentValue {
508 /// # type Target = f32;
509 /// #
510 /// # fn deref(&self) -> &Self::Target {
511 /// # &self.0
512 /// # }
513 /// # }
514 /// #
515 /// # impl DerefMut for ParentValue {
516 /// # fn deref_mut(&mut self) -> &mut Self::Target {
517 /// # &mut self.0
518 /// # }
519 /// # }
520 /// #
521 /// # #[derive(Component, Debug, PartialEq, Eq, PartialOrd, Ord)]
522 /// # struct Length(usize);
523 /// #
524 /// # #[derive(Component, Debug, PartialEq, Eq, PartialOrd, Ord)]
525 /// # struct Width(usize);
526 /// #
527 /// # #[derive(Component, Debug, PartialEq, Eq, PartialOrd, Ord)]
528 /// # struct Height(usize);
529 /// #
530 /// # #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
531 /// # struct ParentEntity(Entity);
532 /// #
533 /// # #[derive(Component, Clone, Copy)]
534 /// # struct ChildPartCount(usize);
535 /// #
536 /// # impl Deref for ChildPartCount {
537 /// # type Target = usize;
538 /// #
539 /// # fn deref(&self) -> &Self::Target {
540 /// # &self.0
541 /// # }
542 /// # }
543 /// # let mut world = World::new();
544 /// // We can ensure that a query always returns in the same order.
545 /// fn system_1(query: Query<(Entity, &PartIndex)>) {
546 /// let parts: Vec<(Entity, &PartIndex)> = query.iter().sort::<&PartIndex>().collect();
547 /// }
548 ///
549 /// // We can freely rearrange query components in the key.
550 /// fn system_2(query: Query<(&Length, &Width, &Height), With<PartMarker>>) {
551 /// for (length, width, height) in query.iter().sort::<(&Height, &Length, &Width)>() {
552 /// println!("height: {height:?}, width: {width:?}, length: {length:?}")
553 /// }
554 /// }
555 ///
556 /// // We can sort by Entity without including it in the original Query.
557 /// // Here, we match iteration orders between query iterators.
558 /// fn system_3(
559 /// part_query: Query<(&PartValue, &ParentEntity)>,
560 /// mut parent_query: Query<(&ChildPartCount, &mut ParentValue)>,
561 /// ) {
562 /// let part_values = &mut part_query
563 /// .into_iter()
564 /// .sort::<&ParentEntity>()
565 /// .map(|(&value, parent_entity)| *value);
566 ///
567 /// for (&child_count, mut parent_value) in parent_query.iter_mut().sort::<Entity>() {
568 /// **parent_value = part_values.take(*child_count).sum();
569 /// }
570 /// }
571 /// #
572 /// # let mut schedule = Schedule::default();
573 /// # schedule.add_systems((system_1, system_2, system_3));
574 /// # schedule.run(&mut world);
575 /// ```
576 pub fn sort<L: ReadOnlyQueryData + SingleEntityQueryData + 'w>(
577 self,
578 ) -> QuerySortedIter<
579 'w,
580 's,
581 D,
582 F,
583 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
584 >
585 where
586 for<'lw, 'ls> L::Item<'lw, 'ls>: Ord,
587 {
588 self.sort_impl::<L>(|keyed_query| keyed_query.sort())
589 }
590
591 /// Sorts all query items into a new iterator, using the query lens as a key.
592 ///
593 /// This sort is unstable (i.e., may reorder equal elements).
594 ///
595 /// This uses [`slice::sort_unstable`] internally.
596 ///
597 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
598 /// This includes the allowed parameter type changes listed under [allowed transmutes]..
599 /// However, the lens uses the filter of the original query when present.
600 ///
601 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
602 /// of query transmutes does not support nested queries.
603 /// This restriction may be lifted in the future.
604 ///
605 /// The sort is not cached across system runs.
606 ///
607 /// If the [`QueryData`] does not implement [`IterQueryData`],
608 /// then it is not sound to yield multiple items concurrently
609 /// and the resulting [`QuerySortedIter`] will not implement [`Iterator`].
610 /// To iterate over the items in that case,
611 /// use the [`QuerySortedIter::fetch_next()`](crate::query::QuerySortedIter::fetch_next) method,
612 /// which ensures only one item is alive at a time.
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 /// #
624 /// # let mut world = World::new();
625 /// #
626 /// # #[derive(Component)]
627 /// # struct PartMarker;
628 /// #
629 /// #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
630 /// enum Flying {
631 /// Enabled,
632 /// Disabled
633 /// };
634 ///
635 /// // We perform an unstable sort by a Component with few values.
636 /// fn system_1(query: Query<&Flying, With<PartMarker>>) {
637 /// let part_values: Vec<&Flying> = query.iter().sort_unstable::<&Flying>().collect();
638 /// }
639 /// #
640 /// # let mut schedule = Schedule::default();
641 /// # schedule.add_systems((system_1));
642 /// # schedule.run(&mut world);
643 /// ```
644 pub fn sort_unstable<L: ReadOnlyQueryData + SingleEntityQueryData + 'w>(
645 self,
646 ) -> QuerySortedIter<
647 'w,
648 's,
649 D,
650 F,
651 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
652 >
653 where
654 for<'lw, 'ls> L::Item<'lw, 'ls>: Ord,
655 {
656 self.sort_impl::<L>(|keyed_query| keyed_query.sort_unstable())
657 }
658
659 /// Sorts all query items into a new iterator with a comparator function over the query lens.
660 ///
661 /// This sort is stable (i.e., does not reorder equal elements).
662 ///
663 /// This uses [`slice::sort_by`] internally.
664 ///
665 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
666 /// This includes the allowed parameter type changes listed under [allowed transmutes].
667 /// However, the lens uses the filter of the original query when present.
668 ///
669 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
670 /// of query transmutes does not support nested queries.
671 /// This restriction may be lifted in the future.
672 ///
673 /// The sort is not cached across system runs.
674 ///
675 /// If the [`QueryData`] does not implement [`IterQueryData`],
676 /// then it is not sound to yield multiple items concurrently
677 /// and the resulting [`QuerySortedIter`] will not implement [`Iterator`].
678 /// To iterate over the items in that case,
679 /// use the [`QuerySortedIter::fetch_next()`](crate::query::QuerySortedIter::fetch_next) method,
680 /// which ensures only one item is alive at a time.
681 ///
682 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
683 ///
684 /// # Panics
685 ///
686 /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
687 ///
688 /// # Example
689 /// ```
690 /// # use bevy_ecs::prelude::*;
691 /// # use std::ops::Deref;
692 /// #
693 /// # impl Deref for PartValue {
694 /// # type Target = f32;
695 /// #
696 /// # fn deref(&self) -> &Self::Target {
697 /// # &self.0
698 /// # }
699 /// # }
700 /// #
701 /// # let mut world = World::new();
702 /// #
703 /// #[derive(Component)]
704 /// struct PartValue(f32);
705 ///
706 /// // We can use a cmp function on components do not implement Ord.
707 /// fn system_1(query: Query<&PartValue>) {
708 /// // Sort part values according to `f32::total_comp`.
709 /// let part_values: Vec<&PartValue> = query
710 /// .iter()
711 /// .sort_by::<&PartValue>(|value_1, value_2| value_1.total_cmp(*value_2))
712 /// .collect();
713 /// }
714 /// #
715 /// # let mut schedule = Schedule::default();
716 /// # schedule.add_systems((system_1));
717 /// # schedule.run(&mut world);
718 /// ```
719 pub fn sort_by<L: ReadOnlyQueryData + SingleEntityQueryData + 'w>(
720 self,
721 mut compare: impl FnMut(&L::Item<'_, '_>, &L::Item<'_, '_>) -> Ordering,
722 ) -> QuerySortedIter<
723 'w,
724 's,
725 D,
726 F,
727 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
728 > {
729 self.sort_impl::<L>(move |keyed_query| {
730 keyed_query.sort_by(|(key_1, _), (key_2, _)| compare(key_1, key_2));
731 })
732 }
733
734 /// Sorts all query items into a new iterator with a comparator function over the query lens.
735 ///
736 /// This sort is unstable (i.e., may reorder equal elements).
737 ///
738 /// This uses [`slice::sort_unstable_by`] internally.
739 ///
740 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
741 /// This includes the allowed parameter type changes listed under [allowed transmutes].
742 /// However, the lens uses the filter of the original query when present.
743 ///
744 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
745 /// of query transmutes does not support nested queries.
746 /// This restriction may be lifted in the future.
747 ///
748 /// The sort is not cached across system runs.
749 ///
750 /// If the [`QueryData`] does not implement [`IterQueryData`],
751 /// then it is not sound to yield multiple items concurrently
752 /// and the resulting [`QuerySortedIter`] will not implement [`Iterator`].
753 /// To iterate over the items in that case,
754 /// use the [`QuerySortedIter::fetch_next()`](crate::query::QuerySortedIter::fetch_next) method,
755 /// which ensures only one item is alive at a time.
756 ///
757 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
758 ///
759 /// # Panics
760 ///
761 /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
762 pub fn sort_unstable_by<L: ReadOnlyQueryData + SingleEntityQueryData + 'w>(
763 self,
764 mut compare: impl FnMut(&L::Item<'_, '_>, &L::Item<'_, '_>) -> Ordering,
765 ) -> QuerySortedIter<
766 'w,
767 's,
768 D,
769 F,
770 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
771 > {
772 self.sort_impl::<L>(move |keyed_query| {
773 keyed_query.sort_unstable_by(|(key_1, _), (key_2, _)| compare(key_1, key_2));
774 })
775 }
776
777 /// Sorts all query items into a new iterator with a key extraction function over the query lens.
778 ///
779 /// This sort is stable (i.e., does not reorder equal elements).
780 ///
781 /// This uses [`slice::sort_by_key`] internally.
782 ///
783 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
784 /// This includes the allowed parameter type changes listed under [allowed transmutes].
785 /// However, the lens uses the filter of the original query when present.
786 ///
787 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
788 /// of query transmutes does not support nested queries.
789 /// This restriction may be lifted in the future.
790 ///
791 /// The sort is not cached across system runs.
792 ///
793 /// If the [`QueryData`] does not implement [`IterQueryData`],
794 /// then it is not sound to yield multiple items concurrently
795 /// and the resulting [`QuerySortedIter`] will not implement [`Iterator`].
796 /// To iterate over the items in that case,
797 /// use the [`QuerySortedIter::fetch_next()`](crate::query::QuerySortedIter::fetch_next) method,
798 /// which ensures only one item is alive at a time.
799 ///
800 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
801 ///
802 /// # Panics
803 ///
804 /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
805 ///
806 /// # Example
807 /// ```
808 /// # use bevy_ecs::prelude::*;
809 /// # use std::ops::Deref;
810 /// #
811 /// # #[derive(Component)]
812 /// # struct PartMarker;
813 /// #
814 /// # impl Deref for PartValue {
815 /// # type Target = f32;
816 /// #
817 /// # fn deref(&self) -> &Self::Target {
818 /// # &self.0
819 /// # }
820 /// # }
821 /// #
822 /// # let mut world = World::new();
823 /// #
824 /// #[derive(Component)]
825 /// struct AvailableMarker;
826 ///
827 /// #[derive(Component, PartialEq, Eq, PartialOrd, Ord, Copy, Clone)]
828 /// enum Rarity {
829 /// Common,
830 /// Rare,
831 /// Epic,
832 /// Legendary
833 /// };
834 ///
835 /// #[derive(Component)]
836 /// struct PartValue(f32);
837 ///
838 /// // We can sort with the internals of components that do not implement Ord.
839 /// fn system_1(query: Query<(Entity, &PartValue)>) {
840 /// // Sort by the sines of the part values.
841 /// let parts: Vec<(Entity, &PartValue)> = query
842 /// .iter()
843 /// .sort_by_key::<&PartValue, _>(|value| value.sin() as usize)
844 /// .collect();
845 /// }
846 ///
847 /// // We can define our own custom comparison functions over an EntityRef.
848 /// fn system_2(query: Query<EntityRef, With<PartMarker>>) {
849 /// // Sort by whether parts are available and their rarity.
850 /// // We want the available legendaries to come first, so we reverse the iterator.
851 /// let parts: Vec<EntityRef> = query.iter()
852 /// .sort_by_key::<EntityRef, _>(|entity_ref| {
853 /// (
854 /// entity_ref.contains::<AvailableMarker>(),
855 /// entity_ref.get::<Rarity>().copied()
856 /// )
857 /// })
858 /// .rev()
859 /// .collect();
860 /// }
861 /// # let mut schedule = Schedule::default();
862 /// # schedule.add_systems((system_1, system_2));
863 /// # schedule.run(&mut world);
864 /// ```
865 pub fn sort_by_key<L: ReadOnlyQueryData + SingleEntityQueryData + 'w, K>(
866 self,
867 mut f: impl FnMut(&L::Item<'_, '_>) -> K,
868 ) -> QuerySortedIter<
869 'w,
870 's,
871 D,
872 F,
873 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
874 >
875 where
876 K: Ord,
877 {
878 self.sort_impl::<L>(move |keyed_query| keyed_query.sort_by_key(|(lens, _)| f(lens)))
879 }
880
881 /// Sorts all query items into a new iterator with a key extraction function over the query lens.
882 ///
883 /// This sort is unstable (i.e., may reorder equal elements).
884 ///
885 /// This uses [`slice::sort_unstable_by_key`] internally.
886 ///
887 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
888 /// This includes the allowed parameter type changes listed under [allowed transmutes].
889 /// However, the lens uses the filter of the original query when present.
890 ///
891 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
892 /// of query transmutes does not support nested queries.
893 /// This restriction may be lifted in the future.
894 ///
895 /// The sort is not cached across system runs.
896 ///
897 /// If the [`QueryData`] does not implement [`IterQueryData`],
898 /// then it is not sound to yield multiple items concurrently
899 /// and the resulting [`QuerySortedIter`] will not implement [`Iterator`].
900 /// To iterate over the items in that case,
901 /// use the [`QuerySortedIter::fetch_next()`](crate::query::QuerySortedIter::fetch_next) method,
902 /// which ensures only one item is alive at a time.
903 ///
904 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
905 ///
906 /// # Panics
907 ///
908 /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
909 pub fn sort_unstable_by_key<L: ReadOnlyQueryData + SingleEntityQueryData + 'w, K>(
910 self,
911 mut f: impl FnMut(&L::Item<'_, '_>) -> K,
912 ) -> QuerySortedIter<
913 'w,
914 's,
915 D,
916 F,
917 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
918 >
919 where
920 K: Ord,
921 {
922 self.sort_impl::<L>(move |keyed_query| {
923 keyed_query.sort_unstable_by_key(|(lens, _)| f(lens));
924 })
925 }
926
927 /// Sort all query items into a new iterator with a key extraction function over the query lens.
928 ///
929 /// This sort is stable (i.e., does not reorder equal elements).
930 ///
931 /// This uses [`slice::sort_by_cached_key`] internally.
932 ///
933 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
934 /// This includes the allowed parameter type changes listed under [allowed transmutes].
935 /// However, the lens uses the filter of the original query when present.
936 ///
937 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
938 /// of query transmutes does not support nested queries.
939 /// This restriction may be lifted in the future.
940 ///
941 /// The sort is not cached across system runs.
942 ///
943 /// If the [`QueryData`] does not implement [`IterQueryData`],
944 /// then it is not sound to yield multiple items concurrently
945 /// and the resulting [`QuerySortedIter`] will not implement [`Iterator`].
946 /// To iterate over the items in that case,
947 /// use the [`QuerySortedIter::fetch_next()`](crate::query::QuerySortedIter::fetch_next) method,
948 /// which ensures only one item is alive at a time.
949 ///
950 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
951 ///
952 /// # Panics
953 ///
954 /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
955 pub fn sort_by_cached_key<L: ReadOnlyQueryData + SingleEntityQueryData + 'w, K>(
956 self,
957 mut f: impl FnMut(&L::Item<'_, '_>) -> K,
958 ) -> QuerySortedIter<
959 'w,
960 's,
961 D,
962 F,
963 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
964 >
965 where
966 K: Ord,
967 {
968 self.sort_impl::<L>(move |keyed_query| keyed_query.sort_by_cached_key(|(lens, _)| f(lens)))
969 }
970
971 /// Shared implementation for the various `sort` methods.
972 /// This uses the lens to collect the items for sorting, but delegates the actual sorting to the provided closure.
973 ///
974 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
975 /// This includes the allowed parameter type changes listed under [allowed transmutes].
976 /// However, the lens uses the filter of the original query when present.
977 ///
978 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
979 /// of query transmutes does not support nested queries.
980 /// This restriction may be lifted in the future.
981 ///
982 /// The sort is not cached across system runs.
983 ///
984 /// If the [`QueryData`] does not implement [`IterQueryData`],
985 /// then it is not sound to yield multiple items concurrently
986 /// and the resulting [`QuerySortedIter`] will not implement [`Iterator`].
987 /// To iterate over the items in that case,
988 /// use the [`QuerySortedIter::fetch_next()`](crate::query::QuerySortedIter::fetch_next) method,
989 /// which ensures only one item is alive at a time.
990 ///
991 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
992 ///
993 /// # Panics
994 ///
995 /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
996 fn sort_impl<L: ReadOnlyQueryData + SingleEntityQueryData + 'w>(
997 self,
998 f: impl FnOnce(&mut Vec<(L::Item<'_, '_>, NeutralOrd<Entity>)>),
999 ) -> QuerySortedIter<
1000 'w,
1001 's,
1002 D,
1003 F,
1004 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
1005 > {
1006 // On the first successful iteration of `QueryIterationCursor`, `archetype_entities` or `table_entities`
1007 // will be set to a non-zero value. The correctness of this method relies on this.
1008 // I.e. this sort method will execute if and only if `next` on `QueryIterationCursor` of a
1009 // non-empty `QueryIter` has not yet been called. When empty, this sort method will not panic.
1010 if !self.cursor.archetype_entities.is_empty() || !self.cursor.table_entities.is_empty() {
1011 panic!("it is not valid to call sort() after next()")
1012 }
1013
1014 let world = self.world;
1015
1016 let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);
1017
1018 // SAFETY:
1019 // `self.world` has permission to access the required components.
1020 // The original query iter has not been iterated on, so no items are aliased from it.
1021 // `QueryIter::new` ensures `world` is the same one used to initialize `query_state`.
1022 let query_lens = unsafe { query_lens_state.query_unchecked_manual(world) }.into_iter();
1023 let mut keyed_query: Vec<_> = query_lens
1024 .map(|(key, entity)| (key, NeutralOrd(entity)))
1025 .collect();
1026 f(&mut keyed_query);
1027 let entity_iter = keyed_query
1028 .into_iter()
1029 .map(|(.., entity)| entity.0)
1030 .collect::<Vec<_>>()
1031 .into_iter();
1032 // SAFETY:
1033 // `self.world` has permission to access the required components.
1034 // Each lens query item is dropped before the respective actual query item is accessed.
1035 unsafe {
1036 QuerySortedIter::new(
1037 world,
1038 self.query_state,
1039 entity_iter,
1040 world.last_change_tick(),
1041 world.change_tick(),
1042 )
1043 }
1044 }
1045}
1046
1047impl<'w, 's, D: IterQueryData, F: QueryFilter> Iterator for QueryIter<'w, 's, D, F> {
1048 type Item = D::Item<'w, 's>;
1049
1050 #[inline(always)]
1051 fn next(&mut self) -> Option<Self::Item> {
1052 // SAFETY:
1053 // - `tables` and `archetypes` belong to the same world that the cursor was initialized for.
1054 // - `query_state` is the state that was passed to `QueryIterationCursor::init`.
1055 // - `D: IterQueryData`
1056 unsafe {
1057 self.cursor
1058 .next(self.tables, self.archetypes, self.query_state)
1059 }
1060 }
1061
1062 fn size_hint(&self) -> (usize, Option<usize>) {
1063 let max_size = self.cursor.max_remaining(self.tables, self.archetypes);
1064 let archetype_query = D::IS_ARCHETYPAL && F::IS_ARCHETYPAL;
1065 let min_size = if archetype_query { max_size } else { 0 };
1066 (min_size as usize, Some(max_size as usize))
1067 }
1068
1069 #[inline]
1070 fn fold<B, Func>(mut self, init: B, mut func: Func) -> B
1071 where
1072 Func: FnMut(B, Self::Item) -> B,
1073 {
1074 let mut accum = init;
1075 // Empty any remaining uniterated values from the current table/archetype
1076 while self.cursor.current_row != self.cursor.current_len {
1077 let Some(item) = self.next() else { break };
1078 accum = func(accum, item);
1079 }
1080
1081 for id in self.cursor.storage_id_iter.clone().copied() {
1082 // SAFETY:
1083 // - The range(None) is equivalent to [0, storage.entity_count)
1084 accum = unsafe { self.fold_over_storage_range(accum, &mut func, id, None) };
1085 }
1086 accum
1087 }
1088}
1089
1090// This is correct as [`QueryIter`] always returns `None` once exhausted.
1091impl<'w, 's, D: IterQueryData, F: QueryFilter> FusedIterator for QueryIter<'w, 's, D, F> {}
1092
1093// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once.
1094unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator for QueryIter<'w, 's, Entity, F> {}
1095
1096// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once.
1097unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator for QueryIter<'w, 's, EntityRef<'_>, F> {}
1098
1099// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once.
1100unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator for QueryIter<'w, 's, EntityMut<'_>, F> {}
1101
1102// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once.
1103unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator
1104 for QueryIter<'w, 's, FilteredEntityRef<'_, '_>, F>
1105{
1106}
1107
1108// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once.
1109unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator
1110 for QueryIter<'w, 's, FilteredEntityMut<'_, '_>, F>
1111{
1112}
1113
1114// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once.
1115unsafe impl<'w, 's, F: QueryFilter, B: Bundle> EntitySetIterator
1116 for QueryIter<'w, 's, EntityRefExcept<'_, '_, B>, F>
1117{
1118}
1119
1120// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once.
1121unsafe impl<'w, 's, F: QueryFilter, B: Bundle> EntitySetIterator
1122 for QueryIter<'w, 's, EntityMutExcept<'_, '_, B>, F>
1123{
1124}
1125
1126impl<'w, 's, D: QueryData, F: QueryFilter> Debug for QueryIter<'w, 's, D, F> {
1127 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
1128 f.debug_struct("QueryIter").finish()
1129 }
1130}
1131
1132impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter> Clone for QueryIter<'w, 's, D, F> {
1133 fn clone(&self) -> Self {
1134 self.remaining()
1135 }
1136}
1137
1138/// Iterator for contiguous chunks of memory
1139pub struct QueryContiguousIter<'w, 's, D: ContiguousQueryData, F: ArchetypeFilter> {
1140 tables: &'w Tables,
1141 storage_id_iter: core::slice::Iter<'s, StorageId>,
1142 query_state: &'s QueryState<D, F>,
1143 fetch: D::Fetch<'w>,
1144 // NOTE: no need for F::Fetch because it always returns true
1145}
1146
1147impl<'w, 's, D: ContiguousQueryData, F: ArchetypeFilter> QueryContiguousIter<'w, 's, D, F> {
1148 /// # Safety
1149 /// - `world` must have permission to access any of the components registered in `query_state`.
1150 /// - `world` must be the same one used to initialize `query_state`.
1151 pub(crate) unsafe fn new(
1152 world: UnsafeWorldCell<'w>,
1153 query_state: &'s QueryState<D, F>,
1154 last_run: Tick,
1155 this_run: Tick,
1156 ) -> Option<Self> {
1157 query_state.is_dense.then(|| Self {
1158 // SAFETY: We only access table data that has been registered in `query_state`
1159 tables: unsafe { &world.storages().tables },
1160 storage_id_iter: query_state.matched_storage_ids.iter(),
1161 // SAFETY: The invariants are upheld by the caller.
1162 fetch: unsafe { D::init_fetch(world, &query_state.fetch_state, last_run, this_run) },
1163 query_state,
1164 })
1165 }
1166}
1167
1168impl<'w, 's, D: ContiguousQueryData, F: ArchetypeFilter> Iterator
1169 for QueryContiguousIter<'w, 's, D, F>
1170{
1171 type Item = D::Contiguous<'w, 's>;
1172
1173 #[inline(always)]
1174 fn next(&mut self) -> Option<Self::Item> {
1175 loop {
1176 // SAFETY: Query is dense
1177 let table_id = unsafe { self.storage_id_iter.next()?.table_id };
1178 // SAFETY: `table_id` was returned by `self.storage_id_iter` which always returns a
1179 // valid id
1180 let table = unsafe { self.tables.get(table_id).debug_checked_unwrap() };
1181 if table.is_empty() {
1182 continue;
1183 }
1184 // SAFETY:
1185 // - `table` is from the same world as `self.query_state` (`self.storage_id_iter` is
1186 // from the same world as `self.query_state`, see [`Self::new`])
1187 // - `self.fetch` was initialized with `self.query_state` (in [`Self::new`])
1188 unsafe {
1189 D::set_table(&mut self.fetch, &self.query_state.fetch_state, table);
1190 }
1191
1192 // no filtering because `F` implements `ArchetypeFilter` which ensures that `QueryFilter::fetch`
1193 // always returns true
1194
1195 // SAFETY:
1196 // - [`D::set_table`] is executed prior.
1197 // - `table.entities()` return a valid entity array
1198 // - the caller of [`Self::new`] ensures that the world has permission to access any of
1199 // the components registered in `self.query_state`
1200 let item = unsafe {
1201 D::fetch_contiguous(
1202 &self.query_state.fetch_state,
1203 &mut self.fetch,
1204 table.entities(),
1205 )
1206 };
1207
1208 return Some(item);
1209 }
1210 }
1211
1212 fn size_hint(&self) -> (usize, Option<usize>) {
1213 (0, self.storage_id_iter.size_hint().1)
1214 }
1215}
1216
1217// [`QueryIterationCursor::next_contiguous`] always returns None when exhausted
1218impl<'w, 's, D: ContiguousQueryData, F: ArchetypeFilter> FusedIterator
1219 for QueryContiguousIter<'w, 's, D, F>
1220{
1221}
1222
1223/// An [`Iterator`] over sorted query results of a [`Query`](crate::system::Query).
1224///
1225/// This struct is created by the [`QueryIter::sort`], [`QueryIter::sort_unstable`],
1226/// [`QueryIter::sort_by`], [`QueryIter::sort_unstable_by`], [`QueryIter::sort_by_key`],
1227/// [`QueryIter::sort_unstable_by_key`], and [`QueryIter::sort_by_cached_key`] methods.
1228pub struct QuerySortedIter<'w, 's, D: QueryData, F: QueryFilter, I>
1229where
1230 I: Iterator<Item = Entity>,
1231{
1232 entity_iter: I,
1233 entities: &'w Entities,
1234 tables: &'w Tables,
1235 archetypes: &'w Archetypes,
1236 fetch: D::Fetch<'w>,
1237 query_state: &'s QueryState<D, F>,
1238}
1239
1240impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> QuerySortedIter<'w, 's, D, F, I>
1241where
1242 I: Iterator<Item = Entity>,
1243{
1244 /// # Safety
1245 /// - `world` must have permission to access any of the components registered in `query_state`.
1246 /// - `world` must be the same one used to initialize `query_state`.
1247 /// - `entity_list` must only contain unique entities or be empty.
1248 pub(crate) unsafe fn new<EntityList: IntoIterator<IntoIter = I>>(
1249 world: UnsafeWorldCell<'w>,
1250 query_state: &'s QueryState<D, F>,
1251 entity_list: EntityList,
1252 last_run: Tick,
1253 this_run: Tick,
1254 ) -> QuerySortedIter<'w, 's, D, F, I> {
1255 let fetch = D::init_fetch(world, &query_state.fetch_state, last_run, this_run);
1256 QuerySortedIter {
1257 query_state,
1258 entities: world.entities(),
1259 archetypes: world.archetypes(),
1260 // SAFETY: We only access table data that has been registered in `query_state`.
1261 // This means `world` has permission to access the data we use.
1262 tables: &world.storages().tables,
1263 fetch,
1264 entity_iter: entity_list.into_iter(),
1265 }
1266 }
1267
1268 /// Get the next result from the query.
1269 ///
1270 /// If the [`QueryData`] does not implement [`IterQueryData`],
1271 /// then it is not sound to yield multiple items concurrently
1272 /// and the resulting [`QuerySortedIter`] will not implement [`Iterator`].
1273 /// In that case, this method can be used to iterate over the items
1274 /// while ensuring only one is alive at a time.
1275 ///
1276 /// Most queries do implement [`IterQueryData`],
1277 /// and can use the ordinary [`Iterator::next`]
1278 /// method or a `for` loop.
1279 ///
1280 /// # Example
1281 ///
1282 /// ```
1283 /// # use bevy_ecs::prelude::*;
1284 /// # #[derive(Component, Ord, PartialOrd, Eq, PartialEq)]
1285 /// # struct C;
1286 /// fn system(mut query: Query<&mut C>) {
1287 /// let mut iter = query.iter_mut().sort::<&C>();
1288 /// while let Some(mut c) = iter.fetch_next() {
1289 /// //
1290 /// }
1291 /// }
1292 /// # bevy_ecs::system::assert_is_system(system);
1293 /// ```
1294 pub fn fetch_next(&mut self) -> Option<D::Item<'_, 's>> {
1295 while let Some(entity) = self.entity_iter.next() {
1296 // SAFETY:
1297 // - `entity` is passed from `entity_iter` the first time.
1298 // - `self` is mutably borrowed, so there are no other items alive for any entity.
1299 if let Some(item) = unsafe { self.fetch_next_impl(entity) } {
1300 return Some(D::shrink(item));
1301 }
1302 }
1303 None
1304 }
1305
1306 /// Get the next result from the back of the query.
1307 ///
1308 /// If the [`QueryData`] does not implement [`IterQueryData`],
1309 /// then it is not sound to yield multiple items concurrently
1310 /// and the resulting [`QuerySortedIter`] will not implement [`Iterator`].
1311 /// In that case, this method can be used to iterate over the items
1312 /// while ensuring only one is alive at a time.
1313 ///
1314 /// Most queries do implement [`IterQueryData`],
1315 /// and can use the ordinary [`Iterator::next`]
1316 /// method or a `for` loop.
1317 pub fn fetch_next_back(&mut self) -> Option<D::Item<'_, 's>>
1318 where
1319 I: DoubleEndedIterator,
1320 {
1321 while let Some(entity) = self.entity_iter.next_back() {
1322 // SAFETY:
1323 // - `entity` is passed from `entity_iter` the first time.
1324 // - `self` is mutably borrowed, so there are no other items alive for any entity.
1325 if let Some(item) = unsafe { self.fetch_next_impl(entity) } {
1326 return Some(D::shrink(item));
1327 }
1328 }
1329 None
1330 }
1331
1332 /// # Safety
1333 ///
1334 /// - `entity` must stem from `self.entity_iter`
1335 /// - If `D` does not impl `ReadOnlyQueryData`, then there must not be any other `Item`s alive for the current entity
1336 /// - If `D` does not impl `IterQueryData`, then there must not be any other `Item`s alive for *any* entity
1337 #[inline(always)]
1338 unsafe fn fetch_next_impl(&mut self, entity: Entity) -> Option<D::Item<'w, 's>> {
1339 let (location, archetype, table);
1340 // SAFETY:
1341 // `tables` and `archetypes` belong to the same world that the [`QueryIter`]
1342 // was initialized for.
1343 unsafe {
1344 location = self.entities.get_spawned(entity).debug_checked_unwrap();
1345 archetype = self
1346 .archetypes
1347 .get(location.archetype_id)
1348 .debug_checked_unwrap();
1349 table = self.tables.get(location.table_id).debug_checked_unwrap();
1350 }
1351
1352 // SAFETY: `archetype` is from the world that `fetch` was created for,
1353 // `fetch_state` is the state that `fetch` was initialized with
1354 unsafe {
1355 D::set_archetype(
1356 &mut self.fetch,
1357 &self.query_state.fetch_state,
1358 archetype,
1359 table,
1360 );
1361 }
1362
1363 // The entity list has already been filtered by the query lens, so we forego filtering here.
1364 // SAFETY:
1365 // - set_archetype was called prior, `location.archetype_row` is an archetype index in range of the current archetype
1366 // - Caller ensures there are no conflicting items alive
1367 unsafe {
1368 D::fetch(
1369 &self.query_state.fetch_state,
1370 &mut self.fetch,
1371 entity,
1372 location.table_row,
1373 )
1374 }
1375 }
1376}
1377
1378impl<'w, 's, D: IterQueryData, F: QueryFilter, I: Iterator> Iterator
1379 for QuerySortedIter<'w, 's, D, F, I>
1380where
1381 I: Iterator<Item = Entity>,
1382{
1383 type Item = D::Item<'w, 's>;
1384
1385 #[inline(always)]
1386 fn next(&mut self) -> Option<Self::Item> {
1387 while let Some(entity) = self.entity_iter.next() {
1388 // SAFETY:
1389 // - `entity` is passed from `entity_iter` the first time.
1390 // - `D: IterQueryData`
1391 if let Some(item) = unsafe { self.fetch_next_impl(entity) } {
1392 return Some(item);
1393 }
1394 }
1395 None
1396 }
1397
1398 fn size_hint(&self) -> (usize, Option<usize>) {
1399 let (min_size, max_size) = self.entity_iter.size_hint();
1400 let archetype_query = D::IS_ARCHETYPAL;
1401 let min_size = if archetype_query { min_size } else { 0 };
1402 (min_size, max_size)
1403 }
1404}
1405
1406impl<'w, 's, D: IterQueryData, F: QueryFilter, I: Iterator> DoubleEndedIterator
1407 for QuerySortedIter<'w, 's, D, F, I>
1408where
1409 I: DoubleEndedIterator<Item = Entity>,
1410{
1411 #[inline(always)]
1412 fn next_back(&mut self) -> Option<Self::Item> {
1413 while let Some(entity) = self.entity_iter.next_back() {
1414 // SAFETY:
1415 // - `entity` is passed from `entity_iter` the first time.
1416 // - `D: IterQueryData`
1417 if let Some(item) = unsafe { self.fetch_next_impl(entity) } {
1418 return Some(item);
1419 }
1420 }
1421 None
1422 }
1423}
1424
1425impl<
1426 D: ArchetypeQueryData + IterQueryData,
1427 F: QueryFilter,
1428 I: ExactSizeIterator<Item = Entity>,
1429 > ExactSizeIterator for QuerySortedIter<'_, '_, D, F, I>
1430{
1431}
1432
1433// This is correct as [`QuerySortedIter`] returns `None` once exhausted if `entity_iter` does.
1434impl<'w, 's, D: IterQueryData, F: QueryFilter, I: Iterator> FusedIterator
1435 for QuerySortedIter<'w, 's, D, F, I>
1436where
1437 I: FusedIterator<Item = Entity>,
1438{
1439}
1440
1441// SAFETY:
1442// `I` stems from a collected and sorted `EntitySetIterator` ([`QueryIter`]).
1443// Fetching unique entities maintains uniqueness.
1444unsafe impl<'w, 's, F: QueryFilter, I: Iterator<Item = Entity>> EntitySetIterator
1445 for QuerySortedIter<'w, 's, Entity, F, I>
1446{
1447}
1448
1449impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item = Entity>> Debug
1450 for QuerySortedIter<'w, 's, D, F, I>
1451{
1452 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
1453 f.debug_struct("QuerySortedIter").finish()
1454 }
1455}
1456
1457/// An [`Iterator`] over the query items generated from an iterator of [`Entity`]s.
1458///
1459/// Items are returned in the order of the provided iterator.
1460/// Entities that don't match the query are skipped.
1461///
1462/// 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.
1463pub struct QueryManyIter<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityEquivalent>>
1464{
1465 world: UnsafeWorldCell<'w>,
1466 entity_iter: I,
1467 entities: &'w Entities,
1468 tables: &'w Tables,
1469 archetypes: &'w Archetypes,
1470 fetch: D::Fetch<'w>,
1471 filter: F::Fetch<'w>,
1472 query_state: &'s QueryState<D, F>,
1473}
1474
1475impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityEquivalent>>
1476 QueryManyIter<'w, 's, D, F, I>
1477{
1478 /// # Safety
1479 /// - `world` must have permission to access any of the components registered in `query_state`.
1480 /// - `world` must be the same one used to initialize `query_state`.
1481 pub(crate) unsafe fn new<EntityList: IntoIterator<IntoIter = I>>(
1482 world: UnsafeWorldCell<'w>,
1483 query_state: &'s QueryState<D, F>,
1484 entity_list: EntityList,
1485 last_run: Tick,
1486 this_run: Tick,
1487 ) -> QueryManyIter<'w, 's, D, F, I> {
1488 let fetch = D::init_fetch(world, &query_state.fetch_state, last_run, this_run);
1489 let filter = F::init_fetch(world, &query_state.filter_state, last_run, this_run);
1490 QueryManyIter {
1491 world,
1492 query_state,
1493 entities: world.entities(),
1494 archetypes: world.archetypes(),
1495 // SAFETY: We only access table data that has been registered in `query_state`.
1496 // This means `world` has permission to access the data we use.
1497 tables: &world.storages().tables,
1498 fetch,
1499 filter,
1500 entity_iter: entity_list.into_iter(),
1501 }
1502 }
1503
1504 /// # Safety
1505 ///
1506 /// - All arguments must stem from the same valid `QueryManyIter`.
1507 /// - If `D` does not impl `ReadOnlyQueryData`, then there must not be any other `Item`s alive for the current entity
1508 /// - If `D` does not impl `IterQueryData`, then there must not be any other `Item`s alive for *any* entity
1509 #[inline(always)]
1510 unsafe fn fetch_next_aliased_unchecked(
1511 entity_iter: impl Iterator<Item: EntityEquivalent>,
1512 entities: &'w Entities,
1513 tables: &'w Tables,
1514 archetypes: &'w Archetypes,
1515 fetch: &mut D::Fetch<'w>,
1516 filter: &mut F::Fetch<'w>,
1517 query_state: &'s QueryState<D, F>,
1518 ) -> Option<D::Item<'w, 's>> {
1519 for entity_borrow in entity_iter {
1520 let entity = entity_borrow.entity();
1521 let Ok(location) = entities.get_spawned(entity) else {
1522 continue;
1523 };
1524
1525 if !query_state
1526 .matched_archetypes
1527 .contains(location.archetype_id.index())
1528 {
1529 continue;
1530 }
1531
1532 let archetype = archetypes.get(location.archetype_id).debug_checked_unwrap();
1533 let table = tables.get(location.table_id).debug_checked_unwrap();
1534
1535 // SAFETY: `archetype` is from the world that `fetch/filter` were created for,
1536 // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
1537 unsafe {
1538 D::set_archetype(fetch, &query_state.fetch_state, archetype, table);
1539 }
1540 // SAFETY: `table` is from the world that `fetch/filter` were created for,
1541 // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
1542 unsafe {
1543 F::set_archetype(filter, &query_state.filter_state, archetype, table);
1544 }
1545
1546 // SAFETY: set_archetype was called prior.
1547 // `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
1548 if unsafe {
1549 F::filter_fetch(
1550 &query_state.filter_state,
1551 filter,
1552 entity,
1553 location.table_row,
1554 )
1555 } {
1556 // SAFETY:
1557 // - set_archetype was called prior, `location.archetype_row` is an archetype index in range of the current archetype
1558 // - Caller ensures there are no conflicting items alive
1559 let item = unsafe {
1560 D::fetch(&query_state.fetch_state, fetch, entity, location.table_row)
1561 };
1562 if let Some(item) = item {
1563 return Some(item);
1564 }
1565 }
1566 }
1567 None
1568 }
1569
1570 /// Get the next result from the query
1571 #[inline(always)]
1572 pub fn fetch_next(&mut self) -> Option<D::Item<'_, 's>> {
1573 // SAFETY:
1574 // All arguments stem from self.
1575 // We are limiting the returned reference to self,
1576 // making sure this method cannot be called multiple times without getting rid
1577 // of any previously returned unique references first, thus preventing aliasing.
1578 unsafe {
1579 Self::fetch_next_aliased_unchecked(
1580 &mut self.entity_iter,
1581 self.entities,
1582 self.tables,
1583 self.archetypes,
1584 &mut self.fetch,
1585 &mut self.filter,
1586 self.query_state,
1587 )
1588 .map(D::shrink)
1589 }
1590 }
1591
1592 /// Sorts all query items into a new iterator, using the query lens as a key.
1593 ///
1594 /// This sort is stable (i.e., does not reorder equal elements).
1595 ///
1596 /// This uses [`slice::sort`] internally.
1597 ///
1598 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
1599 /// This includes the allowed parameter type changes listed under [allowed transmutes].
1600 /// However, the lens uses the filter of the original query when present.
1601 ///
1602 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
1603 /// of query transmutes does not support nested queries.
1604 /// This restriction may be lifted in the future.
1605 ///
1606 /// The sort is not cached across system runs.
1607 ///
1608 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
1609 ///
1610 /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been
1611 /// called on [`QueryManyIter`] before.
1612 ///
1613 /// # Examples
1614 /// ```rust
1615 /// # use bevy_ecs::prelude::*;
1616 /// # use std::{ops::{Deref, DerefMut}, iter::Sum};
1617 /// #
1618 /// # #[derive(Component)]
1619 /// # struct PartMarker;
1620 /// #
1621 /// # #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
1622 /// # struct PartIndex(usize);
1623 /// #
1624 /// # #[derive(Component, Clone, Copy)]
1625 /// # struct PartValue(usize);
1626 /// #
1627 /// # impl Deref for PartValue {
1628 /// # type Target = usize;
1629 /// #
1630 /// # fn deref(&self) -> &Self::Target {
1631 /// # &self.0
1632 /// # }
1633 /// # }
1634 /// #
1635 /// # impl DerefMut for PartValue {
1636 /// # fn deref_mut(&mut self) -> &mut Self::Target {
1637 /// # &mut self.0
1638 /// # }
1639 /// # }
1640 /// #
1641 /// # #[derive(Component, Debug, PartialEq, Eq, PartialOrd, Ord)]
1642 /// # struct Length(usize);
1643 /// #
1644 /// # #[derive(Component, Debug, PartialEq, Eq, PartialOrd, Ord)]
1645 /// # struct Width(usize);
1646 /// #
1647 /// # #[derive(Component, Debug, PartialEq, Eq, PartialOrd, Ord)]
1648 /// # struct Height(usize);
1649 /// #
1650 /// # #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
1651 /// # struct ParentEntity(Entity);
1652 /// #
1653 /// # let mut world = World::new();
1654 /// // We can ensure that a query always returns in the same order.
1655 /// fn system_1(query: Query<(Entity, &PartIndex)>) {
1656 /// # let entity_list: Vec<Entity> = Vec::new();
1657 /// let parts: Vec<(Entity, &PartIndex)> = query.iter_many(entity_list).sort::<&PartIndex>().collect();
1658 /// }
1659 ///
1660 /// // We can freely rearrange query components in the key.
1661 /// fn system_2(query: Query<(&Length, &Width, &Height), With<PartMarker>>) {
1662 /// # let entity_list: Vec<Entity> = Vec::new();
1663 /// for (length, width, height) in query.iter_many(entity_list).sort::<(&Height, &Length, &Width)>() {
1664 /// println!("height: {height:?}, width: {width:?}, length: {length:?}")
1665 /// }
1666 /// }
1667 ///
1668 /// // You can use `fetch_next_back` to obtain mutable references in reverse order.
1669 /// fn system_3(
1670 /// mut query: Query<&mut PartValue>,
1671 /// ) {
1672 /// # let entity_list: Vec<Entity> = Vec::new();
1673 /// // We need to collect the internal iterator before iterating mutably
1674 /// let mut parent_query_iter = query.iter_many_mut(entity_list)
1675 /// .sort::<Entity>();
1676 ///
1677 /// let mut scratch_value = 0;
1678 /// while let Some(mut part_value) = parent_query_iter.fetch_next_back()
1679 /// {
1680 /// // some order-dependent operation, here bitwise XOR
1681 /// **part_value ^= scratch_value;
1682 /// scratch_value = **part_value;
1683 /// }
1684 /// }
1685 /// #
1686 /// # let mut schedule = Schedule::default();
1687 /// # schedule.add_systems((system_1, system_2, system_3));
1688 /// # schedule.run(&mut world);
1689 /// ```
1690 pub fn sort<L: ReadOnlyQueryData + SingleEntityQueryData + 'w>(
1691 self,
1692 ) -> QuerySortedManyIter<
1693 'w,
1694 's,
1695 D,
1696 F,
1697 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
1698 >
1699 where
1700 for<'lw, 'ls> L::Item<'lw, 'ls>: Ord,
1701 {
1702 self.sort_impl::<L>(|keyed_query| keyed_query.sort())
1703 }
1704
1705 /// Sorts all query items into a new iterator, using the query lens as a key.
1706 ///
1707 /// This sort is unstable (i.e., may reorder equal elements).
1708 ///
1709 /// This uses [`slice::sort_unstable`] internally.
1710 ///
1711 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
1712 /// This includes the allowed parameter type changes listed under [allowed transmutes]..
1713 /// However, the lens uses the filter of the original query when present.
1714 ///
1715 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
1716 /// of query transmutes does not support nested queries.
1717 /// This restriction may be lifted in the future.
1718 ///
1719 /// The sort is not cached across system runs.
1720 ///
1721 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
1722 ///
1723 /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been
1724 /// called on [`QueryManyIter`] before.
1725 ///
1726 /// # Example
1727 /// ```
1728 /// # use bevy_ecs::prelude::*;
1729 /// #
1730 /// # let mut world = World::new();
1731 /// #
1732 /// # #[derive(Component)]
1733 /// # struct PartMarker;
1734 /// #
1735 /// # let entity_list: Vec<Entity> = Vec::new();
1736 /// #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
1737 /// enum Flying {
1738 /// Enabled,
1739 /// Disabled
1740 /// };
1741 ///
1742 /// // We perform an unstable sort by a Component with few values.
1743 /// fn system_1(query: Query<&Flying, With<PartMarker>>) {
1744 /// # let entity_list: Vec<Entity> = Vec::new();
1745 /// let part_values: Vec<&Flying> = query.iter_many(entity_list).sort_unstable::<&Flying>().collect();
1746 /// }
1747 /// #
1748 /// # let mut schedule = Schedule::default();
1749 /// # schedule.add_systems((system_1));
1750 /// # schedule.run(&mut world);
1751 /// ```
1752 pub fn sort_unstable<L: ReadOnlyQueryData + SingleEntityQueryData + 'w>(
1753 self,
1754 ) -> QuerySortedManyIter<
1755 'w,
1756 's,
1757 D,
1758 F,
1759 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
1760 >
1761 where
1762 for<'lw, 'ls> L::Item<'lw, 'ls>: Ord,
1763 {
1764 self.sort_impl::<L>(|keyed_query| keyed_query.sort_unstable())
1765 }
1766
1767 /// Sorts all query items into a new iterator with a comparator function over the query lens.
1768 ///
1769 /// This sort is stable (i.e., does not reorder equal elements).
1770 ///
1771 /// This uses [`slice::sort_by`] internally.
1772 ///
1773 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
1774 /// This includes the allowed parameter type changes listed under [allowed transmutes].
1775 /// However, the lens uses the filter of the original query when present.
1776 ///
1777 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
1778 /// of query transmutes does not support nested queries.
1779 /// This restriction may be lifted in the future.
1780 ///
1781 /// The sort is not cached across system runs.
1782 ///
1783 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
1784 ///
1785 /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been
1786 /// called on [`QueryManyIter`] before.
1787 ///
1788 /// # Example
1789 /// ```
1790 /// # use bevy_ecs::prelude::*;
1791 /// # use std::ops::Deref;
1792 /// #
1793 /// # impl Deref for PartValue {
1794 /// # type Target = f32;
1795 /// #
1796 /// # fn deref(&self) -> &Self::Target {
1797 /// # &self.0
1798 /// # }
1799 /// # }
1800 /// #
1801 /// # let mut world = World::new();
1802 /// # let entity_list: Vec<Entity> = Vec::new();
1803 /// #
1804 /// #[derive(Component)]
1805 /// struct PartValue(f32);
1806 ///
1807 /// // We can use a cmp function on components do not implement Ord.
1808 /// fn system_1(query: Query<&PartValue>) {
1809 /// # let entity_list: Vec<Entity> = Vec::new();
1810 /// // Sort part values according to `f32::total_comp`.
1811 /// let part_values: Vec<&PartValue> = query
1812 /// .iter_many(entity_list)
1813 /// .sort_by::<&PartValue>(|value_1, value_2| value_1.total_cmp(*value_2))
1814 /// .collect();
1815 /// }
1816 /// #
1817 /// # let mut schedule = Schedule::default();
1818 /// # schedule.add_systems((system_1));
1819 /// # schedule.run(&mut world);
1820 /// ```
1821 pub fn sort_by<L: ReadOnlyQueryData + SingleEntityQueryData + 'w>(
1822 self,
1823 mut compare: impl FnMut(&L::Item<'_, '_>, &L::Item<'_, '_>) -> Ordering,
1824 ) -> QuerySortedManyIter<
1825 'w,
1826 's,
1827 D,
1828 F,
1829 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
1830 > {
1831 self.sort_impl::<L>(move |keyed_query| {
1832 keyed_query.sort_by(|(key_1, _), (key_2, _)| compare(key_1, key_2));
1833 })
1834 }
1835
1836 /// Sorts all query items into a new iterator with a comparator function over the query lens.
1837 ///
1838 /// This sort is unstable (i.e., may reorder equal elements).
1839 ///
1840 /// This uses [`slice::sort_unstable_by`] internally.
1841 ///
1842 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
1843 /// This includes the allowed parameter type changes listed under [allowed transmutes].
1844 /// However, the lens uses the filter of the original query when present.
1845 ///
1846 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
1847 /// of query transmutes does not support nested queries.
1848 /// This restriction may be lifted in the future.
1849 ///
1850 /// The sort is not cached across system runs.
1851 ///
1852 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
1853 ///
1854 /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been
1855 /// called on [`QueryManyIter`] before.
1856 pub fn sort_unstable_by<L: ReadOnlyQueryData + SingleEntityQueryData + 'w>(
1857 self,
1858 mut compare: impl FnMut(&L::Item<'_, '_>, &L::Item<'_, '_>) -> Ordering,
1859 ) -> QuerySortedManyIter<
1860 'w,
1861 's,
1862 D,
1863 F,
1864 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
1865 > {
1866 self.sort_impl::<L>(move |keyed_query| {
1867 keyed_query.sort_unstable_by(|(key_1, _), (key_2, _)| compare(key_1, key_2));
1868 })
1869 }
1870
1871 /// Sorts all query items into a new iterator with a key extraction function over the query lens.
1872 ///
1873 /// This sort is stable (i.e., does not reorder equal elements).
1874 ///
1875 /// This uses [`slice::sort_by_key`] internally.
1876 ///
1877 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
1878 /// This includes the allowed parameter type changes listed under [allowed transmutes].
1879 /// However, the lens uses the filter of the original query when present.
1880 ///
1881 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
1882 /// of query transmutes does not support nested queries.
1883 /// This restriction may be lifted in the future.
1884 ///
1885 /// The sort is not cached across system runs.
1886 ///
1887 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
1888 ///
1889 /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been
1890 /// called on [`QueryManyIter`] before.
1891 ///
1892 /// # Example
1893 /// ```
1894 /// # use bevy_ecs::prelude::*;
1895 /// # use std::ops::Deref;
1896 /// #
1897 /// # #[derive(Component)]
1898 /// # struct PartMarker;
1899 /// #
1900 /// # impl Deref for PartValue {
1901 /// # type Target = f32;
1902 /// #
1903 /// # fn deref(&self) -> &Self::Target {
1904 /// # &self.0
1905 /// # }
1906 /// # }
1907 /// #
1908 /// # let mut world = World::new();
1909 /// # let entity_list: Vec<Entity> = Vec::new();
1910 /// #
1911 /// #[derive(Component)]
1912 /// struct AvailableMarker;
1913 ///
1914 /// #[derive(Component, PartialEq, Eq, PartialOrd, Ord, Copy, Clone)]
1915 /// enum Rarity {
1916 /// Common,
1917 /// Rare,
1918 /// Epic,
1919 /// Legendary
1920 /// };
1921 ///
1922 /// #[derive(Component)]
1923 /// struct PartValue(f32);
1924 ///
1925 /// // We can sort with the internals of components that do not implement Ord.
1926 /// fn system_1(query: Query<(Entity, &PartValue)>) {
1927 /// # let entity_list: Vec<Entity> = Vec::new();
1928 /// // Sort by the sines of the part values.
1929 /// let parts: Vec<(Entity, &PartValue)> = query
1930 /// .iter_many(entity_list)
1931 /// .sort_by_key::<&PartValue, _>(|value| value.sin() as usize)
1932 /// .collect();
1933 /// }
1934 ///
1935 /// // We can define our own custom comparison functions over an EntityRef.
1936 /// fn system_2(query: Query<EntityRef, With<PartMarker>>) {
1937 /// # let entity_list: Vec<Entity> = Vec::new();
1938 /// // Sort by whether parts are available and their rarity.
1939 /// // We want the available legendaries to come first, so we reverse the iterator.
1940 /// let parts: Vec<EntityRef> = query.iter_many(entity_list)
1941 /// .sort_by_key::<EntityRef, _>(|entity_ref| {
1942 /// (
1943 /// entity_ref.contains::<AvailableMarker>(),
1944 // entity_ref.get::<Rarity>().copied()
1945 /// )
1946 /// })
1947 /// .rev()
1948 /// .collect();
1949 /// }
1950 /// # let mut schedule = Schedule::default();
1951 /// # schedule.add_systems((system_1, system_2));
1952 /// # schedule.run(&mut world);
1953 /// ```
1954 pub fn sort_by_key<L: ReadOnlyQueryData + SingleEntityQueryData + 'w, K>(
1955 self,
1956 mut f: impl FnMut(&L::Item<'_, '_>) -> K,
1957 ) -> QuerySortedManyIter<
1958 'w,
1959 's,
1960 D,
1961 F,
1962 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
1963 >
1964 where
1965 K: Ord,
1966 {
1967 self.sort_impl::<L>(move |keyed_query| keyed_query.sort_by_key(|(lens, _)| f(lens)))
1968 }
1969
1970 /// Sorts all query items into a new iterator with a key extraction function over the query lens.
1971 ///
1972 /// This sort is unstable (i.e., may reorder equal elements).
1973 ///
1974 /// This uses [`slice::sort_unstable_by_key`] internally.
1975 ///
1976 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
1977 /// This includes the allowed parameter type changes listed under [allowed transmutes].
1978 /// However, the lens uses the filter of the original query when present.
1979 ///
1980 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
1981 /// of query transmutes does not support nested queries.
1982 /// This restriction may be lifted in the future.
1983 ///
1984 /// The sort is not cached across system runs.
1985 ///
1986 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
1987 ///
1988 /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been
1989 /// called on [`QueryManyIter`] before.
1990 pub fn sort_unstable_by_key<L: ReadOnlyQueryData + SingleEntityQueryData + 'w, K>(
1991 self,
1992 mut f: impl FnMut(&L::Item<'_, '_>) -> K,
1993 ) -> QuerySortedManyIter<
1994 'w,
1995 's,
1996 D,
1997 F,
1998 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
1999 >
2000 where
2001 K: Ord,
2002 {
2003 self.sort_impl::<L>(move |keyed_query| {
2004 keyed_query.sort_unstable_by_key(|(lens, _)| f(lens));
2005 })
2006 }
2007
2008 /// Sort all query items into a new iterator with a key extraction function over the query lens.
2009 ///
2010 /// This sort is stable (i.e., does not reorder equal elements).
2011 ///
2012 /// This uses [`slice::sort_by_cached_key`] internally.
2013 ///
2014 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
2015 /// This includes the allowed parameter type changes listed under [allowed transmutes].
2016 /// However, the lens uses the filter of the original query when present.
2017 ///
2018 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
2019 /// of query transmutes does not support nested queries.
2020 /// This restriction may be lifted in the future.
2021 ///
2022 /// The sort is not cached across system runs.
2023 ///
2024 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
2025 ///
2026 /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been
2027 /// called on [`QueryManyIter`] before.
2028 pub fn sort_by_cached_key<L: ReadOnlyQueryData + SingleEntityQueryData + 'w, K>(
2029 self,
2030 mut f: impl FnMut(&L::Item<'_, '_>) -> K,
2031 ) -> QuerySortedManyIter<
2032 'w,
2033 's,
2034 D,
2035 F,
2036 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
2037 >
2038 where
2039 K: Ord,
2040 {
2041 self.sort_impl::<L>(move |keyed_query| keyed_query.sort_by_cached_key(|(lens, _)| f(lens)))
2042 }
2043
2044 /// Shared implementation for the various `sort` methods.
2045 /// This uses the lens to collect the items for sorting, but delegates the actual sorting to the provided closure.
2046 ///
2047 /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
2048 /// This includes the allowed parameter type changes listed under [allowed transmutes].
2049 /// However, the lens uses the filter of the original query when present.
2050 ///
2051 /// The lens needs to be a [`SingleEntityQueryData`] because the current implementation
2052 /// of query transmutes does not support nested queries.
2053 /// This restriction may be lifted in the future.
2054 ///
2055 /// The sort is not cached across system runs.
2056 ///
2057 /// [allowed transmutes]: crate::system::Query#allowed-transmutes
2058 ///
2059 /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been
2060 /// called on [`QueryManyIter`] before.
2061 fn sort_impl<L: ReadOnlyQueryData + SingleEntityQueryData + 'w>(
2062 self,
2063 f: impl FnOnce(&mut Vec<(L::Item<'_, '_>, NeutralOrd<Entity>)>),
2064 ) -> QuerySortedManyIter<
2065 'w,
2066 's,
2067 D,
2068 F,
2069 impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
2070 > {
2071 let world = self.world;
2072
2073 let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);
2074
2075 // SAFETY:
2076 // `self.world` has permission to access the required components.
2077 // The original query iter has not been iterated on, so no items are aliased from it.
2078 // `QueryIter::new` ensures `world` is the same one used to initialize `query_state`.
2079 let query_lens = unsafe { query_lens_state.query_unchecked_manual(world) }
2080 .iter_many_inner(self.entity_iter);
2081 let mut keyed_query: Vec<_> = query_lens
2082 .map(|(key, entity)| (key, NeutralOrd(entity)))
2083 .collect();
2084 f(&mut keyed_query);
2085 // Re-collect into a `Vec` to eagerly drop the lens items.
2086 // They must be dropped before `fetch_next` is called since they may alias
2087 // with other data items if there are duplicate entities in `entity_iter`.
2088 let entity_iter = keyed_query
2089 .into_iter()
2090 .map(|(.., entity)| entity.0)
2091 .collect::<Vec<_>>()
2092 .into_iter();
2093 // SAFETY:
2094 // `self.world` has permission to access the required components.
2095 // Each lens query item is dropped before the respective actual query item is accessed.
2096 unsafe {
2097 QuerySortedManyIter::new(
2098 world,
2099 self.query_state,
2100 entity_iter,
2101 world.last_change_tick(),
2102 world.change_tick(),
2103 )
2104 }
2105 }
2106}
2107
2108impl<'w, 's, D: QueryData, F: QueryFilter, I: DoubleEndedIterator<Item: EntityEquivalent>>
2109 QueryManyIter<'w, 's, D, F, I>
2110{
2111 /// Get the next result from the back of the query
2112 #[inline(always)]
2113 pub fn fetch_next_back(&mut self) -> Option<D::Item<'_, 's>> {
2114 // SAFETY:
2115 // All arguments stem from self.
2116 // We are limiting the returned reference to self,
2117 // making sure this method cannot be called multiple times without getting rid
2118 // of any previously returned unique references first, thus preventing aliasing.
2119 unsafe {
2120 Self::fetch_next_aliased_unchecked(
2121 self.entity_iter.by_ref().rev(),
2122 self.entities,
2123 self.tables,
2124 self.archetypes,
2125 &mut self.fetch,
2126 &mut self.filter,
2127 self.query_state,
2128 )
2129 .map(D::shrink)
2130 }
2131 }
2132}
2133
2134impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator<Item: EntityEquivalent>> Iterator
2135 for QueryManyIter<'w, 's, D, F, I>
2136{
2137 type Item = D::Item<'w, 's>;
2138
2139 #[inline(always)]
2140 fn next(&mut self) -> Option<Self::Item> {
2141 // SAFETY:
2142 // All arguments stem from self.
2143 // It is safe to alias for ReadOnlyWorldQuery.
2144 unsafe {
2145 Self::fetch_next_aliased_unchecked(
2146 &mut self.entity_iter,
2147 self.entities,
2148 self.tables,
2149 self.archetypes,
2150 &mut self.fetch,
2151 &mut self.filter,
2152 self.query_state,
2153 )
2154 }
2155 }
2156
2157 fn size_hint(&self) -> (usize, Option<usize>) {
2158 let (_, max_size) = self.entity_iter.size_hint();
2159 (0, max_size)
2160 }
2161}
2162
2163impl<
2164 'w,
2165 's,
2166 D: ReadOnlyQueryData,
2167 F: QueryFilter,
2168 I: DoubleEndedIterator<Item: EntityEquivalent>,
2169 > DoubleEndedIterator for QueryManyIter<'w, 's, D, F, I>
2170{
2171 #[inline(always)]
2172 fn next_back(&mut self) -> Option<Self::Item> {
2173 // SAFETY:
2174 // All arguments stem from self.
2175 // It is safe to alias for ReadOnlyWorldQuery.
2176 unsafe {
2177 Self::fetch_next_aliased_unchecked(
2178 self.entity_iter.by_ref().rev(),
2179 self.entities,
2180 self.tables,
2181 self.archetypes,
2182 &mut self.fetch,
2183 &mut self.filter,
2184 self.query_state,
2185 )
2186 }
2187 }
2188}
2189
2190// This is correct as [`QueryManyIter`] always returns `None` once exhausted.
2191impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator<Item: EntityEquivalent>>
2192 FusedIterator for QueryManyIter<'w, 's, D, F, I>
2193{
2194}
2195
2196// SAFETY: Fetching unique entities maintains uniqueness.
2197unsafe impl<'w, 's, F: QueryFilter, I: EntitySetIterator> EntitySetIterator
2198 for QueryManyIter<'w, 's, Entity, F, I>
2199{
2200}
2201
2202impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityEquivalent>> Debug
2203 for QueryManyIter<'w, 's, D, F, I>
2204{
2205 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
2206 f.debug_struct("QueryManyIter").finish()
2207 }
2208}
2209
2210/// An [`Iterator`] over the query items generated from an iterator of unique [`Entity`]s.
2211///
2212/// Items are returned in the order of the provided iterator.
2213/// Entities that don't match the query are skipped.
2214///
2215/// In contrast with [`QueryManyIter`], this allows for mutable iteration without a [`fetch_next`] method.
2216///
2217/// This struct is created by the [`iter_many_unique`] and [`iter_many_unique_mut`] methods on [`Query`].
2218///
2219/// [`fetch_next`]: QueryManyIter::fetch_next
2220/// [`iter_many_unique`]: crate::system::Query::iter_many
2221/// [`iter_many_unique_mut`]: crate::system::Query::iter_many_mut
2222/// [`Query`]: crate::system::Query
2223pub struct QueryManyUniqueIter<'w, 's, D: IterQueryData, F: QueryFilter, I: EntitySetIterator>(
2224 QueryManyIter<'w, 's, D, F, I>,
2225);
2226
2227impl<'w, 's, D: IterQueryData, F: QueryFilter, I: EntitySetIterator>
2228 QueryManyUniqueIter<'w, 's, D, F, I>
2229{
2230 /// # Safety
2231 /// - `world` must have permission to access any of the components registered in `query_state`.
2232 /// - `world` must be the same one used to initialize `query_state`.
2233 pub(crate) unsafe fn new<EntityList: EntitySet<IntoIter = I>>(
2234 world: UnsafeWorldCell<'w>,
2235 query_state: &'s QueryState<D, F>,
2236 entity_list: EntityList,
2237 last_run: Tick,
2238 this_run: Tick,
2239 ) -> QueryManyUniqueIter<'w, 's, D, F, I> {
2240 QueryManyUniqueIter(QueryManyIter::new(
2241 world,
2242 query_state,
2243 entity_list,
2244 last_run,
2245 this_run,
2246 ))
2247 }
2248}
2249
2250impl<'w, 's, D: IterQueryData, F: QueryFilter, I: EntitySetIterator> Iterator
2251 for QueryManyUniqueIter<'w, 's, D, F, I>
2252{
2253 type Item = D::Item<'w, 's>;
2254
2255 #[inline(always)]
2256 fn next(&mut self) -> Option<Self::Item> {
2257 // SAFETY:
2258 // - Entities are guaranteed to be unique, thus do not alias.
2259 // - `D: IterQueryData`
2260 unsafe {
2261 QueryManyIter::<'w, 's, D, F, I>::fetch_next_aliased_unchecked(
2262 &mut self.0.entity_iter,
2263 self.0.entities,
2264 self.0.tables,
2265 self.0.archetypes,
2266 &mut self.0.fetch,
2267 &mut self.0.filter,
2268 self.0.query_state,
2269 )
2270 }
2271 }
2272
2273 fn size_hint(&self) -> (usize, Option<usize>) {
2274 let (_, max_size) = self.0.entity_iter.size_hint();
2275 (0, max_size)
2276 }
2277}
2278
2279// This is correct as [`QueryManyIter`] always returns `None` once exhausted.
2280impl<'w, 's, D: IterQueryData, F: QueryFilter, I: EntitySetIterator> FusedIterator
2281 for QueryManyUniqueIter<'w, 's, D, F, I>
2282{
2283}
2284
2285// SAFETY: Fetching unique entities maintains uniqueness.
2286unsafe impl<'w, 's, F: QueryFilter, I: EntitySetIterator> EntitySetIterator
2287 for QueryManyUniqueIter<'w, 's, Entity, F, I>
2288{
2289}
2290
2291impl<'w, 's, D: IterQueryData, F: QueryFilter, I: EntitySetIterator> Debug
2292 for QueryManyUniqueIter<'w, 's, D, F, I>
2293{
2294 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
2295 f.debug_struct("QueryManyUniqueIter").finish()
2296 }
2297}
2298
2299/// An [`Iterator`] over sorted query results of a [`QueryManyIter`].
2300///
2301/// This struct is created by the [`sort`](QueryManyIter), [`sort_unstable`](QueryManyIter),
2302/// [`sort_by`](QueryManyIter), [`sort_unstable_by`](QueryManyIter), [`sort_by_key`](QueryManyIter),
2303/// [`sort_unstable_by_key`](QueryManyIter), and [`sort_by_cached_key`](QueryManyIter) methods of [`QueryManyIter`].
2304pub struct QuerySortedManyIter<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item = Entity>> {
2305 entity_iter: I,
2306 entities: &'w Entities,
2307 tables: &'w Tables,
2308 archetypes: &'w Archetypes,
2309 fetch: D::Fetch<'w>,
2310 query_state: &'s QueryState<D, F>,
2311}
2312
2313impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item = Entity>>
2314 QuerySortedManyIter<'w, 's, D, F, I>
2315{
2316 /// # Safety
2317 /// - `world` must have permission to access any of the components registered in `query_state`.
2318 /// - `world` must be the same one used to initialize `query_state`.
2319 /// - `entity_list` must only contain unique entities or be empty.
2320 pub(crate) unsafe fn new<EntityList: IntoIterator<IntoIter = I>>(
2321 world: UnsafeWorldCell<'w>,
2322 query_state: &'s QueryState<D, F>,
2323 entity_list: EntityList,
2324 last_run: Tick,
2325 this_run: Tick,
2326 ) -> QuerySortedManyIter<'w, 's, D, F, I> {
2327 let fetch = D::init_fetch(world, &query_state.fetch_state, last_run, this_run);
2328 QuerySortedManyIter {
2329 query_state,
2330 entities: world.entities(),
2331 archetypes: world.archetypes(),
2332 // SAFETY: We only access table data that has been registered in `query_state`.
2333 // This means `world` has permission to access the data we use.
2334 tables: &world.storages().tables,
2335 fetch,
2336 entity_iter: entity_list.into_iter(),
2337 }
2338 }
2339
2340 /// # Safety
2341 ///
2342 /// - `entity` must stem from `self.entity_iter`
2343 /// - If `D` does not impl `ReadOnlyQueryData`, then there must not be any other `Item`s alive for the current entity
2344 /// - If `D` does not impl `IterQueryData`, then there must not be any other `Item`s alive for *any* entity
2345 #[inline(always)]
2346 unsafe fn fetch_next_aliased_unchecked(&mut self, entity: Entity) -> Option<D::Item<'w, 's>> {
2347 let (location, archetype, table);
2348 // SAFETY:
2349 // `tables` and `archetypes` belong to the same world that the [`QueryIter`]
2350 // was initialized for.
2351 unsafe {
2352 location = self.entities.get_spawned(entity).debug_checked_unwrap();
2353 archetype = self
2354 .archetypes
2355 .get(location.archetype_id)
2356 .debug_checked_unwrap();
2357 table = self.tables.get(location.table_id).debug_checked_unwrap();
2358 }
2359
2360 // SAFETY: `archetype` is from the world that `fetch` was created for,
2361 // `fetch_state` is the state that `fetch` was initialized with
2362 unsafe {
2363 D::set_archetype(
2364 &mut self.fetch,
2365 &self.query_state.fetch_state,
2366 archetype,
2367 table,
2368 );
2369 }
2370
2371 // The entity list has already been filtered by the query lens, so we forego filtering here.
2372 // SAFETY:
2373 // - set_archetype was called prior, `location.archetype_row` is an archetype index in range of the current archetype
2374 // - Caller ensures there are no conflicting items alive
2375 unsafe {
2376 D::fetch(
2377 &self.query_state.fetch_state,
2378 &mut self.fetch,
2379 entity,
2380 location.table_row,
2381 )
2382 }
2383 }
2384
2385 /// Get the next result from the query
2386 #[inline(always)]
2387 pub fn fetch_next(&mut self) -> Option<D::Item<'_, 's>> {
2388 while let Some(entity) = self.entity_iter.next() {
2389 // SAFETY:
2390 // We have collected the entity_iter once to drop all internal lens query item
2391 // references.
2392 // We are limiting the returned reference to self,
2393 // making sure this method cannot be called multiple times without getting rid
2394 // of any previously returned unique references first, thus preventing aliasing.
2395 // `entity` is passed from `entity_iter` the first time.
2396 if let Some(item) = unsafe { self.fetch_next_aliased_unchecked(entity) } {
2397 return Some(D::shrink(item));
2398 }
2399 }
2400 None
2401 }
2402}
2403
2404impl<'w, 's, D: QueryData, F: QueryFilter, I: DoubleEndedIterator<Item = Entity>>
2405 QuerySortedManyIter<'w, 's, D, F, I>
2406{
2407 /// Get the next result from the query
2408 #[inline(always)]
2409 pub fn fetch_next_back(&mut self) -> Option<D::Item<'_, 's>> {
2410 while let Some(entity) = self.entity_iter.next_back() {
2411 // SAFETY:
2412 // We have collected the entity_iter once to drop all internal lens query item
2413 // references.
2414 // We are limiting the returned reference to self,
2415 // making sure this method cannot be called multiple times without getting rid
2416 // of any previously returned unique references first, thus preventing aliasing.
2417 // `entity` is passed from `entity_iter` the first time.
2418 if let Some(item) = unsafe { self.fetch_next_aliased_unchecked(entity) } {
2419 return Some(D::shrink(item));
2420 }
2421 }
2422 None
2423 }
2424}
2425
2426impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator<Item = Entity>> Iterator
2427 for QuerySortedManyIter<'w, 's, D, F, I>
2428{
2429 type Item = D::Item<'w, 's>;
2430
2431 #[inline(always)]
2432 fn next(&mut self) -> Option<Self::Item> {
2433 while let Some(entity) = self.entity_iter.next() {
2434 // SAFETY:
2435 // It is safe to alias for ReadOnlyWorldQuery.
2436 // `entity` is passed from `entity_iter` the first time.
2437 if let Some(item) = unsafe { self.fetch_next_aliased_unchecked(entity) } {
2438 return Some(D::shrink(item));
2439 }
2440 }
2441 None
2442 }
2443
2444 fn size_hint(&self) -> (usize, Option<usize>) {
2445 let (min_size, max_size) = self.entity_iter.size_hint();
2446 let archetype_query = D::IS_ARCHETYPAL;
2447 let min_size = if archetype_query { min_size } else { 0 };
2448 (min_size, max_size)
2449 }
2450}
2451
2452impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: DoubleEndedIterator<Item = Entity>>
2453 DoubleEndedIterator for QuerySortedManyIter<'w, 's, D, F, I>
2454{
2455 #[inline(always)]
2456 fn next_back(&mut self) -> Option<Self::Item> {
2457 while let Some(entity) = self.entity_iter.next_back() {
2458 // SAFETY:
2459 // It is safe to alias for ReadOnlyWorldQuery.
2460 // `entity` is passed from `entity_iter` the first time.
2461 if let Some(item) = unsafe { self.fetch_next_aliased_unchecked(entity) } {
2462 return Some(D::shrink(item));
2463 }
2464 }
2465 None
2466 }
2467}
2468
2469impl<
2470 D: ReadOnlyQueryData + ArchetypeQueryData,
2471 F: QueryFilter,
2472 I: ExactSizeIterator<Item = Entity>,
2473 > ExactSizeIterator for QuerySortedManyIter<'_, '_, D, F, I>
2474{
2475}
2476
2477impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item = Entity>> Debug
2478 for QuerySortedManyIter<'w, 's, D, F, I>
2479{
2480 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
2481 f.debug_struct("QuerySortedManyIter").finish()
2482 }
2483}
2484
2485/// An iterator over `K`-sized combinations of query items without repetition.
2486///
2487/// A combination is an arrangement of a collection of items where order does not matter.
2488///
2489/// `K` is the number of items that make up each subset, and the number of items returned by the iterator.
2490/// `N` is the number of total entities output by the query.
2491///
2492/// For example, given the list [1, 2, 3, 4], where `K` is 2, the combinations without repeats are
2493/// [1, 2], [1, 3], [1, 4], [2, 3], [2, 4], [3, 4].
2494/// And in this case, `N` would be defined as 4 since the size of the input list is 4.
2495///
2496/// The number of combinations depend on how `K` relates to the number of entities matching the [`Query`]:
2497/// - if `K = N`, only one combination exists,
2498/// - if `K < N`, there are <sub>N</sub>C<sub>K</sub> combinations (see the [performance section] of `Query`),
2499/// - if `K > N`, there are no combinations.
2500///
2501/// The output combination is not guaranteed to have any order of iteration.
2502///
2503/// # Usage
2504///
2505/// This type is returned by calling [`Query::iter_combinations`] or [`Query::iter_combinations_mut`].
2506///
2507/// It implements [`Iterator`] only if it iterates over read-only query items ([learn more]).
2508///
2509/// In the case of mutable query items, it can be iterated by calling [`fetch_next`] in a `while let` loop.
2510///
2511/// # Examples
2512///
2513/// The following example shows how to traverse the iterator when the query items are read-only.
2514///
2515/// ```
2516/// # use bevy_ecs::prelude::*;
2517/// # #[derive(Component)]
2518/// # struct ComponentA;
2519/// #
2520/// fn some_system(query: Query<&ComponentA>) {
2521/// for [a1, a2] in query.iter_combinations() {
2522/// // ...
2523/// }
2524/// }
2525/// ```
2526///
2527/// 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.
2528///
2529/// ```
2530/// # use bevy_ecs::prelude::*;
2531/// # #[derive(Component)]
2532/// # struct ComponentA;
2533/// #
2534/// fn some_system(mut query: Query<&mut ComponentA>) {
2535/// let mut combinations = query.iter_combinations_mut();
2536/// while let Some([a1, a2]) = combinations.fetch_next() {
2537/// // ...
2538/// }
2539/// }
2540/// ```
2541///
2542/// [`fetch_next`]: Self::fetch_next
2543/// [learn more]: Self#impl-Iterator
2544/// [performance section]: crate::system::Query#performance
2545/// [`Query`]: crate::system::Query
2546/// [`Query::iter_combinations`]: crate::system::Query::iter_combinations
2547/// [`Query::iter_combinations_mut`]: crate::system::Query::iter_combinations_mut
2548pub struct QueryCombinationIter<'w, 's, D: IterQueryData, F: QueryFilter, const K: usize> {
2549 tables: &'w Tables,
2550 archetypes: &'w Archetypes,
2551 query_state: &'s QueryState<D, F>,
2552 cursors: [QueryIterationCursor<'w, 's, D, F>; K],
2553}
2554
2555impl<'w, 's, D: IterQueryData, F: QueryFilter, const K: usize>
2556 QueryCombinationIter<'w, 's, D, F, K>
2557{
2558 /// # Safety
2559 /// - `world` must have permission to access any of the components registered in `query_state`.
2560 /// - `world` must be the same one used to initialize `query_state`.
2561 pub(crate) unsafe fn new(
2562 world: UnsafeWorldCell<'w>,
2563 query_state: &'s QueryState<D, F>,
2564 last_run: Tick,
2565 this_run: Tick,
2566 ) -> Self {
2567 assert!(K != 0, "K should not equal to zero");
2568 // Initialize array with cursors.
2569 // There is no FromIterator on arrays, so instead initialize it manually with MaybeUninit
2570
2571 let mut array: MaybeUninit<[QueryIterationCursor<'w, 's, D, F>; K]> = MaybeUninit::uninit();
2572 let ptr = array
2573 .as_mut_ptr()
2574 .cast::<QueryIterationCursor<'w, 's, D, F>>();
2575 ptr.write(QueryIterationCursor::init(
2576 world,
2577 query_state,
2578 last_run,
2579 this_run,
2580 ));
2581 for slot in (1..K).map(|offset| ptr.add(offset)) {
2582 slot.write(QueryIterationCursor::init_empty(
2583 world,
2584 query_state,
2585 last_run,
2586 this_run,
2587 ));
2588 }
2589
2590 QueryCombinationIter {
2591 query_state,
2592 // SAFETY: We only access table data that has been registered in `query_state`.
2593 tables: unsafe { &world.storages().tables },
2594 archetypes: world.archetypes(),
2595 cursors: array.assume_init(),
2596 }
2597 }
2598
2599 /// # Safety
2600 /// The lifetime here is not restrictive enough for Fetch with &mut access,
2601 /// as calling `fetch_next_aliased_unchecked` multiple times can produce multiple
2602 /// references to the same component, leading to unique reference aliasing.
2603 /// .
2604 /// It is always safe for shared access.
2605 #[inline]
2606 unsafe fn fetch_next_aliased_unchecked(&mut self) -> Option<[D::Item<'w, 's>; K]> {
2607 // PERF: can speed up the following code using `cursor.remaining()` instead of `next_item.is_none()`
2608 // when D::IS_ARCHETYPAL && F::IS_ARCHETYPAL
2609 //
2610 // let `i` be the index of `c`, the last cursor in `self.cursors` that
2611 // returns `K-i` or more elements.
2612 // Make cursor in index `j` for all `j` in `[i, K)` a copy of `c` advanced `j-i+1` times.
2613 // If no such `c` exists, return `None`
2614 'outer: for i in (0..K).rev() {
2615 match self.cursors[i].next(self.tables, self.archetypes, self.query_state) {
2616 Some(_) => {
2617 for j in (i + 1)..K {
2618 self.cursors[j] = self.cursors[j - 1].clone();
2619 match self.cursors[j].next(self.tables, self.archetypes, self.query_state) {
2620 Some(_) => {}
2621 None if i > 0 => continue 'outer,
2622 None => return None,
2623 }
2624 }
2625 break;
2626 }
2627 None if i > 0 => continue,
2628 None => return None,
2629 }
2630 }
2631
2632 let mut values = MaybeUninit::<[D::Item<'w, 's>; K]>::uninit();
2633
2634 let ptr = values.as_mut_ptr().cast::<D::Item<'w, 's>>();
2635 for (offset, cursor) in self.cursors.iter_mut().enumerate() {
2636 ptr.add(offset)
2637 .write(cursor.peek_last(self.query_state).unwrap());
2638 }
2639
2640 Some(values.assume_init())
2641 }
2642
2643 /// Get next combination of queried components
2644 #[inline]
2645 pub fn fetch_next(&mut self) -> Option<[D::Item<'_, 's>; K]> {
2646 // SAFETY: we are limiting the returned reference to self,
2647 // making sure this method cannot be called multiple times without getting rid
2648 // of any previously returned unique references first, thus preventing aliasing.
2649 unsafe {
2650 self.fetch_next_aliased_unchecked()
2651 .map(|array| array.map(D::shrink))
2652 }
2653 }
2654}
2655
2656// Iterator type is intentionally implemented only for read-only access.
2657// Doing so for mutable references would be unsound, because calling `next`
2658// multiple times would allow multiple owned references to the same data to exist.
2659impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, const K: usize> Iterator
2660 for QueryCombinationIter<'w, 's, D, F, K>
2661{
2662 type Item = [D::Item<'w, 's>; K];
2663
2664 #[inline]
2665 fn next(&mut self) -> Option<Self::Item> {
2666 // Safety: it is safe to alias for ReadOnlyWorldQuery
2667 unsafe { QueryCombinationIter::fetch_next_aliased_unchecked(self) }
2668 }
2669
2670 fn size_hint(&self) -> (usize, Option<usize>) {
2671 // binomial coefficient: (n ; k) = n! / k!(n-k)! = (n*n-1*...*n-k+1) / k!
2672 // See https://en.wikipedia.org/wiki/Binomial_coefficient
2673 // See https://blog.plover.com/math/choose.html for implementation
2674 // It was chosen to reduce overflow potential.
2675 fn choose(n: usize, k: usize) -> Option<usize> {
2676 if k > n || n == 0 {
2677 return Some(0);
2678 }
2679 let k = k.min(n - k);
2680 let ks = 1..=k;
2681 let ns = (n - k + 1..=n).rev();
2682 ks.zip(ns)
2683 .try_fold(1_usize, |acc, (k, n)| Some(acc.checked_mul(n)? / k))
2684 }
2685 // sum_i=0..k choose(cursors[i].remaining, k-i)
2686 let max_combinations = self
2687 .cursors
2688 .iter()
2689 .enumerate()
2690 .try_fold(0, |acc, (i, cursor)| {
2691 let n = cursor.max_remaining(self.tables, self.archetypes);
2692 Some(acc + choose(n as usize, K - i)?)
2693 });
2694
2695 let archetype_query = D::IS_ARCHETYPAL && F::IS_ARCHETYPAL;
2696 let known_max = max_combinations.unwrap_or(usize::MAX);
2697 let min_combinations = if archetype_query { known_max } else { 0 };
2698 (min_combinations, max_combinations)
2699 }
2700}
2701
2702impl<'w, 's, D: ArchetypeQueryData + IterQueryData, F: ArchetypeFilter> ExactSizeIterator
2703 for QueryIter<'w, 's, D, F>
2704{
2705 fn len(&self) -> usize {
2706 self.size_hint().0
2707 }
2708}
2709
2710// This is correct as [`QueryCombinationIter`] always returns `None` once exhausted.
2711impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, const K: usize> FusedIterator
2712 for QueryCombinationIter<'w, 's, D, F, K>
2713{
2714}
2715
2716impl<'w, 's, D: IterQueryData, F: QueryFilter, const K: usize> Debug
2717 for QueryCombinationIter<'w, 's, D, F, K>
2718{
2719 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
2720 f.debug_struct("QueryCombinationIter").finish()
2721 }
2722}
2723
2724struct QueryIterationCursor<'w, 's, D: QueryData, F: QueryFilter> {
2725 // whether the query iteration is dense or not. Mirrors QueryState's `is_dense` field.
2726 is_dense: bool,
2727 storage_id_iter: core::slice::Iter<'s, StorageId>,
2728 table_entities: &'w [Entity],
2729 archetype_entities: &'w [ArchetypeEntity],
2730 fetch: D::Fetch<'w>,
2731 filter: F::Fetch<'w>,
2732 // length of the table or length of the archetype, depending on whether both `D`'s and `F`'s fetches are dense
2733 current_len: u32,
2734 // either table row or archetype index, depending on whether both `D`'s and `F`'s fetches are dense
2735 current_row: u32,
2736}
2737
2738impl<D: QueryData, F: QueryFilter> Clone for QueryIterationCursor<'_, '_, D, F> {
2739 fn clone(&self) -> Self {
2740 Self {
2741 is_dense: self.is_dense,
2742 storage_id_iter: self.storage_id_iter.clone(),
2743 table_entities: self.table_entities,
2744 archetype_entities: self.archetype_entities,
2745 fetch: self.fetch.clone(),
2746 filter: self.filter.clone(),
2747 current_len: self.current_len,
2748 current_row: self.current_row,
2749 }
2750 }
2751}
2752
2753impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> {
2754 /// # Safety
2755 /// - `world` must have permission to access any of the components registered in `query_state`.
2756 /// - `world` must be the same one used to initialize `query_state`.
2757 unsafe fn init_empty(
2758 world: UnsafeWorldCell<'w>,
2759 query_state: &'s QueryState<D, F>,
2760 last_run: Tick,
2761 this_run: Tick,
2762 ) -> Self {
2763 QueryIterationCursor {
2764 storage_id_iter: [].iter(),
2765 ..Self::init(world, query_state, last_run, this_run)
2766 }
2767 }
2768
2769 /// # Safety
2770 /// - `world` must have permission to access any of the components registered in `query_state`.
2771 /// - `world` must be the same one used to initialize `query_state`.
2772 unsafe fn init(
2773 world: UnsafeWorldCell<'w>,
2774 query_state: &'s QueryState<D, F>,
2775 last_run: Tick,
2776 this_run: Tick,
2777 ) -> Self {
2778 let fetch = D::init_fetch(world, &query_state.fetch_state, last_run, this_run);
2779 let filter = F::init_fetch(world, &query_state.filter_state, last_run, this_run);
2780 QueryIterationCursor {
2781 fetch,
2782 filter,
2783 table_entities: &[],
2784 archetype_entities: &[],
2785 storage_id_iter: query_state.matched_storage_ids.iter(),
2786 is_dense: query_state.is_dense,
2787 current_len: 0,
2788 current_row: 0,
2789 }
2790 }
2791
2792 fn reborrow(&mut self) -> QueryIterationCursor<'_, 's, D, F> {
2793 QueryIterationCursor {
2794 is_dense: self.is_dense,
2795 fetch: D::shrink_fetch(self.fetch.clone()),
2796 filter: F::shrink_fetch(self.filter.clone()),
2797 table_entities: self.table_entities,
2798 archetype_entities: self.archetype_entities,
2799 storage_id_iter: self.storage_id_iter.clone(),
2800 current_len: self.current_len,
2801 current_row: self.current_row,
2802 }
2803 }
2804
2805 /// Retrieve item returned from most recent `next` call again.
2806 ///
2807 /// # Safety
2808 /// The result of `next` and any previous calls to `peek_last` with this row must have been
2809 /// dropped to prevent aliasing mutable references.
2810 #[inline]
2811 unsafe fn peek_last(&mut self, query_state: &'s QueryState<D, F>) -> Option<D::Item<'w, 's>>
2812 where
2813 D: IterQueryData,
2814 {
2815 if self.current_row > 0 {
2816 let index = self.current_row - 1;
2817 if self.is_dense {
2818 // SAFETY: This must have been called previously in `next` as `current_row > 0`
2819 let entity = unsafe { self.table_entities.get_unchecked(index as usize) };
2820 // SAFETY:
2821 // - `set_table` must have been called previously either in `next` or before it.
2822 // - `*entity` and `index` are in the current table.
2823 // - `D: IterQueryData`
2824 unsafe {
2825 D::fetch(
2826 &query_state.fetch_state,
2827 &mut self.fetch,
2828 *entity,
2829 // SAFETY: This is from an exclusive range, so it can't be max.
2830 TableRow::new(NonMaxU32::new_unchecked(index)),
2831 )
2832 }
2833 } else {
2834 // SAFETY: This must have been called previously in `next` as `current_row > 0`
2835 let archetype_entity =
2836 unsafe { self.archetype_entities.get_unchecked(index as usize) };
2837 // SAFETY:
2838 // - `set_archetype` must have been called previously either in `next` or before it.
2839 // - `archetype_entity.id()` and `archetype_entity.table_row()` are in the current archetype.
2840 // - `D: IterQueryData`
2841 unsafe {
2842 D::fetch(
2843 &query_state.fetch_state,
2844 &mut self.fetch,
2845 archetype_entity.id(),
2846 archetype_entity.table_row(),
2847 )
2848 }
2849 }
2850 } else {
2851 None
2852 }
2853 }
2854
2855 /// How many values will this cursor return at most?
2856 ///
2857 /// Note that if `D::IS_ARCHETYPAL && F::IS_ARCHETYPAL`, the return value
2858 /// will be **the exact count of remaining values**.
2859 fn max_remaining(&self, tables: &'w Tables, archetypes: &'w Archetypes) -> u32 {
2860 let ids = self.storage_id_iter.clone();
2861 let remaining_matched: u32 = if self.is_dense {
2862 // SAFETY: The if check ensures that storage_id_iter stores TableIds
2863 unsafe { ids.map(|id| tables[id.table_id].entity_count()).sum() }
2864 } else {
2865 // SAFETY: The if check ensures that storage_id_iter stores ArchetypeIds
2866 unsafe { ids.map(|id| archetypes[id.archetype_id].len()).sum() }
2867 };
2868 remaining_matched + self.current_len - self.current_row
2869 }
2870
2871 // NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
2872 // QueryIter, QueryIterationCursor, QuerySortedIter, QueryManyIter, QuerySortedManyIter, QueryCombinationIter,
2873 // QueryState::par_fold_init_unchecked_manual, QueryState::par_many_fold_init_unchecked_manual,
2874 // QueryState::par_many_unique_fold_init_unchecked_manual, QueryContiguousIter::next
2875 /// # Safety
2876 /// - `tables` and `archetypes` must belong to the same world that the [`QueryIterationCursor`]
2877 /// was initialized for.
2878 /// - `query_state` must be the same [`QueryState`] that was passed to `init` or `init_empty`.
2879 /// - If `D` does not impl `ReadOnlyQueryData`, then there must not be any other `Item`s alive for the current entity
2880 /// - If `D` does not impl `IterQueryData`, then there must not be any other `Item`s alive for *any* entity
2881 #[inline(always)]
2882 unsafe fn next(
2883 &mut self,
2884 tables: &'w Tables,
2885 archetypes: &'w Archetypes,
2886 query_state: &'s QueryState<D, F>,
2887 ) -> Option<D::Item<'w, 's>> {
2888 if self.is_dense {
2889 // NOTE: if you are changing this branch you would probably have to change
2890 // QueryContiguousIter::next as well
2891 loop {
2892 // we are on the beginning of the query, or finished processing a table, so skip to the next
2893 if self.current_row == self.current_len {
2894 let table_id = self.storage_id_iter.next()?.table_id;
2895 let table = tables.get(table_id).debug_checked_unwrap();
2896 if table.is_empty() {
2897 continue;
2898 }
2899 // SAFETY: `table` is from the world that `fetch/filter` were created for,
2900 // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
2901 unsafe {
2902 D::set_table(&mut self.fetch, &query_state.fetch_state, table);
2903 F::set_table(&mut self.filter, &query_state.filter_state, table);
2904 }
2905 self.table_entities = table.entities();
2906 self.current_len = table.entity_count();
2907 self.current_row = 0;
2908 }
2909
2910 // SAFETY: set_table was called prior.
2911 // `current_row` is a table row in range of the current table, because if it was not, then the above would have been executed.
2912 let entity =
2913 unsafe { self.table_entities.get_unchecked(self.current_row as usize) };
2914 // SAFETY: The row is less than the u32 len, so it must not be max.
2915 let row = unsafe { TableRow::new(NonMaxU32::new_unchecked(self.current_row)) };
2916 self.current_row += 1;
2917
2918 if !F::filter_fetch(&query_state.filter_state, &mut self.filter, *entity, row) {
2919 continue;
2920 }
2921
2922 // SAFETY:
2923 // - set_table was called prior.
2924 // - `current_row` must be a table row in range of the current table,
2925 // because if it was not, then the above would have been executed.
2926 // - fetch is only called once for each `entity`.
2927 // - caller ensures no conflicting `Item`s are alive
2928 let item =
2929 unsafe { D::fetch(&query_state.fetch_state, &mut self.fetch, *entity, row) };
2930 if let Some(item) = item {
2931 return Some(item);
2932 }
2933 }
2934 } else {
2935 loop {
2936 if self.current_row == self.current_len {
2937 let archetype_id = self.storage_id_iter.next()?.archetype_id;
2938 let archetype = archetypes.get(archetype_id).debug_checked_unwrap();
2939 if archetype.is_empty() {
2940 continue;
2941 }
2942 let table = tables.get(archetype.table_id()).debug_checked_unwrap();
2943 // SAFETY: `archetype` and `tables` are from the world that `fetch/filter` were created for,
2944 // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
2945 unsafe {
2946 D::set_archetype(
2947 &mut self.fetch,
2948 &query_state.fetch_state,
2949 archetype,
2950 table,
2951 );
2952 F::set_archetype(
2953 &mut self.filter,
2954 &query_state.filter_state,
2955 archetype,
2956 table,
2957 );
2958 }
2959 self.archetype_entities = archetype.entities();
2960 self.current_len = archetype.len();
2961 self.current_row = 0;
2962 }
2963
2964 // SAFETY: set_archetype was called prior.
2965 // `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.
2966 let archetype_entity = unsafe {
2967 self.archetype_entities
2968 .get_unchecked(self.current_row as usize)
2969 };
2970 self.current_row += 1;
2971
2972 if !F::filter_fetch(
2973 &query_state.filter_state,
2974 &mut self.filter,
2975 archetype_entity.id(),
2976 archetype_entity.table_row(),
2977 ) {
2978 continue;
2979 }
2980
2981 // SAFETY:
2982 // - set_archetype was called prior.
2983 // - `current_row` must be an archetype index row in range of the current archetype,
2984 // because if it was not, then the if above would have been executed.
2985 // - fetch is only called once for each `archetype_entity`.
2986 // - caller ensures no conflicting `Item`s are alive
2987 let item = unsafe {
2988 D::fetch(
2989 &query_state.fetch_state,
2990 &mut self.fetch,
2991 archetype_entity.id(),
2992 archetype_entity.table_row(),
2993 )
2994 };
2995 if let Some(item) = item {
2996 return Some(item);
2997 }
2998 }
2999 }
3000 }
3001}
3002
3003// A wrapper struct that gives its data a neutral ordering.
3004#[derive(Copy, Clone)]
3005struct NeutralOrd<T>(T);
3006
3007impl<T> PartialEq for NeutralOrd<T> {
3008 fn eq(&self, _other: &Self) -> bool {
3009 true
3010 }
3011}
3012
3013impl<T> Eq for NeutralOrd<T> {}
3014
3015#[expect(
3016 clippy::non_canonical_partial_ord_impl,
3017 reason = "`PartialOrd` and `Ord` on this struct must only ever return `Ordering::Equal`, so we prefer clarity"
3018)]
3019impl<T> PartialOrd for NeutralOrd<T> {
3020 fn partial_cmp(&self, _other: &Self) -> Option<Ordering> {
3021 Some(Ordering::Equal)
3022 }
3023}
3024
3025impl<T> Ord for NeutralOrd<T> {
3026 fn cmp(&self, _other: &Self) -> Ordering {
3027 Ordering::Equal
3028 }
3029}
3030
3031#[cfg(test)]
3032#[expect(clippy::print_stdout, reason = "Allowed in tests.")]
3033mod tests {
3034 use alloc::vec::Vec;
3035 use std::println;
3036
3037 use crate::component::Component;
3038 use crate::entity::Entity;
3039 use crate::prelude::{With, World};
3040
3041 #[derive(Component, Debug, PartialEq, PartialOrd, Clone, Copy)]
3042 struct A(f32);
3043 #[derive(Component, Debug, Eq, PartialEq, Clone, Copy)]
3044 #[component(storage = "SparseSet")]
3045 struct Sparse(usize);
3046
3047 #[derive(Component)]
3048 struct Marker;
3049
3050 #[test]
3051 #[cfg_attr(miri, ignore = "This test takes ~70s on CI")]
3052 fn query_iter_sorts() {
3053 let mut world = World::new();
3054 for i in 0..100 {
3055 world.spawn((A(i as f32), Marker));
3056 world.spawn((A(i as f32), Sparse(i), Marker));
3057 world.spawn((Sparse(i), Marker));
3058 }
3059
3060 let mut query = world.query_filtered::<Entity, With<Marker>>();
3061
3062 let sort = query.iter(&world).sort::<Entity>().collect::<Vec<_>>();
3063 assert_eq!(sort.len(), 300);
3064
3065 let sort_unstable = query
3066 .iter(&world)
3067 .sort_unstable::<Entity>()
3068 .collect::<Vec<_>>();
3069
3070 let sort_by = query
3071 .iter(&world)
3072 .sort_by::<Entity>(Ord::cmp)
3073 .collect::<Vec<_>>();
3074
3075 let sort_unstable_by = query
3076 .iter(&world)
3077 .sort_unstable_by::<Entity>(Ord::cmp)
3078 .collect::<Vec<_>>();
3079
3080 let sort_by_key = query
3081 .iter(&world)
3082 .sort_by_key::<Entity, _>(|&e| e)
3083 .collect::<Vec<_>>();
3084
3085 let sort_unstable_by_key = query
3086 .iter(&world)
3087 .sort_unstable_by_key::<Entity, _>(|&e| e)
3088 .collect::<Vec<_>>();
3089
3090 let sort_by_cached_key = query
3091 .iter(&world)
3092 .sort_by_cached_key::<Entity, _>(|&e| e)
3093 .collect::<Vec<_>>();
3094
3095 let mut sort_v2 = query.iter(&world).collect::<Vec<_>>();
3096 sort_v2.sort();
3097
3098 let mut sort_unstable_v2 = query.iter(&world).collect::<Vec<_>>();
3099 sort_unstable_v2.sort_unstable();
3100
3101 let mut sort_by_v2 = query.iter(&world).collect::<Vec<_>>();
3102 sort_by_v2.sort_by(Ord::cmp);
3103
3104 let mut sort_unstable_by_v2 = query.iter(&world).collect::<Vec<_>>();
3105 sort_unstable_by_v2.sort_unstable_by(Ord::cmp);
3106
3107 let mut sort_by_key_v2 = query.iter(&world).collect::<Vec<_>>();
3108 sort_by_key_v2.sort_by_key(|&e| e);
3109
3110 let mut sort_unstable_by_key_v2 = query.iter(&world).collect::<Vec<_>>();
3111 sort_unstable_by_key_v2.sort_unstable_by_key(|&e| e);
3112
3113 let mut sort_by_cached_key_v2 = query.iter(&world).collect::<Vec<_>>();
3114 sort_by_cached_key_v2.sort_by_cached_key(|&e| e);
3115
3116 assert_eq!(sort, sort_v2);
3117 assert_eq!(sort_unstable, sort_unstable_v2);
3118 assert_eq!(sort_by, sort_by_v2);
3119 assert_eq!(sort_unstable_by, sort_unstable_by_v2);
3120 assert_eq!(sort_by_key, sort_by_key_v2);
3121 assert_eq!(sort_unstable_by_key, sort_unstable_by_key_v2);
3122 assert_eq!(sort_by_cached_key, sort_by_cached_key_v2);
3123 }
3124
3125 #[test]
3126 #[should_panic]
3127 fn query_iter_sort_after_next() {
3128 let mut world = World::new();
3129 world.spawn((A(0.),));
3130 world.spawn((A(1.1),));
3131 world.spawn((A(2.22),));
3132
3133 {
3134 let mut query = world.query::<&A>();
3135 let mut iter = query.iter(&world);
3136 println!(
3137 "archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
3138 iter.cursor.archetype_entities.len(),
3139 iter.cursor.table_entities.len(),
3140 iter.cursor.current_len,
3141 iter.cursor.current_row
3142 );
3143 _ = iter.next();
3144 println!(
3145 "archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
3146 iter.cursor.archetype_entities.len(),
3147 iter.cursor.table_entities.len(),
3148 iter.cursor.current_len,
3149 iter.cursor.current_row
3150 );
3151 println!("{}", iter.sort::<Entity>().len());
3152 }
3153 }
3154
3155 #[test]
3156 #[should_panic]
3157 fn query_iter_sort_after_next_dense() {
3158 let mut world = World::new();
3159 world.spawn((Sparse(11),));
3160 world.spawn((Sparse(22),));
3161 world.spawn((Sparse(33),));
3162
3163 {
3164 let mut query = world.query::<&Sparse>();
3165 let mut iter = query.iter(&world);
3166 println!(
3167 "before_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
3168 iter.cursor.archetype_entities.len(),
3169 iter.cursor.table_entities.len(),
3170 iter.cursor.current_len,
3171 iter.cursor.current_row
3172 );
3173 _ = iter.next();
3174 println!(
3175 "after_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
3176 iter.cursor.archetype_entities.len(),
3177 iter.cursor.table_entities.len(),
3178 iter.cursor.current_len,
3179 iter.cursor.current_row
3180 );
3181 println!("{}", iter.sort::<Entity>().len());
3182 }
3183 }
3184
3185 #[test]
3186 fn empty_query_iter_sort_after_next_does_not_panic() {
3187 let mut world = World::new();
3188 {
3189 let mut query = world.query::<(&A, &Sparse)>();
3190 let mut iter = query.iter(&world);
3191 println!(
3192 "before_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
3193 iter.cursor.archetype_entities.len(),
3194 iter.cursor.table_entities.len(),
3195 iter.cursor.current_len,
3196 iter.cursor.current_row
3197 );
3198 _ = iter.next();
3199 println!(
3200 "after_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
3201 iter.cursor.archetype_entities.len(),
3202 iter.cursor.table_entities.len(),
3203 iter.cursor.current_len,
3204 iter.cursor.current_row
3205 );
3206 println!("{}", iter.sort::<Entity>().len());
3207 }
3208 }
3209
3210 #[test]
3211 fn query_iter_cursor_state_non_empty_after_next() {
3212 let mut world = World::new();
3213 world.spawn((A(0.), Sparse(11)));
3214 world.spawn((A(1.1), Sparse(22)));
3215 world.spawn((A(2.22), Sparse(33)));
3216 {
3217 let mut query = world.query::<(&A, &Sparse)>();
3218 let mut iter = query.iter(&world);
3219 println!(
3220 "before_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
3221 iter.cursor.archetype_entities.len(),
3222 iter.cursor.table_entities.len(),
3223 iter.cursor.current_len,
3224 iter.cursor.current_row
3225 );
3226 assert!(iter.cursor.table_entities.len() | iter.cursor.archetype_entities.len() == 0);
3227 _ = iter.next();
3228 println!(
3229 "after_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}",
3230 iter.cursor.archetype_entities.len(),
3231 iter.cursor.table_entities.len(),
3232 iter.cursor.current_len,
3233 iter.cursor.current_row
3234 );
3235 assert!(
3236 (
3237 iter.cursor.table_entities.len(),
3238 iter.cursor.archetype_entities.len()
3239 ) != (0, 0)
3240 );
3241 }
3242 }
3243
3244 #[test]
3245 fn query_iter_many_sorts() {
3246 let mut world = World::new();
3247
3248 let entity_list: &Vec<_> = &world
3249 .spawn_batch([A(0.), A(1.), A(2.), A(3.), A(4.)])
3250 .collect();
3251
3252 let mut query = world.query::<Entity>();
3253
3254 let sort = query
3255 .iter_many(&world, entity_list)
3256 .sort::<Entity>()
3257 .collect::<Vec<_>>();
3258
3259 let sort_unstable = query
3260 .iter_many(&world, entity_list)
3261 .sort_unstable::<Entity>()
3262 .collect::<Vec<_>>();
3263
3264 let sort_by = query
3265 .iter_many(&world, entity_list)
3266 .sort_by::<Entity>(Ord::cmp)
3267 .collect::<Vec<_>>();
3268
3269 let sort_unstable_by = query
3270 .iter_many(&world, entity_list)
3271 .sort_unstable_by::<Entity>(Ord::cmp)
3272 .collect::<Vec<_>>();
3273
3274 let sort_by_key = query
3275 .iter_many(&world, entity_list)
3276 .sort_by_key::<Entity, _>(|&e| e)
3277 .collect::<Vec<_>>();
3278
3279 let sort_unstable_by_key = query
3280 .iter_many(&world, entity_list)
3281 .sort_unstable_by_key::<Entity, _>(|&e| e)
3282 .collect::<Vec<_>>();
3283
3284 let sort_by_cached_key = query
3285 .iter_many(&world, entity_list)
3286 .sort_by_cached_key::<Entity, _>(|&e| e)
3287 .collect::<Vec<_>>();
3288
3289 let mut sort_v2 = query.iter_many(&world, entity_list).collect::<Vec<_>>();
3290 sort_v2.sort();
3291
3292 let mut sort_unstable_v2 = query.iter_many(&world, entity_list).collect::<Vec<_>>();
3293 sort_unstable_v2.sort_unstable();
3294
3295 let mut sort_by_v2 = query.iter_many(&world, entity_list).collect::<Vec<_>>();
3296 sort_by_v2.sort_by(Ord::cmp);
3297
3298 let mut sort_unstable_by_v2 = query.iter_many(&world, entity_list).collect::<Vec<_>>();
3299 sort_unstable_by_v2.sort_unstable_by(Ord::cmp);
3300
3301 let mut sort_by_key_v2 = query.iter_many(&world, entity_list).collect::<Vec<_>>();
3302 sort_by_key_v2.sort_by_key(|&e| e);
3303
3304 let mut sort_unstable_by_key_v2 = query.iter_many(&world, entity_list).collect::<Vec<_>>();
3305 sort_unstable_by_key_v2.sort_unstable_by_key(|&e| e);
3306
3307 let mut sort_by_cached_key_v2 = query.iter_many(&world, entity_list).collect::<Vec<_>>();
3308 sort_by_cached_key_v2.sort_by_cached_key(|&e| e);
3309
3310 assert_eq!(sort, sort_v2);
3311 assert_eq!(sort_unstable, sort_unstable_v2);
3312 assert_eq!(sort_by, sort_by_v2);
3313 assert_eq!(sort_unstable_by, sort_unstable_by_v2);
3314 assert_eq!(sort_by_key, sort_by_key_v2);
3315 assert_eq!(sort_unstable_by_key, sort_unstable_by_key_v2);
3316 assert_eq!(sort_by_cached_key, sort_by_cached_key_v2);
3317 }
3318
3319 #[test]
3320 fn query_iter_many_sort_doesnt_panic_after_next() {
3321 let mut world = World::new();
3322
3323 let entity_list: &Vec<_> = &world
3324 .spawn_batch([A(0.), A(1.), A(2.), A(3.), A(4.)])
3325 .collect();
3326
3327 let mut query = world.query::<Entity>();
3328 let mut iter = query.iter_many(&world, entity_list);
3329
3330 _ = iter.next();
3331
3332 iter.sort::<Entity>();
3333
3334 let mut query_2 = world.query::<&mut A>();
3335 let mut iter_2 = query_2.iter_many_mut(&mut world, entity_list);
3336
3337 _ = iter_2.fetch_next();
3338
3339 iter_2.sort::<Entity>();
3340 }
3341
3342 // This test should be run with miri to check for UB caused by aliasing.
3343 // The lens items created during the sort must not be live at the same time as the mutable references returned from the iterator.
3344 #[test]
3345 fn query_iter_many_sorts_duplicate_entities_no_ub() {
3346 #[derive(Component, Ord, PartialOrd, Eq, PartialEq)]
3347 struct C(usize);
3348
3349 let mut world = World::new();
3350 let id = world.spawn(C(10)).id();
3351 let mut query_state = world.query::<&mut C>();
3352
3353 {
3354 let mut query = query_state.iter_many_mut(&mut world, [id, id]).sort::<&C>();
3355 while query.fetch_next().is_some() {}
3356 }
3357 {
3358 let mut query = query_state
3359 .iter_many_mut(&mut world, [id, id])
3360 .sort_unstable::<&C>();
3361 while query.fetch_next().is_some() {}
3362 }
3363 {
3364 let mut query = query_state
3365 .iter_many_mut(&mut world, [id, id])
3366 .sort_by::<&C>(|l, r| Ord::cmp(l, r));
3367 while query.fetch_next().is_some() {}
3368 }
3369 {
3370 let mut query = query_state
3371 .iter_many_mut(&mut world, [id, id])
3372 .sort_unstable_by::<&C>(|l, r| Ord::cmp(l, r));
3373 while query.fetch_next().is_some() {}
3374 }
3375 {
3376 let mut query = query_state
3377 .iter_many_mut(&mut world, [id, id])
3378 .sort_by_key::<&C, _>(|d| d.0);
3379 while query.fetch_next().is_some() {}
3380 }
3381 {
3382 let mut query = query_state
3383 .iter_many_mut(&mut world, [id, id])
3384 .sort_unstable_by_key::<&C, _>(|d| d.0);
3385 while query.fetch_next().is_some() {}
3386 }
3387 {
3388 let mut query = query_state
3389 .iter_many_mut(&mut world, [id, id])
3390 .sort_by_cached_key::<&C, _>(|d| d.0);
3391 while query.fetch_next().is_some() {}
3392 }
3393 }
3394}