Mirrors/bevy
2
0
Fork 0
mirror of https://github.com/bevyengine/bevy synced 2024-12-18 17:13:10 +00:00
Code Issues Projects Releases Packages Wiki Activity

implement the full set of sorts on QueryManyIter (#13443)

Browse source 
# Objective

~Blocked on #13417~

Motivation is the same as in #13417. If users can sort `QueryIter`, to
only makes sense to also allow them to use this functionality on
`QueryManyIter`.

## Solution

Also implement the sorts on `QueryManyIter`. 

The implementation of the sorts themselves are mostly the same as with
`QueryIter` in #13417.
They differ in that they re-use the `entity_iter` passed to the
`iter_many`, and internally call `iter_many_unchecked_manual` on the
lens `QueryState` with it.
These methods also return a different struct, `QuerySortedManyIter`,
because there is no longer a guarantee of unique entities.
`QuerySortedManyIter` implements the various `Iterator` traits for
read-only iteration, as `QueryManyIter` does + `DoubleEndedIterator`.
For mutable iteration, there is both a `fetch_next` and a
`fetch_next_back` method. However, they only become available after the
user calls `collect_inner` on `QuerySortedManyIter` first. This collects
the inner `entity_iter` (this is the sorted one, **not** the original
the user passed) to drop all query lens items to avoid aliasing.
When TAITs are available this `collect_inner` could be hidden away,
until then it is unfortunately not possible to elide this without either
regressing read-only iteration, or introducing a whole new type, mostly
being a copy of `QuerySortedIter`.

As a follow-up we could add a `entities_all_unique` method to check
whether the entity list consists of only unique entities, and then
return a `QuerySortedIter` from it (under opaque impl Trait if need be),
*allowing mutable `Iterator` trait iteration* over what was originally
an `iter_many` call.
Such a method can also be added to `QueryManyIter`, albeit needing a
separate, new return type.

## Testing

I've switched the third example/doc test under `sort` out for one that
shows the collect_inner/fetch_next_back functionality, otherwise the
examples are the same as in #13417, adjusted to use `iter_many` instead
of `iter`.

The `query-iter-many-sorts` test checks for equivalence to the
underlying sorts.
The test after shows that these sorts *do not* panic after
`fetch`/`fetch_next` calls.

## Changelog

Added `sort`, `sort_unstable`, `sort_by`, `sort_unstable_by`,
`sort_by_key`, `sort_by_cached_key` to `QueryManyIter`.
Added `QuerySortedManyIter`.
This commit is contained in:
Vic 2024-12-03 18:02:37 +01:00 committed by GitHub
parent d4883a9b5f
commit aa600ae95e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 914 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

920
crates/bevy_ecs/src/query/iter.rs
View file

@ -16,6 +16,7 @@ use core::{
}; };
use super::{QueryData, QueryFilter, ReadOnlyQueryData}; use super::{QueryData, QueryFilter, ReadOnlyQueryData};
use alloc::vec::IntoIter;
/// An [`Iterator`] over query results of a [`Query`](crate::system::Query). /// An [`Iterator`] over query results of a [`Query`](crate::system::Query).
/// ///
@ -1251,7 +1252,11 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item = Entity>> Debug
/// Entities that don't match the query are skipped. /// 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. /// 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<Item: Borrow<Entity>>> { pub struct QueryManyIter<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>
where
I::Item: Borrow<Entity>,
{
world: UnsafeWorldCell<'w>,
entity_iter: I, entity_iter: I,
entities: &'w Entities, entities: &'w Entities,
tables: &'w Tables, tables: &'w Tables,
@ -1277,6 +1282,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: Borrow<Entity>>>
let fetch = D::init_fetch(world, &query_state.fetch_state, last_run, this_run); 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); let filter = F::init_fetch(world, &query_state.filter_state, last_run, this_run);
QueryManyIter { QueryManyIter {
world,
query_state, query_state,
entities: world.entities(), entities: world.entities(),
archetypes: world.archetypes(), archetypes: world.archetypes(),
@ -1367,6 +1373,632 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: Borrow<Entity>>>
.map(D::shrink) .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<Entity> = 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<PartMarker>>) {
/// # let entity_list: Vec<Entity> = 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<Entity> = Vec::new();
/// // We need to collect the internal iterator before iterating mutably
/// let mut parent_query_iter = query.iter_many_mut(entity_list)
/// .sort::<Entity>()
/// .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<L: ReadOnlyQueryData + 'w>(
self,
) -> QuerySortedManyIter<
'w,
's,
D,
F,
impl ExactSizeIterator<Item = Entity> + 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<Entity> = 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<PartMarker>>) {
/// # let entity_list: Vec<Entity> = 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<L: ReadOnlyQueryData + 'w>(
self,
) -> QuerySortedManyIter<
'w,
's,
D,
F,
impl ExactSizeIterator<Item = Entity> + 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<Entity> = 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<Entity> = 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<L: ReadOnlyQueryData + 'w>(
self,
mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
) -> QuerySortedManyIter<
'w,
's,
D,
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
> {
let world = self.world;
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);
// 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<L: ReadOnlyQueryData + 'w>(
self,
mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
) -> QuerySortedManyIter<
'w,
's,
D,
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
> {
let world = self.world;
let query_lens_state = self.query_state.transmute_filtered::<(L, Entity), F>(world);
// 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<Entity> = 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<Entity> = 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<EntityRef, With<PartMarker>>) {
/// # let entity_list: Vec<Entity> = 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<EntityRef> = query.iter_many(entity_list)
/// .sort_by_key::<EntityRef, _>(|entity_ref| {
/// (
/// entity_ref.contains::<AvailableMarker>(),
/// entity_ref.get::<Rarity>()
/// )
/// })
/// .rev()
/// .collect();
/// }
/// # let mut schedule = Schedule::default();
/// # schedule.add_systems((system_1, system_2));
/// # schedule.run(&mut world);
/// ```
pub fn sort_by_key<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
) -> QuerySortedManyIter<
'w,
's,
D,
F,
impl ExactSizeIterator<Item = Entity> + 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<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
) -> QuerySortedManyIter<
'w,
's,
D,
F,
impl ExactSizeIterator<Item = Entity> + 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<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
) -> QuerySortedManyIter<
'w,
's,
D,
F,
impl ExactSizeIterator<Item = Entity> + 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_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 {
QuerySortedManyIter::new(
world,
self.query_state,
entity_iter,
world.last_change_tick(),
world.change_tick(),
)
}
}
} }
impl<'w, 's, D: QueryData, F: QueryFilter, I: DoubleEndedIterator<Item: Borrow<Entity>>> impl<'w, 's, D: QueryData, F: QueryFilter, I: DoubleEndedIterator<Item: Borrow<Entity>>>
@ -1465,6 +2097,183 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: Borrow<Entity>>> 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<Item = Entity>> {
entity_iter: I,
entities: &'w Entities,
tables: &'w Tables,
archetypes: &'w Archetypes,
fetch: D::Fetch<'w>,
query_state: &'s QueryState<D, F>,
}
impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item = Entity>>
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<EntityList: IntoIterator<IntoIter = I>>(
world: UnsafeWorldCell<'w>,
query_state: &'s QueryState<D, F>,
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<Entity>> {
QuerySortedManyIter {
entity_iter: self.entity_iter.collect::<Vec<_>>().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<Entity>> {
/// 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<D::Item<'_>> {
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<D::Item<'_>> {
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<Item = Entity>> Iterator
for QuerySortedManyIter<'w, 's, D, F, I>
{
type Item = D::Item<'w>;
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
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<usize>) {
self.entity_iter.size_hint()
}
}
impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: DoubleEndedIterator<Item = Entity>>
DoubleEndedIterator for QuerySortedManyIter<'w, 's, D, F, I>
{
#[inline(always)]
fn next_back(&mut self) -> Option<Self::Item> {
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<Item = Entity>>
ExactSizeIterator for QuerySortedManyIter<'w, 's, D, F, I>
{
}
impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item = Entity>> 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. /// 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. /// 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: // 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 /// # Safety
/// `tables` and `archetypes` must belong to the same world that the [`QueryIterationCursor`] /// `tables` and `archetypes` must belong to the same world that the [`QueryIterationCursor`]
/// was initialized for. /// was initialized for.
@ -1998,7 +2807,7 @@ mod tests {
#[allow(clippy::unnecessary_sort_by)] #[allow(clippy::unnecessary_sort_by)]
#[test] #[test]
fn query_sorts() { fn query_iter_sorts() {
let mut world = World::new(); let mut world = World::new();
let mut query = world.query::<Entity>(); let mut query = world.query::<Entity>();
@ -2067,7 +2876,7 @@ mod tests {
#[test] #[test]
#[should_panic] #[should_panic]
fn query_sort_after_next() { fn query_iter_sort_after_next() {
let mut world = World::new(); let mut world = World::new();
world.spawn((A(0.),)); world.spawn((A(0.),));
world.spawn((A(1.1),)); world.spawn((A(1.1),));
@ -2097,7 +2906,7 @@ mod tests {
#[test] #[test]
#[should_panic] #[should_panic]
fn query_sort_after_next_dense() { fn query_iter_sort_after_next_dense() {
let mut world = World::new(); let mut world = World::new();
world.spawn((Sparse(11),)); world.spawn((Sparse(11),));
world.spawn((Sparse(22),)); world.spawn((Sparse(22),));
@ -2126,7 +2935,7 @@ mod tests {
} }
#[test] #[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 world = World::new();
{ {
let mut query = world.query::<(&A, &Sparse)>(); 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::<Entity>();
let sort = query
.iter_many(&world, entity_list)
.sort::<Entity>()
.collect::<Vec<_>>();
let sort_unstable = query
.iter_many(&world, entity_list)
.sort_unstable::<Entity>()
.collect::<Vec<_>>();
let sort_by = query
.iter_many(&world, entity_list)
.sort_by::<Entity>(Ord::cmp)
.collect::<Vec<_>>();
let sort_unstable_by = query
.iter_many(&world, entity_list)
.sort_unstable_by::<Entity>(Ord::cmp)
.collect::<Vec<_>>();
let sort_by_key = query
.iter_many(&world, entity_list)
.sort_by_key::<Entity, _>(|&e| e)
.collect::<Vec<_>>();
let sort_unstable_by_key = query
.iter_many(&world, entity_list)
.sort_unstable_by_key::<Entity, _>(|&e| e)
.collect::<Vec<_>>();
let sort_by_cached_key = query
.iter_many(&world, entity_list)
.sort_by_cached_key::<Entity, _>(|&e| e)
.collect::<Vec<_>>();
let mut sort_v2 = query.iter_many(&world, entity_list).collect::<Vec<_>>();
sort_v2.sort();
let mut sort_unstable_v2 = query.iter_many(&world, entity_list).collect::<Vec<_>>();
sort_unstable_v2.sort_unstable();
let mut sort_by_v2 = query.iter_many(&world, entity_list).collect::<Vec<_>>();
sort_by_v2.sort_by(Ord::cmp);
let mut sort_unstable_by_v2 = query.iter_many(&world, entity_list).collect::<Vec<_>>();
sort_unstable_by_v2.sort_unstable_by(Ord::cmp);
let mut sort_by_key_v2 = query.iter_many(&world, entity_list).collect::<Vec<_>>();
sort_by_key_v2.sort_by_key(|&e| e);
let mut sort_unstable_by_key_v2 = query.iter_many(&world, entity_list).collect::<Vec<_>>();
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::<Vec<_>>();
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::<Entity>();
let mut iter = query.iter_many(&world, entity_list);
_ = iter.next();
iter.sort::<Entity>();
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::<Entity>();
}
} }