From 41cce65acd0f48617f491a36f095137631e5d42f Mon Sep 17 00:00:00 2001 From: Victoron <59878206+Victoronz@users.noreply.github.com> Date: Mon, 20 May 2024 21:23:33 +0200 Subject: [PATCH] implement sort, sort_unstable, sort_by, sort_unstable_by, sort_by_key, and sort_by_cached_key for 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 ae6eab9162321..9909106093dc9 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) + ); + } + } +}