mirror of
https://github.com/bevyengine/bevy
synced 2024-12-18 17:13:10 +00:00
implement the full set of sorts on QueryManyIter (#13443)
# 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
GitHub
parent
d4883a9b5f
commit
aa600ae95e
1 changed files with 914 additions and 6 deletions
|
|
@ -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<Item = Entity>> 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<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,
|
||||
entities: &'w Entities,
|
||||
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 filter = F::init_fetch(world, &query_state.filter_state, last_run, this_run);
|
||||
QueryManyIter {
|
||||
world,
|
||||
query_state,
|
||||
entities: world.entities(),
|
||||
archetypes: world.archetypes(),
|
||||
|
|
@ -1367,6 +1373,632 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: Borrow<Entity>>>
|
|||
.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>>>
|
||||
|
|
@ -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.
|
||||
///
|
||||
/// 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::<Entity>();
|
||||
|
|
@ -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::<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>();
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue