From 399fd23797bf2c45ec6a889633d7a90878ea0073 Mon Sep 17 00:00:00 2001 From: Vic <59878206+Victoronz@users.noreply.github.com> Date: Tue, 21 May 2024 20:35:19 +0200 Subject: [PATCH] implement the full set of sort methods on QueryIter (#13417) # Objective Currently, a query iterator can be collected into a `Vec` and sorted, but this can be quite unwieldy, especially when many `Component`s are involved. The `itertools` crate helps somewhat, but the need to write a closure over all of `QueryData` can sometimes hurt ergonomics, anywhere from slightly to strongly. A key extraction function only partially helps, as `sort_by_key` does not allow returning non-`Copy` data. `sort_by` does not suffer from the `Copy` restriction, but now the user has to write out a `cmp` function over two `QueryData::Item`s when it could have just been handled by the `Ord` impl for the key. `sort` requires the entire `Iterator` Item to be `Ord`, which is rarely usable without manual helper functionality. If the user wants to hide away unused components with a `..` range, they need to track item tuple order across their function. Mutable `QueryData` can also introduce further complexity. Additionally, sometimes users solely include `Component`s /`Entity` to guarantee iteration order. For a user to write a function to abstract away repeated sorts over various `QueryData` types they use would require reaching for the `all_tuples!` macro, and continue tracking tuple order afterwards. Fixes https://github.com/bevyengine/bevy/issues/1470. ## Solution Custom sort methods on `QueryIter`, which take a query lens as a generic argument, like `transmute_lens` in `Query`. This allows users to choose what part of their queries they pass to their sort function calls, serving as a kind of "key extraction function" before the sort call. F.e. allowing users to implement `Ord` for a Component, then call `query.iter().sort::()` This works independent of mutability in `QueryData`, `QueryData` tuple order, or the underlying `iter/iter_mut` call. Non-`Copy` components could also be used this way, an internal `Arc` being an example. If `Ord` impls on components do not suffice, other sort methods can be used. Notably useful when combined with `EntityRef` or `EntityMut`. Another boon from using underlying `transmute` functionality, is that with the [allowed transmutes](http://dev-docs.bevyengine.org/bevy/ecs/prelude/struct.Query.html#allowed-transmutes), it is possible to sort a `Query` with `Entity` even if it wasn't included in the original `Query`. The additional generic parameter on the methods other than `sort` and `sort_unstable` currently cannot be removed due to Rust limitations, however their types can be inferred. The new methods do not conflict with the `itertools` sort methods, as those use the "sorted" prefix. This is implemented barely touching existing code. That change to existing code being that `QueryIter` now holds on to the reference to `UnsafeWorldCell` that is used to initialize it. A lens query is constructed with `Entity` attached at the end, sorted, and turned into an iterator. The iterator maps away the lens query, leaving only an iterator of `Entity`, which is used by `QuerySortedIter` to retrieve the actual items. `QuerySortedIter` resembles a combination of `QueryManyIter` and `QueryIter`, but it uses an entity list that is guaranteed to contain unique entities, and implements `ExactSizeIterator`, `DoubleEndedIterator`, `FusedIterator` regardless of mutability or filter kind (archetypal/non-archetypal). The sort methods are not allowed to be called after `next`, and will panic otherwise. This is checked using `QueryIterationCursor` state, which is unique on initialization. Empty queries are an exception to this, as they do not return any item in the first place. That is because tracking how many iterations have already passed would require regressing either normal query iteration a slight bit, or sorted iteration by a lot. Besides, that would not be the intended use of these methods. ## Testing To ensure that `next` being called before `sort` results in a panic, I added some tests. I also test that empty `QueryIter`s do not exhibit this restriction. The query sorts test checks for equivalence to the underlying sorts. This change requires that `Query<(Entity, Entity)>` remains legal, if that is not already guaranteed, which is also ensured by the aforementioned test. ## Next Steps Implement the set of sort methods for `QueryManyIter` as well. - This will mostly work the same, other than needing to return a new `QuerySortedManyIter` to account for iteration over lists of entities that are not guaranteed to be unique. This new query iterator will need a bit of internal restructuring to allow for double-ended mutable iteration, while not regressing read-only iteration. The implementations for each pair of - `sort`, `sort_unstable`, - `sort_by`, sort_unstable_by, - `sort_by_key,` `sort_by_cached_key` are the same aside from the panic message and the sort call, so they could be merged with an inner function. That would require the use of higher-ranked trait bounds on `WorldQuery::Item<'1>`, and is unclear to me whether it is currently doable. Iteration in QuerySortedIter might have space for improvement. When sorting by `Entity`, an `(Entity, Entity)` lens `QueryData` is constructed, is that worth remedying? When table sorts are implemented, a fast path could be introduced to these sort methods. ## Future Possibilities Implementing `Ord` for EntityLocation might be useful. Some papercuts in ergonomics can be improved by future Rust features: - The additional generic parameter aside from the query lens can be removed once this feature is stable: `Fn -> impl Trait` (`impl Trait` in `Fn` trait return position) - With type parameter defaults, the query lens generic can be defaulted to `QueryData::Item`, allowing the sort methods to look and behave like `slice::sort` when no query lens is specified. - With TAIT, the iterator generic on `QuerySortedIter` and thus the huge visible `impl Iterator` type in the sort function signatures can be removed. - With specialization, the bound on `L` could be relaxed to `QueryData` when the underlying iterator is mutable. ## Changelog Added `sort`, `sort_unstable`, `sort_by`, `sort_unstable_by`, `sort_by_key`, `sort_by_cached_key` to `QueryIter`. --- crates/bevy_ecs/src/query/iter.rs | 985 +++++++++++++++++++++++++++++- 1 file changed, 983 insertions(+), 2 deletions(-) diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index ae6eab9162..9909106093 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -6,7 +6,7 @@ use crate::{ storage::{Table, TableRow, Tables}, world::unsafe_world_cell::UnsafeWorldCell, }; -use std::{borrow::Borrow, iter::FusedIterator, mem::MaybeUninit, ops::Range}; +use std::{borrow::Borrow, cmp::Ordering, iter::FusedIterator, mem::MaybeUninit, ops::Range}; use super::{QueryData, QueryFilter, ReadOnlyQueryData}; @@ -15,6 +15,7 @@ use super::{QueryData, QueryFilter, ReadOnlyQueryData}; /// This struct is created by the [`Query::iter`](crate::system::Query::iter) and /// [`Query::iter_mut`](crate::system::Query::iter_mut) methods. pub struct QueryIter<'w, 's, D: QueryData, F: QueryFilter> { + world: UnsafeWorldCell<'w>, tables: &'w Tables, archetypes: &'w Archetypes, query_state: &'s QueryState, @@ -32,6 +33,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { this_run: Tick, ) -> Self { QueryIter { + world, query_state, // SAFETY: We only access table data that has been registered in `query_state`. tables: unsafe { &world.storages().tables }, @@ -157,6 +159,641 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { } accum } + + /// Sorts all query items into a new iterator, using the query lens as a key. + /// + /// This sort is stable (i.e., does not reorder equal elements). + /// + /// This uses [`slice::sort`] 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. + /// + /// # Examples + /// ```rust + /// # use bevy_ecs::prelude::*; + /// # use std::{ops::{Deref, DerefMut}, iter::Sum}; + /// # + /// # #[derive(Component)] + /// # struct PartMarker; + /// # + /// # #[derive(Component, PartialEq, Eq, PartialOrd, Ord)] + /// # struct PartIndex(usize); + /// # + /// # #[derive(Component, Clone, Copy)] + /// # struct PartValue(f32); + /// # + /// # impl Deref for PartValue { + /// # type Target = f32; + /// # + /// # fn deref(&self) -> &Self::Target { + /// # &self.0 + /// # } + /// # } + /// # + /// # #[derive(Component)] + /// # struct ParentValue(f32); + /// # + /// # impl Deref for ParentValue { + /// # type Target = f32; + /// # + /// # fn deref(&self) -> &Self::Target { + /// # &self.0 + /// # } + /// # } + /// # + /// # impl DerefMut for ParentValue { + /// # fn deref_mut(&mut self) -> &mut Self::Target { + /// # &mut self.0 + /// # } + /// # } + /// # + /// # #[derive(Component, Debug, PartialEq, Eq, PartialOrd, Ord)] + /// # struct Length(usize); + /// # + /// # #[derive(Component, Debug, PartialEq, Eq, PartialOrd, Ord)] + /// # struct Width(usize); + /// # + /// # #[derive(Component, Debug, PartialEq, Eq, PartialOrd, Ord)] + /// # struct Height(usize); + /// # + /// # #[derive(Component, PartialEq, Eq, PartialOrd, Ord)] + /// # struct ParentEntity(Entity); + /// # + /// # #[derive(Component, Clone, Copy)] + /// # struct ChildPartCount(usize); + /// # + /// # impl Deref for ChildPartCount { + /// # type Target = usize; + /// # + /// # fn deref(&self) -> &Self::Target { + /// # &self.0 + /// # } + /// # } + /// # let mut world = World::new(); + /// // We can ensure that a query always returns in the same order. + /// fn system_1(query: Query<(Entity, &PartIndex)>) { + /// let parts: Vec<(Entity, &PartIndex)> = query.iter().sort::<&PartIndex>().collect(); + /// } + /// + /// // We can freely rearrange query components in the key. + /// fn system_2(query: Query<(&Length, &Width, &Height), With>) { + /// for (length, width, height) in query.iter().sort::<(&Height, &Length, &Width)>() { + /// println!("height: {height:?}, width: {width:?}, length: {length:?}") + /// } + /// } + /// + /// // We can sort by Entity without including it in the original Query. + /// // Here, we match iteration orders between query iterators. + /// fn system_3( + /// part_query: Query<(&PartValue, &ParentEntity)>, + /// mut parent_query: Query<(&ChildPartCount, &mut ParentValue)>, + /// ) { + /// let part_values = &mut part_query + /// .into_iter() + /// .sort::<&ParentEntity>() + /// .map(|(&value, parent_entity)| *value); + /// + /// for (&child_count, mut parent_value) in parent_query.iter_mut().sort::() { + /// **parent_value = part_values.take(*child_count).sum(); + /// } + /// } + /// # + /// # let mut schedule = Schedule::default(); + /// # schedule.add_systems((system_1, system_2, system_3)); + /// # schedule.run(&mut world); + /// ``` + pub fn sort( + self, + ) -> QuerySortedIter< + 'w, + 's, + D, + F, + impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, + > + where + L::Item<'w>: 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 + .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(), + ) + } + } + + /// Sorts all query items into a new iterator, using the query lens as a key. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// This uses [`slice::sort_unstable`] 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. + /// + /// # Example + /// ``` + /// # use bevy_ecs::prelude::*; + /// # + /// # let mut world = World::new(); + /// # + /// # #[derive(Component)] + /// # struct PartMarker; + /// # + /// #[derive(Component, PartialEq, Eq, PartialOrd, Ord)] + /// enum Flying { + /// Enabled, + /// Disabled + /// }; + /// + /// // We perform an unstable sort by a Component with few values. + /// fn system_1(query: Query<&Flying, With>) { + /// let part_values: Vec<&Flying> = query.iter().sort_unstable::<&Flying>().collect(); + /// } + /// # + /// # let mut schedule = Schedule::default(); + /// # schedule.add_systems((system_1)); + /// # schedule.run(&mut world); + /// ``` + pub fn sort_unstable( + self, + ) -> QuerySortedIter< + 'w, + 's, + D, + F, + impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, + > + where + L::Item<'w>: 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 + .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(), + ) + } + } + + /// Sorts all query items into a new iterator with a comparator function over the query lens. + /// + /// This sort is stable (i.e., does not reorder equal elements). + /// + /// This uses [`slice::sort_by`] 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. + /// + /// # Example + /// ``` + /// # use bevy_ecs::prelude::*; + /// # use std::ops::Deref; + /// # + /// # impl Deref for PartValue { + /// # type Target = f32; + /// # + /// # fn deref(&self) -> &Self::Target { + /// # &self.0 + /// # } + /// # } + /// # + /// # let mut world = World::new(); + /// # + /// #[derive(Component)] + /// struct PartValue(f32); + /// + /// // We can use a cmp function on components do not implement Ord. + /// fn system_1(query: Query<&PartValue>) { + /// // Sort part values according to `f32::total_comp`. + /// let part_values: Vec<&PartValue> = query + /// .iter() + /// .sort_by::<&PartValue>(|value_1, value_2| value_1.total_cmp(*value_2)) + /// .collect(); + /// } + /// # + /// # let mut schedule = Schedule::default(); + /// # schedule.add_systems((system_1)); + /// # schedule.run(&mut world); + /// ``` + pub fn sort_by( + self, + mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering, + ) -> QuerySortedIter< + 'w, + 's, + D, + F, + impl ExactSizeIterator + 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.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_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(), + ) + } + } + + /// Sorts all query items into a new iterator with a comparator function over the query lens. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// This uses [`slice::sort_unstable_by`] 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( + self, + mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering, + ) -> QuerySortedIter< + 'w, + 's, + D, + F, + impl ExactSizeIterator + 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.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_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(), + ) + } + } + + /// Sorts 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). + /// + /// This uses [`slice::sort_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. + /// + /// # Example + /// ``` + /// # use bevy_ecs::prelude::*; + /// # use std::ops::Deref; + /// # + /// # #[derive(Component)] + /// # struct PartMarker; + /// # + /// # impl Deref for PartValue { + /// # type Target = f32; + /// # + /// # fn deref(&self) -> &Self::Target { + /// # &self.0 + /// # } + /// # } + /// # + /// # let mut world = World::new(); + /// # + /// #[derive(Component)] + /// struct AvailableMarker; + /// + /// #[derive(Component, PartialEq, Eq, PartialOrd, Ord)] + /// enum Rarity { + /// Common, + /// Rare, + /// Epic, + /// Legendary + /// }; + /// + /// #[derive(Component)] + /// struct PartValue(f32); + /// + /// // We can sort with the internals of components that do not implement Ord. + /// fn system_1(query: Query<(Entity, &PartValue)>) { + /// // Sort by the sines of the part values. + /// let parts: Vec<(Entity, &PartValue)> = query + /// .iter() + /// .sort_by_key::<&PartValue, _>(|value| value.sin() as usize) + /// .collect(); + /// } + /// + /// // We can define our own custom comparison functions over an EntityRef. + /// fn system_2(query: Query>) { + /// // Sort by whether parts are available and their rarity. + /// // We want the available legendaries to come first, so we reverse the iterator. + /// let parts: Vec = query.iter() + /// .sort_by_key::(|entity_ref| { + /// ( + /// entity_ref.contains::(), + /// entity_ref.get::() + /// ) + /// }) + /// .rev() + /// .collect(); + /// } + /// # let mut schedule = Schedule::default(); + /// # schedule.add_systems((system_1, system_2)); + /// # schedule.run(&mut world); + /// ``` + pub fn sort_by_key( + self, + mut f: impl FnMut(&L::Item<'w>) -> K, + ) -> QuerySortedIter< + 'w, + 's, + D, + F, + impl ExactSizeIterator + 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_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). + /// + /// This uses [`slice::sort_by_cached_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_by_cached_key( + self, + mut f: impl FnMut(&L::Item<'w>) -> K, + ) -> QuerySortedIter< + 'w, + 's, + D, + F, + impl ExactSizeIterator + 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_by_cached_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(), + ) + } + } } impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for QueryIter<'w, 's, D, F> { @@ -220,6 +857,134 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for QueryIter<'w, 's, D, F> // This is correct as [`QueryIter`] always returns `None` once exhausted. impl<'w, 's, D: QueryData, F: QueryFilter> FusedIterator for QueryIter<'w, 's, D, F> {} +/// An [`Iterator`] over sorted query results of a [`Query`](crate::system::Query). +/// +/// This struct is created by the [`QueryIter::sort`], [`QueryIter::sort_unstable`], +/// [`QueryIter::sort_by`], [`QueryIter::sort_unstable_by`], [`QueryIter::sort_by_key`], +/// and [`QueryIter::sort_by_cached_key`] methods. +pub struct QuerySortedIter<'w, 's, D: QueryData, F: QueryFilter, I> +where + I: Iterator, +{ + entity_iter: I, + entities: &'w Entities, + tables: &'w Tables, + archetypes: &'w Archetypes, + fetch: D::Fetch<'w>, + query_state: &'s QueryState, +} + +impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> QuerySortedIter<'w, 's, D, F, I> +where + I: Iterator, +{ + /// # Safety + /// - `world` must have permission to access any of the components registered in `query_state`. + /// - `world` must be the same one used to initialize `query_state`. + /// - `entity_list` must only contain unique entities or be empty. + pub(crate) unsafe fn new>( + world: UnsafeWorldCell<'w>, + query_state: &'s QueryState, + entity_list: EntityList, + last_run: Tick, + this_run: Tick, + ) -> QuerySortedIter<'w, 's, D, F, I> { + let fetch = D::init_fetch(world, &query_state.fetch_state, last_run, this_run); + QuerySortedIter { + query_state, + entities: world.entities(), + archetypes: world.archetypes(), + // SAFETY: We only access table data that has been registered in `query_state`. + // This means `world` has permission to access the data we use. + tables: &world.storages().tables, + fetch, + entity_iter: entity_list.into_iter(), + } + } + + /// # Safety + /// `entity` must stem from `self.entity_iter`, and not have been passed before. + #[inline(always)] + unsafe fn fetch_next(&mut self, entity: Entity) -> D::Item<'w> { + let (location, archetype, table); + // SAFETY: + // `tables` and `archetypes` belong to the same world that the [`QueryIter`] + // was initialized for. + unsafe { + location = self.entities.get(entity).debug_checked_unwrap(); + archetype = self + .archetypes + .get(location.archetype_id) + .debug_checked_unwrap(); + table = self.tables.get(location.table_id).debug_checked_unwrap(); + } + + // SAFETY: `archetype` is from the world that `fetch` was created for, + // `fetch_state` is the state that `fetch` was initialized with + unsafe { + D::set_archetype( + &mut self.fetch, + &self.query_state.fetch_state, + archetype, + table, + ); + } + + // The entity list has already been filtered by the query lens, so we forego filtering here. + // SAFETY: + // - set_archetype was called prior, `location.archetype_row` is an archetype index in range of the current archetype + // - fetch is only called once for each entity. + unsafe { D::fetch(&mut self.fetch, entity, location.table_row) } + } +} + +impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Iterator + for QuerySortedIter<'w, 's, D, F, I> +where + I: Iterator, +{ + type Item = D::Item<'w>; + + #[inline(always)] + fn next(&mut self) -> Option { + let entity = self.entity_iter.next()?; + // SAFETY: `entity` is passed from `entity_iter` the first time. + unsafe { self.fetch_next(entity).into() } + } + + fn size_hint(&self) -> (usize, Option) { + self.entity_iter.size_hint() + } +} + +impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> DoubleEndedIterator + for QuerySortedIter<'w, 's, D, F, I> +where + I: DoubleEndedIterator, +{ + #[inline(always)] + fn next_back(&mut self) -> Option { + let entity = self.entity_iter.next_back()?; + // SAFETY: `entity` is passed from `entity_iter` the first time. + unsafe { self.fetch_next(entity).into() } + } +} + +impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> ExactSizeIterator + for QuerySortedIter<'w, 's, D, F, I> +where + I: ExactSizeIterator, +{ +} + +// This is correct as [`QuerySortedIter`] returns `None` once exhausted if `entity_iter` does. +impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> FusedIterator + for QuerySortedIter<'w, 's, D, F, I> +where + I: FusedIterator, +{ +} + /// An [`Iterator`] over the query items generated from an iterator of [`Entity`]s. /// /// Items are returned in the order of the provided iterator. @@ -704,7 +1469,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { } // NOTE: If you are changing query iteration code, remember to update the following places, where relevant: - // QueryIter, QueryIterationCursor, QueryManyIter, QueryCombinationIter, QueryState::par_fold_init_unchecked_manual + // QueryIter, QueryIterationCursor, QuerySortedIter, QueryManyIter, QueryCombinationIter, QueryState::par_fold_init_unchecked_manual /// # Safety /// `tables` and `archetypes` must belong to the same world that the [`QueryIterationCursor`] /// was initialized for. @@ -812,3 +1577,219 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { } } } + +// A wrapper struct that gives its data a neutral ordering. +#[derive(Copy, Clone)] +struct NeutralOrd(T); + +impl PartialEq for NeutralOrd { + fn eq(&self, _other: &Self) -> bool { + true + } +} + +impl Eq for NeutralOrd {} + +#[allow(clippy::non_canonical_partial_ord_impl)] +impl PartialOrd for NeutralOrd { + fn partial_cmp(&self, _other: &Self) -> Option { + Some(Ordering::Equal) + } +} + +impl Ord for NeutralOrd { + fn cmp(&self, _other: &Self) -> Ordering { + Ordering::Equal + } +} + +#[cfg(test)] +mod tests { + #[allow(unused_imports)] + use crate::{self as bevy_ecs, component::Component, entity::Entity, prelude::World}; + + #[derive(Component, Debug, PartialEq, PartialOrd, Clone, Copy)] + struct A(f32); + #[derive(Component, Debug, Eq, PartialEq, Clone, Copy)] + #[component(storage = "SparseSet")] + struct Sparse(usize); + + #[allow(clippy::unnecessary_sort_by)] + #[test] + fn query_sorts() { + let mut world = World::new(); + + let mut query = world.query::(); + + let sort = query.iter(&world).sort::().collect::>(); + + let sort_unstable = query + .iter(&world) + .sort_unstable::() + .collect::>(); + + let sort_by = query + .iter(&world) + .sort_by::(|e1, e2| e1.cmp(e2)) + .collect::>(); + + let sort_unstable_by = query + .iter(&world) + .sort_unstable_by::(|e1, e2| e1.cmp(e2)) + .collect::>(); + + let sort_by_key = query + .iter(&world) + .sort_by_key::(|&e| e) + .collect::>(); + + let sort_by_cached_key = query + .iter(&world) + .sort_by_cached_key::(|&e| e) + .collect::>(); + + let mut sort_v2 = query.iter(&world).collect::>(); + sort_v2.sort(); + + let mut sort_unstable_v2 = query.iter(&world).collect::>(); + sort_unstable_v2.sort_unstable(); + + let mut sort_by_v2 = query.iter(&world).collect::>(); + sort_by_v2.sort_by(|e1, e2| e1.cmp(e2)); + + let mut sort_unstable_by_v2 = query.iter(&world).collect::>(); + sort_unstable_by_v2.sort_unstable_by(|e1, e2| e1.cmp(e2)); + + let mut sort_by_key_v2 = query.iter(&world).collect::>(); + sort_by_key_v2.sort_by_key(|&e| e); + + let mut sort_by_cached_key_v2 = query.iter(&world).collect::>(); + sort_by_cached_key_v2.sort_by_cached_key(|&e| e); + + assert_eq!(sort, sort_v2); + assert_eq!(sort_unstable, sort_unstable_v2); + 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_by_cached_key, sort_by_cached_key_v2); + } + + #[test] + #[should_panic] + fn query_sort_after_next() { + let mut world = World::new(); + world.spawn((A(0.),)); + world.spawn((A(1.1),)); + world.spawn((A(2.22),)); + + { + let mut query = world.query::<&A>(); + let mut iter = query.iter(&world); + println!( + "archetype_entities: {} table_entities: {} current_len: {} current_row: {}", + iter.cursor.archetype_entities.len(), + iter.cursor.table_entities.len(), + iter.cursor.current_len, + iter.cursor.current_row + ); + _ = iter.next(); + println!( + "archetype_entities: {} table_entities: {} current_len: {} current_row: {}", + iter.cursor.archetype_entities.len(), + iter.cursor.table_entities.len(), + iter.cursor.current_len, + iter.cursor.current_row + ); + println!("{}", iter.sort::().len()); + } + } + + #[test] + #[should_panic] + fn query_sort_after_next_dense() { + let mut world = World::new(); + world.spawn((Sparse(11),)); + world.spawn((Sparse(22),)); + world.spawn((Sparse(33),)); + + { + let mut query = world.query::<&Sparse>(); + let mut iter = query.iter(&world); + println!( + "before_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}", + iter.cursor.archetype_entities.len(), + iter.cursor.table_entities.len(), + iter.cursor.current_len, + iter.cursor.current_row + ); + _ = iter.next(); + println!( + "after_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}", + iter.cursor.archetype_entities.len(), + iter.cursor.table_entities.len(), + iter.cursor.current_len, + iter.cursor.current_row + ); + println!("{}", iter.sort::().len()); + } + } + + #[test] + fn empty_query_sort_after_next_does_not_panic() { + let mut world = World::new(); + { + let mut query = world.query::<(&A, &Sparse)>(); + let mut iter = query.iter(&world); + println!( + "before_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}", + iter.cursor.archetype_entities.len(), + iter.cursor.table_entities.len(), + iter.cursor.current_len, + iter.cursor.current_row + ); + _ = iter.next(); + println!( + "after_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}", + iter.cursor.archetype_entities.len(), + iter.cursor.table_entities.len(), + iter.cursor.current_len, + iter.cursor.current_row + ); + println!("{}", iter.sort::().len()); + } + } + + #[test] + fn query_iter_cursor_state_non_empty_after_next() { + let mut world = World::new(); + world.spawn((A(0.), Sparse(11))); + world.spawn((A(1.1), Sparse(22))); + world.spawn((A(2.22), Sparse(33))); + { + let mut query = world.query::<(&A, &Sparse)>(); + let mut iter = query.iter(&world); + println!( + "before_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}", + iter.cursor.archetype_entities.len(), + iter.cursor.table_entities.len(), + iter.cursor.current_len, + iter.cursor.current_row + ); + assert!(iter.cursor.table_entities.len() | iter.cursor.archetype_entities.len() == 0); + _ = iter.next(); + println!( + "after_next_call: archetype_entities: {} table_entities: {} current_len: {} current_row: {}", + iter.cursor.archetype_entities.len(), + iter.cursor.table_entities.len(), + iter.cursor.current_len, + iter.cursor.current_row + ); + assert!( + ( + iter.cursor.table_entities.len(), + iter.cursor.archetype_entities.len() + ) != (0, 0) + ); + } + } +}