Skip to content

Commit 7b98db6

Browse files
Victoronzmockersf
authored andcommitted
add missing sort_unstable_by_key to QueryIter (#14040)
# Objective `QueryIter::sort_unstable_by_key` is missing. ## Solution Add `QueryIter::sort_unstable_by_key`. ## Testing Added the new method to existing test. ## Changelog Added `QueryIter::sort_unstable_by_key`.
1 parent 7c60387 commit 7b98db6

File tree

1 file changed

+80
-0
lines changed

1 file changed

+80
-0
lines changed

crates/bevy_ecs/src/query/iter.rs

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -730,6 +730,77 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
730730
}
731731
}
732732

733+
/// Sorts all query items into a new iterator with a key extraction function over the query lens.
734+
///
735+
/// This sort is unstable (i.e., may reorder equal elements).
736+
///
737+
/// This uses [`slice::sort_unstable_by_key`] internally.
738+
///
739+
/// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
740+
/// This includes the allowed parameter type changes listed under [allowed transmutes].
741+
/// However, the lens uses the filter of the original query when present.
742+
///
743+
/// The sort is not cached across system runs.
744+
///
745+
/// [allowed transmutes]: crate::system::Query#allowed-transmutes
746+
///
747+
/// # Panics
748+
///
749+
/// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
750+
pub fn sort_unstable_by_key<L: ReadOnlyQueryData + 'w, K>(
751+
self,
752+
mut f: impl FnMut(&L::Item<'w>) -> K,
753+
) -> QuerySortedIter<
754+
'w,
755+
's,
756+
D,
757+
F,
758+
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
759+
>
760+
where
761+
K: Ord,
762+
{
763+
// On the first successful iteration of `QueryIterationCursor`, `archetype_entities` or `table_entities`
764+
// will be set to a non-zero value. The correctness of this method relies on this.
765+
// I.e. this sort method will execute if and only if `next` on `QueryIterationCursor` of a
766+
// non-empty `QueryIter` has not yet been called. When empty, this sort method will not panic.
767+
if !self.cursor.archetype_entities.is_empty() || !self.cursor.table_entities.is_empty() {
768+
panic!("it is not valid to call sort() after next()")
769+
}
770+
771+
let world = self.world;
772+
773+
let query_lens_state = self
774+
.query_state
775+
.transmute_filtered::<(L, Entity), F>(world.components());
776+
777+
// SAFETY:
778+
// `self.world` has permission to access the required components.
779+
// The original query iter has not been iterated on, so no items are aliased from it.
780+
let query_lens = unsafe {
781+
query_lens_state.iter_unchecked_manual(
782+
world,
783+
world.last_change_tick(),
784+
world.change_tick(),
785+
)
786+
};
787+
let mut keyed_query: Vec<_> = query_lens.collect();
788+
keyed_query.sort_unstable_by_key(|(lens, _)| f(lens));
789+
let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity);
790+
// SAFETY:
791+
// `self.world` has permission to access the required components.
792+
// Each lens query item is dropped before the respective actual query item is accessed.
793+
unsafe {
794+
QuerySortedIter::new(
795+
world,
796+
self.query_state,
797+
entity_iter,
798+
world.last_change_tick(),
799+
world.change_tick(),
800+
)
801+
}
802+
}
803+
733804
/// Sort all query items into a new iterator with a key extraction function over the query lens.
734805
///
735806
/// This sort is stable (i.e., does not reorder equal elements).
@@ -1681,6 +1752,11 @@ mod tests {
16811752
.sort_by_key::<Entity, _>(|&e| e)
16821753
.collect::<Vec<_>>();
16831754

1755+
let sort_unstable_by_key = query
1756+
.iter(&world)
1757+
.sort_unstable_by_key::<Entity, _>(|&e| e)
1758+
.collect::<Vec<_>>();
1759+
16841760
let sort_by_cached_key = query
16851761
.iter(&world)
16861762
.sort_by_cached_key::<Entity, _>(|&e| e)
@@ -1701,6 +1777,9 @@ mod tests {
17011777
let mut sort_by_key_v2 = query.iter(&world).collect::<Vec<_>>();
17021778
sort_by_key_v2.sort_by_key(|&e| e);
17031779

1780+
let mut sort_unstable_by_key_v2 = query.iter(&world).collect::<Vec<_>>();
1781+
sort_unstable_by_key_v2.sort_unstable_by_key(|&e| e);
1782+
17041783
let mut sort_by_cached_key_v2 = query.iter(&world).collect::<Vec<_>>();
17051784
sort_by_cached_key_v2.sort_by_cached_key(|&e| e);
17061785

@@ -1709,6 +1788,7 @@ mod tests {
17091788
assert_eq!(sort_by, sort_by_v2);
17101789
assert_eq!(sort_unstable_by, sort_unstable_by_v2);
17111790
assert_eq!(sort_by_key, sort_by_key_v2);
1791+
assert_eq!(sort_unstable_by_key, sort_unstable_by_key_v2);
17121792
assert_eq!(sort_by_cached_key, sort_by_cached_key_v2);
17131793
}
17141794

0 commit comments

Comments
 (0)