Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Share implementation of sort methods. #16203

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Share implementation of sort methods. #16203

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
273 changes: 44 additions & 229 deletions crates/bevy_ecs/src/query/iter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -493,45 +493,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.
Expand Down Expand Up @@ -585,45 +547,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.
Expand Down Expand Up @@ -685,43 +609,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.
Expand Down Expand Up @@ -751,43 +641,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));
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Query::sort_unstable_by was calling sort_by instead of sort_unstable_by. That seemed like a bug, so I changed it. If we don't merge this PR, we should probably make one that fixes this line.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch! This is a bug.

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.
Expand Down Expand Up @@ -880,43 +736,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.
Expand Down Expand Up @@ -949,43 +769,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.
Expand Down Expand Up @@ -1018,6 +804,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,
> {
Comment on lines +824 to +833
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you see a way to preserve whether Entity needs to be wrapped in NeutralOrd?

Since this is no longer public, no need to restrict yourself to impl Trait, you can just use plain generics here.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems simpler to use it everywhere. It's a simple newtype wrapper, so it won't affect performance. And the methods that weren't wrapping the Entity don't actually use it during sorting, so it won't affect behavior, either.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fair, it should be okay then!

// 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
Expand All @@ -1040,9 +853,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.
Expand Down