diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index cb1d695a4d58c..4de38258e9edc 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -16,6 +16,7 @@ use core::{ }; use super::{QueryData, QueryFilter, ReadOnlyQueryData}; +use alloc::vec::IntoIter; /// An [`Iterator`] over query results of a [`Query`](crate::system::Query). /// @@ -1251,7 +1252,11 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Debug /// Entities that don't match the query are skipped. /// /// This struct is created by the [`Query::iter_many`](crate::system::Query::iter_many) and [`Query::iter_many_mut`](crate::system::Query::iter_many_mut) methods. -pub struct QueryManyIter<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> { +pub struct QueryManyIter<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> +where + I::Item: Borrow, +{ + world: UnsafeWorldCell<'w>, entity_iter: I, entities: &'w Entities, tables: &'w Tables, @@ -1277,6 +1282,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> let fetch = D::init_fetch(world, &query_state.fetch_state, last_run, this_run); let filter = F::init_fetch(world, &query_state.filter_state, last_run, this_run); QueryManyIter { + world, query_state, entities: world.entities(), archetypes: world.archetypes(), @@ -1334,37 +1340,663 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> F::set_archetype(filter, &query_state.filter_state, archetype, table); } - // SAFETY: set_archetype was called prior. - // `location.archetype_row` is an archetype index row in range of the current archetype, because if it was not, the match above would have `continue`d - if unsafe { F::filter_fetch(filter, entity, location.table_row) } { - // 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. - return Some(unsafe { D::fetch(fetch, entity, location.table_row) }); - } - } - None - } + // SAFETY: set_archetype was called prior. + // `location.archetype_row` is an archetype index row in range of the current archetype, because if it was not, the match above would have `continue`d + if unsafe { F::filter_fetch(filter, entity, location.table_row) } { + // 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. + return Some(unsafe { D::fetch(fetch, entity, location.table_row) }); + } + } + None + } + + /// Get next result from the query + #[inline(always)] + pub fn fetch_next(&mut self) -> Option> { + // SAFETY: + // All arguments stem from self. + // We are limiting the returned reference to self, + // making sure this method cannot be called multiple times without getting rid + // of any previously returned unique references first, thus preventing aliasing. + unsafe { + Self::fetch_next_aliased_unchecked( + &mut self.entity_iter, + self.entities, + self.tables, + self.archetypes, + &mut self.fetch, + &mut self.filter, + self.query_state, + ) + .map(D::shrink) + } + } + + /// 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 + /// + /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been + /// called on [`QueryManyIter`] before. + /// + /// # 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(usize); + /// # + /// # impl Deref for PartValue { + /// # type Target = usize; + /// # + /// # fn deref(&self) -> &Self::Target { + /// # &self.0 + /// # } + /// # } + /// # + /// # impl DerefMut for PartValue { + /// # 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); + /// # + /// # 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 entity_list: Vec = Vec::new(); + /// let parts: Vec<(Entity, &PartIndex)> = query.iter_many(entity_list).sort::<&PartIndex>().collect(); + /// } + /// + /// // We can freely rearrange query components in the key. + /// fn system_2(query: Query<(&Length, &Width, &Height), With>) { + /// # let entity_list: Vec = Vec::new(); + /// for (length, width, height) in query.iter_many(entity_list).sort::<(&Height, &Length, &Width)>() { + /// println!("height: {height:?}, width: {width:?}, length: {length:?}") + /// } + /// } + /// + /// // You can use `fetch_next_back` to obtain mutable references in reverse order. + /// fn system_3( + /// mut query: Query<&mut PartValue>, + /// ) { + /// # let entity_list: Vec = Vec::new(); + /// // We need to collect the internal iterator before iterating mutably + /// let mut parent_query_iter = query.iter_many_mut(entity_list) + /// .sort::() + /// .collect_inner(); + /// + /// let mut scratch_value = 0; + /// while let Some(mut part_value) = parent_query_iter.fetch_next_back() + /// { + /// // some order-dependent operation, here bitwise XOR + /// **part_value ^= scratch_value; + /// scratch_value = **part_value; + /// } + /// } + /// # + /// # let mut schedule = Schedule::default(); + /// # schedule.add_systems((system_1, system_2, system_3)); + /// # schedule.run(&mut world); + /// ``` + pub fn sort( + self, + ) -> QuerySortedManyIter< + 'w, + 's, + D, + F, + impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, + > + 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(); + 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 { + QuerySortedManyIter::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 + /// + /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been + /// called on [`QueryManyIter`] before. + /// + /// # Example + /// ``` + /// # use bevy_ecs::prelude::*; + /// # + /// # let mut world = World::new(); + /// # + /// # #[derive(Component)] + /// # struct PartMarker; + /// # + /// # let entity_list: Vec = Vec::new(); + /// #[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 entity_list: Vec = Vec::new(); + /// let part_values: Vec<&Flying> = query.iter_many(entity_list).sort_unstable::<&Flying>().collect(); + /// } + /// # + /// # let mut schedule = Schedule::default(); + /// # schedule.add_systems((system_1)); + /// # schedule.run(&mut world); + /// ``` + pub fn sort_unstable( + self, + ) -> QuerySortedManyIter< + 'w, + 's, + D, + F, + impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, + > + 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(); + 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 { + QuerySortedManyIter::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 + /// + /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been + /// called on [`QueryManyIter`] before. + /// + /// # 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(); + /// # let entity_list: Vec = Vec::new(); + /// # + /// #[derive(Component)] + /// struct PartValue(f32); + /// + /// // We can use a cmp function on components do not implement Ord. + /// fn system_1(query: Query<&PartValue>) { + /// # let entity_list: Vec = Vec::new(); + /// // Sort part values according to `f32::total_comp`. + /// let part_values: Vec<&PartValue> = query + /// .iter_many(entity_list) + /// .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, + ) -> QuerySortedManyIter< + 'w, + 's, + D, + F, + impl ExactSizeIterator + 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)); + 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 { + QuerySortedManyIter::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 + /// + /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been + /// called on [`QueryManyIter`] before. + pub fn sort_unstable_by( + self, + mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering, + ) -> QuerySortedManyIter< + 'w, + 's, + D, + F, + impl ExactSizeIterator + 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)); + 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 { + QuerySortedManyIter::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 + /// + /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been + /// called on [`QueryManyIter`] before. + /// + /// # 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(); + /// # let entity_list: Vec = Vec::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)>) { + /// # let entity_list: Vec = Vec::new(); + /// // Sort by the sines of the part values. + /// let parts: Vec<(Entity, &PartValue)> = query + /// .iter_many(entity_list) + /// .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>) { + /// # let entity_list: Vec = Vec::new(); + /// // 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_many(entity_list) + /// .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, + ) -> QuerySortedManyIter< + 'w, + 's, + D, + F, + impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, + > + 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)); + 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 { + QuerySortedManyIter::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 unstable (i.e., may reorder equal elements). + /// + /// This uses [`slice::sort_unstable_by_key`] internally. + /// + /// Defining the lens works like [`transmute_lens`](crate::system::Query::transmute_lens). + /// This includes the allowed parameter type changes listed under [allowed transmutes]. + /// However, the lens uses the filter of the original query when present. + /// + /// The sort is not cached across system runs. + /// + /// [allowed transmutes]: crate::system::Query#allowed-transmutes + /// + /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been + /// called on [`QueryManyIter`] before. + pub fn sort_unstable_by_key( + self, + mut f: impl FnMut(&L::Item<'w>) -> K, + ) -> QuerySortedManyIter< + 'w, + 's, + D, + F, + impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, + > + 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)); + 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 { + QuerySortedManyIter::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 + /// + /// Unlike the sort methods on [`QueryIter`], this does NOT panic if `next`/`fetch_next` has been + /// called on [`QueryManyIter`] before. + pub fn sort_by_cached_key( + self, + mut f: impl FnMut(&L::Item<'w>) -> K, + ) -> QuerySortedManyIter< + 'w, + 's, + D, + F, + impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, + > + where + K: Ord, + { + let world = self.world; + + let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world); - /// Get next result from the query - #[inline(always)] - pub fn fetch_next(&mut self) -> Option> { // SAFETY: - // All arguments stem from self. - // We are limiting the returned reference to self, - // making sure this method cannot be called multiple times without getting rid - // of any previously returned unique references first, thus preventing aliasing. + // `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_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 { - Self::fetch_next_aliased_unchecked( - &mut self.entity_iter, - self.entities, - self.tables, - self.archetypes, - &mut self.fetch, - &mut self.filter, + QuerySortedManyIter::new( + world, self.query_state, + entity_iter, + world.last_change_tick(), + world.change_tick(), ) - .map(D::shrink) } } } @@ -1465,6 +2097,183 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> De } } +/// An [`Iterator`] over sorted query results of a [`QueryManyIter`]. +/// +/// This struct is created by the [`sort`](QueryManyIter), [`sort_unstable`](QueryManyIter), +/// [`sort_by`](QueryManyIter), [`sort_unstable_by`](QueryManyIter), [`sort_by_key`](QueryManyIter), +/// [`sort_unstable_by_key`](QueryManyIter), and [`sort_by_cached_key`](QueryManyIter) methods of [`QueryManyIter`]. +pub struct QuerySortedManyIter<'w, 's, D: QueryData, F: QueryFilter, 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> + QuerySortedManyIter<'w, 's, D, F, I> +{ + /// # 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, + ) -> QuerySortedManyIter<'w, 's, D, F, I> { + let fetch = D::init_fetch(world, &query_state.fetch_state, last_run, this_run); + QuerySortedManyIter { + 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 + /// The lifetime here is not restrictive enough for Fetch with &mut access, + /// as calling `fetch_next_aliased_unchecked` multiple times can produce multiple + /// references to the same component, leading to unique reference aliasing. + /// + /// It is always safe for shared access. + /// `entity` must stem from `self.entity_iter`, and not have been passed before. + #[inline(always)] + unsafe fn fetch_next_aliased_unchecked(&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) } + } + + /// Collects the internal [`I`](QuerySortedManyIter) once. + /// [`fetch_next`](QuerySortedManyIter) and [`fetch_next_back`](QuerySortedManyIter) require this to be called first. + #[inline(always)] + pub fn collect_inner(self) -> QuerySortedManyIter<'w, 's, D, F, IntoIter> { + QuerySortedManyIter { + entity_iter: self.entity_iter.collect::>().into_iter(), + entities: self.entities, + tables: self.tables, + archetypes: self.archetypes, + fetch: self.fetch, + query_state: self.query_state, + } + } +} + +impl<'w, 's, D: QueryData, F: QueryFilter> QuerySortedManyIter<'w, 's, D, F, IntoIter> { + /// Get next result from the query + /// [`collect_inner`](QuerySortedManyIter) needs to be called before this method becomes available. + /// This is done to prevent mutable aliasing. + #[inline(always)] + pub fn fetch_next(&mut self) -> Option> { + let entity = self.entity_iter.next()?; + + // SAFETY: + // We have collected the entity_iter once to drop all internal lens query item + // references. + // We are limiting the returned reference to self, + // making sure this method cannot be called multiple times without getting rid + // of any previously returned unique references first, thus preventing aliasing. + // `entity` is passed from `entity_iter` the first time. + unsafe { D::shrink(self.fetch_next_aliased_unchecked(entity)).into() } + } + + /// Get next result from the query + /// [`collect_inner`](QuerySortedManyIter) needs to be called before this method becomes available. + /// This is done to prevent mutable aliasing. + #[inline(always)] + pub fn fetch_next_back(&mut self) -> Option> { + let entity = self.entity_iter.next_back()?; + + // SAFETY: + // We have collected the entity_iter once to drop all internal lens query item + // references. + // We are limiting the returned reference to self, + // making sure this method cannot be called multiple times without getting rid + // of any previously returned unique references first, thus preventing aliasing. + // `entity` is passed from `entity_iter` the first time. + unsafe { D::shrink(self.fetch_next_aliased_unchecked(entity)).into() } + } +} + +impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator> Iterator + for QuerySortedManyIter<'w, 's, D, F, I> +{ + type Item = D::Item<'w>; + + #[inline(always)] + fn next(&mut self) -> Option { + let entity = self.entity_iter.next()?; + // SAFETY: + // It is safe to alias for ReadOnlyWorldQuery. + // `entity` is passed from `entity_iter` the first time. + unsafe { self.fetch_next_aliased_unchecked(entity).into() } + } + + fn size_hint(&self) -> (usize, Option) { + self.entity_iter.size_hint() + } +} + +impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: DoubleEndedIterator> + DoubleEndedIterator for QuerySortedManyIter<'w, 's, D, F, I> +{ + #[inline(always)] + fn next_back(&mut self) -> Option { + let entity = self.entity_iter.next_back()?; + // SAFETY: + // It is safe to alias for ReadOnlyWorldQuery. + // `entity` is passed from `entity_iter` the first time. + unsafe { self.fetch_next_aliased_unchecked(entity).into() } + } +} + +impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: ExactSizeIterator> + ExactSizeIterator for QuerySortedManyIter<'w, 's, D, F, I> +{ +} + +impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Debug + for QuerySortedManyIter<'w, 's, D, F, I> +{ + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("QuerySortedManyIter").finish() + } +} + /// An iterator over `K`-sized combinations of query items without repetition. /// /// A combination is an arrangement of a collection of items where order does not matter. @@ -1841,7 +2650,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, QuerySortedIter, QueryManyIter, QueryCombinationIter, QueryState::par_fold_init_unchecked_manual + // QueryIter, QueryIterationCursor, QuerySortedIter, QueryManyIter, QuerySortedManyIter, QueryCombinationIter, QueryState::par_fold_init_unchecked_manual /// # Safety /// `tables` and `archetypes` must belong to the same world that the [`QueryIterationCursor`] /// was initialized for. @@ -1998,7 +2807,7 @@ mod tests { #[allow(clippy::unnecessary_sort_by)] #[test] - fn query_sorts() { + fn query_iter_sorts() { let mut world = World::new(); let mut query = world.query::(); @@ -2067,7 +2876,7 @@ mod tests { #[test] #[should_panic] - fn query_sort_after_next() { + fn query_iter_sort_after_next() { let mut world = World::new(); world.spawn((A(0.),)); world.spawn((A(1.1),)); @@ -2097,7 +2906,7 @@ mod tests { #[test] #[should_panic] - fn query_sort_after_next_dense() { + fn query_iter_sort_after_next_dense() { let mut world = World::new(); world.spawn((Sparse(11),)); world.spawn((Sparse(22),)); @@ -2126,7 +2935,7 @@ mod tests { } #[test] - fn empty_query_sort_after_next_does_not_panic() { + fn empty_query_iter_sort_after_next_does_not_panic() { let mut world = World::new(); { let mut query = world.query::<(&A, &Sparse)>(); @@ -2183,4 +2992,103 @@ mod tests { ); } } + + #[allow(clippy::unnecessary_sort_by)] + #[test] + fn query_iter_many_sorts() { + let mut world = World::new(); + + let entity_list: &Vec<_> = &world + .spawn_batch([A(0.), A(1.), A(2.), A(3.), A(4.)]) + .collect(); + + let mut query = world.query::(); + + let sort = query + .iter_many(&world, entity_list) + .sort::() + .collect::>(); + + let sort_unstable = query + .iter_many(&world, entity_list) + .sort_unstable::() + .collect::>(); + + let sort_by = query + .iter_many(&world, entity_list) + .sort_by::(Ord::cmp) + .collect::>(); + + let sort_unstable_by = query + .iter_many(&world, entity_list) + .sort_unstable_by::(Ord::cmp) + .collect::>(); + + let sort_by_key = query + .iter_many(&world, entity_list) + .sort_by_key::(|&e| e) + .collect::>(); + + let sort_unstable_by_key = query + .iter_many(&world, entity_list) + .sort_unstable_by_key::(|&e| e) + .collect::>(); + + let sort_by_cached_key = query + .iter_many(&world, entity_list) + .sort_by_cached_key::(|&e| e) + .collect::>(); + + let mut sort_v2 = query.iter_many(&world, entity_list).collect::>(); + sort_v2.sort(); + + let mut sort_unstable_v2 = query.iter_many(&world, entity_list).collect::>(); + sort_unstable_v2.sort_unstable(); + + let mut sort_by_v2 = query.iter_many(&world, entity_list).collect::>(); + sort_by_v2.sort_by(Ord::cmp); + + let mut sort_unstable_by_v2 = query.iter_many(&world, entity_list).collect::>(); + sort_unstable_by_v2.sort_unstable_by(Ord::cmp); + + let mut sort_by_key_v2 = query.iter_many(&world, entity_list).collect::>(); + sort_by_key_v2.sort_by_key(|&e| e); + + let mut sort_unstable_by_key_v2 = query.iter_many(&world, entity_list).collect::>(); + sort_unstable_by_key_v2.sort_unstable_by_key(|&e| e); + + let mut sort_by_cached_key_v2 = query.iter_many(&world, entity_list).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_unstable_by_key, sort_unstable_by_key_v2); + assert_eq!(sort_by_cached_key, sort_by_cached_key_v2); + } + + #[test] + fn query_iter_many_sort_doesnt_panic_after_next() { + let mut world = World::new(); + + let entity_list: &Vec<_> = &world + .spawn_batch([A(0.), A(1.), A(2.), A(3.), A(4.)]) + .collect(); + + let mut query = world.query::(); + let mut iter = query.iter_many(&world, entity_list); + + _ = iter.next(); + + iter.sort::(); + + let mut query_2 = world.query::<&mut A>(); + let mut iter_2 = query_2.iter_many_mut(&mut world, entity_list); + + _ = iter_2.fetch_next(); + + iter_2.sort::(); + } }