Skip to content

Commit

Permalink
add missing sort_unstable_by_key to QueryIter (#14040)
Browse files Browse the repository at this point in the history
# 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`.
  • Loading branch information
Victoronz authored and mockersf committed Jun 27, 2024
1 parent 7c60387 commit 7b98db6
Showing 1 changed file with 80 additions and 0 deletions.
80 changes: 80 additions & 0 deletions crates/bevy_ecs/src/query/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -730,6 +730,77 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
}
}

/// Sorts all query items into a new iterator with a key extraction function over the query lens.
///
/// This sort is unstable (i.e., may reorder equal elements).
///
/// This uses [`slice::sort_unstable_by_key`] internally.
///
/// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens).
/// This includes the allowed parameter type changes listed under [allowed transmutes].
/// However, the lens uses the filter of the original query when present.
///
/// The sort is not cached across system runs.
///
/// [allowed transmutes]: crate::system::Query#allowed-transmutes
///
/// # Panics
///
/// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
pub fn sort_unstable_by_key<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
) -> QuerySortedIter<
'w,
's,
D,
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
>
where
K: Ord,
{
// On the first successful iteration of `QueryIterationCursor`, `archetype_entities` or `table_entities`
// will be set to a non-zero value. The correctness of this method relies on this.
// I.e. this sort method will execute if and only if `next` on `QueryIterationCursor` of a
// non-empty `QueryIter` has not yet been called. When empty, this sort method will not panic.
if !self.cursor.archetype_entities.is_empty() || !self.cursor.table_entities.is_empty() {
panic!("it is not valid to call sort() after next()")
}

let world = self.world;

let query_lens_state = self
.query_state
.transmute_filtered::<(L, Entity), F>(world.components());

// SAFETY:
// `self.world` has permission to access the required components.
// The original query iter has not been iterated on, so no items are aliased from it.
let query_lens = unsafe {
query_lens_state.iter_unchecked_manual(
world,
world.last_change_tick(),
world.change_tick(),
)
};
let mut keyed_query: Vec<_> = query_lens.collect();
keyed_query.sort_unstable_by_key(|(lens, _)| f(lens));
let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity);
// SAFETY:
// `self.world` has permission to access the required components.
// Each lens query item is dropped before the respective actual query item is accessed.
unsafe {
QuerySortedIter::new(
world,
self.query_state,
entity_iter,
world.last_change_tick(),
world.change_tick(),
)
}
}

/// Sort all query items into a new iterator with a key extraction function over the query lens.
///
/// This sort is stable (i.e., does not reorder equal elements).
Expand Down Expand Up @@ -1681,6 +1752,11 @@ mod tests {
.sort_by_key::<Entity, _>(|&e| e)
.collect::<Vec<_>>();

let sort_unstable_by_key = query
.iter(&world)
.sort_unstable_by_key::<Entity, _>(|&e| e)
.collect::<Vec<_>>();

let sort_by_cached_key = query
.iter(&world)
.sort_by_cached_key::<Entity, _>(|&e| e)
Expand All @@ -1701,6 +1777,9 @@ mod tests {
let mut sort_by_key_v2 = query.iter(&world).collect::<Vec<_>>();
sort_by_key_v2.sort_by_key(|&e| e);

let mut sort_unstable_by_key_v2 = query.iter(&world).collect::<Vec<_>>();
sort_unstable_by_key_v2.sort_unstable_by_key(|&e| e);

let mut sort_by_cached_key_v2 = query.iter(&world).collect::<Vec<_>>();
sort_by_cached_key_v2.sort_by_cached_key(|&e| e);

Expand All @@ -1709,6 +1788,7 @@ mod tests {
assert_eq!(sort_by, sort_by_v2);
assert_eq!(sort_unstable_by, sort_unstable_by_v2);
assert_eq!(sort_by_key, sort_by_key_v2);
assert_eq!(sort_unstable_by_key, sort_unstable_by_key_v2);
assert_eq!(sort_by_cached_key, sort_by_cached_key_v2);
}

Expand Down

0 comments on commit 7b98db6

Please sign in to comment.