Share implementation of sort methods. (#16203)

# Objective

The various `Query::sort()` methods have a lot of duplicated code
between them, including some unsafe code. Reduce the duplication to make
the code easier to read and maintain.

## Solution

Extract the duplicated code to a private method, and pass in the sorting
strategy as a closure.

## Testing

I used `cargo-show-asm` to verify that the closures were inlined, but I
didn't run anything through a profiler. The `sort()` method itself even
had identical assembly before and after this change, although the others
did not.
This commit is contained in:
Chris Russell 2025-01-27 23:57:54 -05:00 committed by GitHub
parent c0ccc87738
commit 514a35c656
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -496,45 +496,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
> {
// 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);
// 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
.map(|(key, entity)| (key, NeutralOrd(entity)))
.collect();
keyed_query.sort();
let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity.0);
// 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(),
)
}
self.sort_impl::<L>(|keyed_query| keyed_query.sort())
}
/// Sorts all query items into a new iterator, using the query lens as a key.
@ -588,45 +550,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
> {
// 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);
// 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
.map(|(key, entity)| (key, NeutralOrd(entity)))
.collect();
keyed_query.sort_unstable();
let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity.0);
// 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(),
)
}
self.sort_impl::<L>(|keyed_query| keyed_query.sort_unstable())
}
/// Sorts all query items into a new iterator with a comparator function over the query lens.
@ -688,43 +612,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
> {
// 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);
// 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_by(|(key_1, _), (key_2, _)| compare(key_1, key_2));
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(),
)
}
self.sort_impl::<L>(move |keyed_query| {
keyed_query.sort_by(|(key_1, _), (key_2, _)| compare(key_1, key_2));
})
}
/// Sorts all query items into a new iterator with a comparator function over the query lens.
@ -754,43 +644,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
> {
// 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);
// 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_1, _), (key_2, _)| compare(key_1, key_2));
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(),
)
}
self.sort_impl::<L>(move |keyed_query| {
keyed_query.sort_unstable_by(|(key_1, _), (key_2, _)| compare(key_1, key_2));
})
}
/// Sorts all query items into a new iterator with a key extraction function over the query lens.
@ -883,43 +739,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
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);
// 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_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(),
)
}
self.sort_impl::<L>(move |keyed_query| keyed_query.sort_by_key(|(lens, _)| f(lens)))
}
/// Sorts all query items into a new iterator with a key extraction function over the query lens.
@ -952,43 +772,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
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);
// 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(),
)
}
self.sort_impl::<L>(move |keyed_query| {
keyed_query.sort_unstable_by_key(|(lens, _)| f(lens));
})
}
/// Sort all query items into a new iterator with a key extraction function over the query lens.
@ -1021,6 +807,33 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
where
K: Ord,
{
self.sort_impl::<L>(move |keyed_query| keyed_query.sort_by_cached_key(|(lens, _)| f(lens)))
}
/// Shared implementation for the various `sort` methods.
/// This uses the lens to collect the items for sorting, but delegates the actual sorting to the provided closure.
///
/// 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.
fn sort_impl<L: ReadOnlyQueryData + 'w>(
self,
f: impl FnOnce(&mut Vec<(L::Item<'w>, NeutralOrd<Entity>)>),
) -> QuerySortedIter<
'w,
's,
D,
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
> {
// 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
@ -1043,9 +856,11 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
world.change_tick(),
)
};
let mut keyed_query: Vec<_> = query_lens.collect();
keyed_query.sort_by_cached_key(|(lens, _)| f(lens));
let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity);
let mut keyed_query: Vec<_> = query_lens
.map(|(key, entity)| (key, NeutralOrd(entity)))
.collect();
f(&mut keyed_query);
let entity_iter = keyed_query.into_iter().map(|(.., entity)| entity.0);
// SAFETY:
// `self.world` has permission to access the required components.
// Each lens query item is dropped before the respective actual query item is accessed.
@ -1526,45 +1341,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
where
L::Item<'w>: Ord,
{
let world = self.world;
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);
// 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_many_unchecked_manual(
self.entity_iter,
world,
world.last_change_tick(),
world.change_tick(),
)
};
let mut keyed_query: Vec<_> = query_lens
.map(|(key, entity)| (key, NeutralOrd(entity)))
.collect();
keyed_query.sort();
// Re-collect into a `Vec` to eagerly drop the lens items.
// They must be dropped before `fetch_next` is called since they may alias
// with other data items if there are duplicate entities in `entity_iter`.
let entity_iter = keyed_query
.into_iter()
.map(|(.., entity)| entity.0)
.collect::<Vec<_>>()
.into_iter();
// 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 {
QuerySortedManyIter::new(
world,
self.query_state,
entity_iter,
world.last_change_tick(),
world.change_tick(),
)
}
self.sort_impl::<L>(|keyed_query| keyed_query.sort())
}
/// Sorts all query items into a new iterator, using the query lens as a key.
@ -1622,45 +1399,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
where
L::Item<'w>: Ord,
{
let world = self.world;
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);
// 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_many_unchecked_manual(
self.entity_iter,
world,
world.last_change_tick(),
world.change_tick(),
)
};
let mut keyed_query: Vec<_> = query_lens
.map(|(key, entity)| (key, NeutralOrd(entity)))
.collect();
keyed_query.sort_unstable();
// Re-collect into a `Vec` to eagerly drop the lens items.
// They must be dropped before `fetch_next` is called since they may alias
// with other data items if there are duplicate entities in `entity_iter`.
let entity_iter = keyed_query
.into_iter()
.map(|(.., entity)| entity.0)
.collect::<Vec<_>>()
.into_iter();
// 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 {
QuerySortedManyIter::new(
world,
self.query_state,
entity_iter,
world.last_change_tick(),
world.change_tick(),
)
}
self.sort_impl::<L>(|keyed_query| keyed_query.sort_unstable())
}
/// Sorts all query items into a new iterator with a comparator function over the query lens.
@ -1723,43 +1462,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
> {
let world = self.world;
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);
// 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_many_unchecked_manual(
self.entity_iter,
world,
world.last_change_tick(),
world.change_tick(),
)
};
let mut keyed_query: Vec<_> = query_lens.collect();
keyed_query.sort_by(|(key_1, _), (key_2, _)| compare(key_1, key_2));
// Re-collect into a `Vec` to eagerly drop the lens items.
// They must be dropped before `fetch_next` is called since they may alias
// with other data items if there are duplicate entities in `entity_iter`.
let entity_iter = keyed_query
.into_iter()
.map(|(.., entity)| entity)
.collect::<Vec<_>>()
.into_iter();
// 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 {
QuerySortedManyIter::new(
world,
self.query_state,
entity_iter,
world.last_change_tick(),
world.change_tick(),
)
}
self.sort_impl::<L>(move |keyed_query| {
keyed_query.sort_by(|(key_1, _), (key_2, _)| compare(key_1, key_2));
})
}
/// Sorts all query items into a new iterator with a comparator function over the query lens.
@ -1788,43 +1493,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
> {
let world = self.world;
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);
// 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_many_unchecked_manual(
self.entity_iter,
world,
world.last_change_tick(),
world.change_tick(),
)
};
let mut keyed_query: Vec<_> = query_lens.collect();
keyed_query.sort_unstable_by(|(key_1, _), (key_2, _)| compare(key_1, key_2));
// Re-collect into a `Vec` to eagerly drop the lens items.
// They must be dropped before `fetch_next` is called since they may alias
// with other data items if there are duplicate entities in `entity_iter`.
let entity_iter = keyed_query
.into_iter()
.map(|(.., entity)| entity)
.collect::<Vec<_>>()
.into_iter();
// 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 {
QuerySortedManyIter::new(
world,
self.query_state,
entity_iter,
world.last_change_tick(),
world.change_tick(),
)
}
self.sort_impl::<L>(move |keyed_query| {
keyed_query.sort_unstable_by(|(key_1, _), (key_2, _)| compare(key_1, key_2));
})
}
/// Sorts all query items into a new iterator with a key extraction function over the query lens.
@ -1919,43 +1590,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
where
K: Ord,
{
let world = self.world;
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);
// 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_many_unchecked_manual(
self.entity_iter,
world,
world.last_change_tick(),
world.change_tick(),
)
};
let mut keyed_query: Vec<_> = query_lens.collect();
keyed_query.sort_by_key(|(lens, _)| f(lens));
// Re-collect into a `Vec` to eagerly drop the lens items.
// They must be dropped before `fetch_next` is called since they may alias
// with other data items if there are duplicate entities in `entity_iter`.
let entity_iter = keyed_query
.into_iter()
.map(|(.., entity)| entity)
.collect::<Vec<_>>()
.into_iter();
// 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 {
QuerySortedManyIter::new(
world,
self.query_state,
entity_iter,
world.last_change_tick(),
world.change_tick(),
)
}
self.sort_impl::<L>(move |keyed_query| keyed_query.sort_by_key(|(lens, _)| f(lens)))
}
/// Sorts all query items into a new iterator with a key extraction function over the query lens.
@ -1987,43 +1622,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
where
K: Ord,
{
let world = self.world;
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);
// 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_many_unchecked_manual(
self.entity_iter,
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));
// Re-collect into a `Vec` to eagerly drop the lens items.
// They must be dropped before `fetch_next` is called since they may alias
// with other data items if there are duplicate entities in `entity_iter`.
let entity_iter = keyed_query
.into_iter()
.map(|(.., entity)| entity)
.collect::<Vec<_>>()
.into_iter();
// 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 {
QuerySortedManyIter::new(
world,
self.query_state,
entity_iter,
world.last_change_tick(),
world.change_tick(),
)
}
self.sort_impl::<L>(move |keyed_query| {
keyed_query.sort_unstable_by_key(|(lens, _)| f(lens));
})
}
/// Sort all query items into a new iterator with a key extraction function over the query lens.
@ -2055,6 +1656,32 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
where
K: Ord,
{
self.sort_impl::<L>(move |keyed_query| keyed_query.sort_by_cached_key(|(lens, _)| f(lens)))
}
/// Shared implementation for the various `sort` methods.
/// This uses the lens to collect the items for sorting, but delegates the actual sorting to the provided closure.
///
/// 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
///
/// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been
/// called on [`QueryManyIter`] before.
fn sort_impl<L: ReadOnlyQueryData + 'w>(
self,
f: impl FnOnce(&mut Vec<(L::Item<'w>, NeutralOrd<Entity>)>),
) -> QuerySortedManyIter<
'w,
's,
D,
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
> {
let world = self.world;
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);
@ -2070,14 +1697,16 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
world.change_tick(),
)
};
let mut keyed_query: Vec<_> = query_lens.collect();
keyed_query.sort_by_cached_key(|(lens, _)| f(lens));
let mut keyed_query: Vec<_> = query_lens
.map(|(key, entity)| (key, NeutralOrd(entity)))
.collect();
f(&mut keyed_query);
// Re-collect into a `Vec` to eagerly drop the lens items.
// They must be dropped before `fetch_next` is called since they may alias
// with other data items if there are duplicate entities in `entity_iter`.
let entity_iter = keyed_query
.into_iter()
.map(|(.., entity)| entity)
.map(|(.., entity)| entity.0)
.collect::<Vec<_>>()
.into_iter();
// SAFETY: