diff --git a/crates/bevy_ecs/src/query/fetch.rs b/crates/bevy_ecs/src/query/fetch.rs index 0a13b6181974d..d985ff4be79c7 100644 --- a/crates/bevy_ecs/src/query/fetch.rs +++ b/crates/bevy_ecs/src/query/fetch.rs @@ -3,7 +3,7 @@ use crate::{ bundle::Bundle, change_detection::{MaybeLocation, Ticks, TicksMut}, component::{Component, ComponentId, Components, Mutable, StorageType, Tick}, - entity::{Entities, Entity, EntityLocation}, + entity::{Entities, Entity, EntityEquivalent, EntityLocation}, query::{Access, DebugCheckedUnwrap, FilteredAccess, WorldQuery}, storage::{ComponentSparseSet, Table, TableRow}, world::{ @@ -340,6 +340,20 @@ pub unsafe trait QueryData: WorldQuery { /// This must only be implemented for read-only [`QueryData`]'s. pub unsafe trait ReadOnlyQueryData: QueryData {} +/// A [`QueryData`] type that produces a [`EntityEquivalent`] item +/// equaling the `Entity` it is addressed by. +/// +/// # Safety +/// +/// [`::fetch`](QueryData::fetch) must always return an item whose `Entity` +/// equals the one the function was called with. +/// I.e.: `Self::fetch(fetch, entity, table_row).entity() == entity` always holds. +pub unsafe trait EntityEquivalentQueryData: QueryData +where + for<'w, 's> Self: QueryData: EntityEquivalent>, +{ +} + /// The item type returned when a [`WorldQuery`] is iterated over pub type QueryItem<'w, 's, Q> = ::Item<'w, 's>; /// The read-only variant of the item type returned when a [`QueryData`] is iterated over immutably @@ -435,6 +449,9 @@ unsafe impl QueryData for Entity { /// SAFETY: access is read only unsafe impl ReadOnlyQueryData for Entity {} +/// SAFETY: `entity()` returns `self`, and `fetch` returns `entity` +unsafe impl EntityEquivalentQueryData for Entity {} + impl ReleaseStateQueryData for Entity { fn release_state<'w>(item: Self::Item<'w, '_>) -> Self::Item<'w, 'static> { item @@ -814,6 +831,9 @@ unsafe impl<'a> QueryData for EntityRef<'a> { /// SAFETY: access is read only unsafe impl ReadOnlyQueryData for EntityRef<'_> {} +/// SAFETY: `entity()` returns the `entity` used to get the `UnsafeEntityCell`, and `fetch` constructs that from `entity` +unsafe impl EntityEquivalentQueryData for EntityRef<'_> {} + impl ReleaseStateQueryData for EntityRef<'_> { fn release_state<'w>(item: Self::Item<'w, '_>) -> Self::Item<'w, 'static> { item @@ -914,6 +934,9 @@ unsafe impl<'a> QueryData for EntityMut<'a> { } } +/// SAFETY: `entity()` returns the `entity` used to get the `UnsafeEntityCell`, and `fetch` constructs that from `entity` +unsafe impl EntityEquivalentQueryData for EntityMut<'_> {} + impl ReleaseStateQueryData for EntityMut<'_> { fn release_state<'w>(item: Self::Item<'w, '_>) -> Self::Item<'w, 'static> { item @@ -1044,6 +1067,9 @@ unsafe impl<'a> QueryData for FilteredEntityRef<'a> { /// SAFETY: Access is read-only. unsafe impl ReadOnlyQueryData for FilteredEntityRef<'_> {} +/// SAFETY: `entity()` returns the `entity` used to get the `UnsafeEntityCell`, and `fetch` constructs that from `entity` +unsafe impl EntityEquivalentQueryData for FilteredEntityRef<'_> {} + /// SAFETY: The accesses of `Self::ReadOnly` are a subset of the accesses of `Self` unsafe impl<'a> WorldQuery for FilteredEntityMut<'a> { type Fetch<'w> = (EntityFetch<'w>, Access); @@ -1163,6 +1189,9 @@ unsafe impl<'a> QueryData for FilteredEntityMut<'a> { } } +/// SAFETY: `entity()` returns the `entity` used to get the `UnsafeEntityCell`, and `fetch` constructs that from `entity` +unsafe impl EntityEquivalentQueryData for FilteredEntityMut<'_> {} + /// SAFETY: `EntityRefExcept` guards access to all components in the bundle `B` /// and populates `Access` values so that queries that conflict with this access /// are rejected. @@ -1273,6 +1302,9 @@ where /// components. unsafe impl<'a, B> ReadOnlyQueryData for EntityRefExcept<'a, B> where B: Bundle {} +/// SAFETY: `entity()` returns the `entity` used to get the `UnsafeEntityCell`, and `fetch` constructs that from `entity` +unsafe impl EntityEquivalentQueryData for EntityRefExcept<'_, B> {} + /// SAFETY: `EntityMutExcept` guards access to all components in the bundle `B` /// and populates `Access` values so that queries that conflict with this access /// are rejected. @@ -1380,6 +1412,9 @@ where } } +/// SAFETY: `entity()` returns the `entity` used to get the `UnsafeEntityCell`, and `fetch` constructs that from `entity` +unsafe impl EntityEquivalentQueryData for EntityMutExcept<'_, B> {} + /// SAFETY: /// `update_component_access` does nothing. /// This is sound because `fetch` does not access components. diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index eb49204434b6f..f88cc1dde75a6 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -1,21 +1,21 @@ -use super::{QueryData, QueryFilter, ReadOnlyQueryData}; +use super::{QueryData, QueryFilter, ReadOnlyQueryData, WorldQuery}; use crate::{ archetype::{Archetype, ArchetypeEntity, Archetypes}, - bundle::Bundle, component::Tick, entity::{ContainsEntity, Entities, Entity, EntityEquivalent, EntitySet, EntitySetIterator}, - query::{ArchetypeFilter, DebugCheckedUnwrap, QueryState, StorageId}, - storage::{Table, TableRow, Tables}, - world::{ - unsafe_world_cell::UnsafeWorldCell, EntityMut, EntityMutExcept, EntityRef, EntityRefExcept, - FilteredEntityMut, FilteredEntityRef, + query::{ + ArchetypeFilter, ConsumableQueryState, DebugCheckedUnwrap, EntityEquivalentQueryData, + QueryState, StorageId, }, + storage::{Table, TableRow, Tables}, + world::unsafe_world_cell::UnsafeWorldCell, }; use alloc::vec::Vec; use core::{ cmp::Ordering, fmt::{self, Debug, Formatter}, iter::FusedIterator, + marker::PhantomData, mem::MaybeUninit, ops::Range, }; @@ -25,32 +25,35 @@ use nonmax::NonMaxU32; /// /// This struct is created by the [`Query::iter`](crate::system::Query::iter) and /// [`Query::iter_mut`](crate::system::Query::iter_mut) methods. -pub struct QueryIter<'w, 's, D: QueryData, F: QueryFilter> { +pub struct QueryIter<'w, 's, S: ConsumableQueryState<'s>> { world: UnsafeWorldCell<'w>, tables: &'w Tables, archetypes: &'w Archetypes, - query_state: &'s QueryState, - cursor: QueryIterationCursor<'w, 's, D, F>, + query_state: S, + cursor: QueryIterationCursor<'w, 's, S>, } -impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { +impl<'w, 's, D: QueryData, F: QueryFilter, S: ConsumableQueryState<'s, Data = D, Filter = F>> + QueryIter<'w, 's, S> +{ /// # 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`. pub(crate) unsafe fn new( world: UnsafeWorldCell<'w>, - query_state: &'s QueryState, + query_state: S, last_run: Tick, this_run: Tick, ) -> Self { + // SAFETY: The invariants are upheld by the caller. + let cursor = unsafe { QueryIterationCursor::init(world, &query_state, last_run, this_run) }; QueryIter { world, query_state, // SAFETY: We only access table data that has been registered in `query_state`. tables: unsafe { &world.storages().tables }, archetypes: world.archetypes(), - // SAFETY: The invariants are upheld by the caller. - cursor: unsafe { QueryIterationCursor::init(world, query_state, last_run, this_run) }, + cursor, } } @@ -79,9 +82,10 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { /// } /// } /// ``` - pub fn remaining(&self) -> QueryIter<'w, 's, D, F> + pub fn remaining(&self) -> Self where D: ReadOnlyQueryData, + S: Copy, { QueryIter { world: self.world, @@ -116,12 +120,12 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { /// } /// } /// ``` - pub fn remaining_mut(&mut self) -> QueryIter<'_, 's, D, F> { + pub fn remaining_mut(&mut self) -> QueryIter<'_, '_, &QueryState> { QueryIter { world: self.world, tables: self.tables, archetypes: self.archetypes, - query_state: self.query_state, + query_state: &self.query_state, cursor: self.cursor.reborrow(), } } @@ -239,12 +243,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { // SAFETY: set_table was called prior. // Caller assures `row` in range of the current archetype. - let item = D::fetch( - &self.query_state.fetch_state, - &mut self.cursor.fetch, - *entity, - row, - ); + let item = self.query_state.fetch(&mut self.cursor.fetch, *entity, row); accum = func(accum, item); } @@ -308,8 +307,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { // SAFETY: set_archetype was called prior, `index` is an archetype index in range of the current archetype // Caller assures `index` in range of the current archetype. let item = unsafe { - D::fetch( - &self.query_state.fetch_state, + self.query_state.fetch( &mut self.cursor.fetch, archetype_entity.id(), archetype_entity.table_row(), @@ -384,12 +382,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { // SAFETY: set_table was called prior. // Caller assures `row` in range of the current archetype. - let item = D::fetch( - &self.query_state.fetch_state, - &mut self.cursor.fetch, - entity, - row, - ); + let item = self.query_state.fetch(&mut self.cursor.fetch, entity, row); accum = func(accum, item); } @@ -513,8 +506,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { ) -> QuerySortedIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > where @@ -570,8 +562,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { ) -> QuerySortedIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > where @@ -635,8 +626,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { ) -> QuerySortedIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > { self.sort_impl::(move |keyed_query| { @@ -667,8 +657,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { ) -> QuerySortedIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > { self.sort_impl::(move |keyed_query| { @@ -759,8 +748,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { ) -> QuerySortedIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > where @@ -792,8 +780,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { ) -> QuerySortedIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > where @@ -827,8 +814,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { ) -> QuerySortedIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > where @@ -857,8 +843,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { ) -> QuerySortedIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > { // On the first successful iteration of `QueryIterationCursor`, `archetype_entities` or `table_entities` @@ -902,7 +887,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { } } -impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for QueryIter<'w, 's, D, F> { +impl<'w, 's, D: QueryData, F: QueryFilter, S: ConsumableQueryState<'s, Data = D, Filter = F>> + Iterator for QueryIter<'w, 's, S> +{ type Item = D::Item<'w, 's>; #[inline(always)] @@ -912,7 +899,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for QueryIter<'w, 's, D, F> // `query_state` is the state that was passed to `QueryIterationCursor::init`. unsafe { self.cursor - .next(self.tables, self.archetypes, self.query_state) + .next(self.tables, self.archetypes, &self.query_state) } } @@ -935,7 +922,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for QueryIter<'w, 's, D, F> accum = func(accum, item); } - for id in self.cursor.storage_id_iter.clone().copied() { + for id in self.cursor.storage_id_iter.clone() { // SAFETY: // - The range(None) is equivalent to [0, storage.entity_count) accum = unsafe { self.fold_over_storage_range(accum, &mut func, id, None) }; @@ -945,48 +932,24 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for QueryIter<'w, 's, D, F> } // This is correct as [`QueryIter`] always returns `None` once exhausted. -impl<'w, 's, D: QueryData, F: QueryFilter> FusedIterator for QueryIter<'w, 's, D, F> {} - -// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once. -unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator for QueryIter<'w, 's, Entity, F> {} - -// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once. -unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator for QueryIter<'w, 's, EntityRef<'_>, F> {} - -// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once. -unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator for QueryIter<'w, 's, EntityMut<'_>, F> {} - -// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once. -unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator - for QueryIter<'w, 's, FilteredEntityRef<'_>, F> -{ -} - -// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once. -unsafe impl<'w, 's, F: QueryFilter> EntitySetIterator - for QueryIter<'w, 's, FilteredEntityMut<'_>, F> -{ -} - -// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once. -unsafe impl<'w, 's, F: QueryFilter, B: Bundle> EntitySetIterator - for QueryIter<'w, 's, EntityRefExcept<'_, B>, F> -{ -} +impl<'w, 's, S: ConsumableQueryState<'s>> FusedIterator for QueryIter<'w, 's, S> {} -// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once. -unsafe impl<'w, 's, F: QueryFilter, B: Bundle> EntitySetIterator - for QueryIter<'w, 's, EntityMutExcept<'_, B>, F> +// SAFETY: [`QueryIter`] is guaranteed to return every matching entity once and only once, +// and [`EntityEquivalentQueryData`] ensures that they return the original `Entity`. +unsafe impl<'w, 's, S: ConsumableQueryState<'s, Data: EntityEquivalentQueryData>> EntitySetIterator + for QueryIter<'w, 's, S> { } -impl<'w, 's, D: QueryData, F: QueryFilter> Debug for QueryIter<'w, 's, D, F> { +impl<'w, 's, S: ConsumableQueryState<'s>> Debug for QueryIter<'w, 's, S> { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { f.debug_struct("QueryIter").finish() } } -impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter> Clone for QueryIter<'w, 's, D, F> { +impl<'w, 's, S: ConsumableQueryState<'s, Data: ReadOnlyQueryData> + Copy> Clone + for QueryIter<'w, 's, S> +{ fn clone(&self) -> Self { self.remaining() } @@ -997,7 +960,7 @@ impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter> Clone for QueryIter<'w, 's, D /// This struct is created by the [`QueryIter::sort`], [`QueryIter::sort_unstable`], /// [`QueryIter::sort_by`], [`QueryIter::sort_unstable_by`], [`QueryIter::sort_by_key`], /// [`QueryIter::sort_unstable_by_key`], and [`QueryIter::sort_by_cached_key`] methods. -pub struct QuerySortedIter<'w, 's, D: QueryData, F: QueryFilter, I> +pub struct QuerySortedIter<'w, 's, S: ConsumableQueryState<'s>, I> where I: Iterator, { @@ -1005,11 +968,13 @@ where entities: &'w Entities, tables: &'w Tables, archetypes: &'w Archetypes, - fetch: D::Fetch<'w>, - query_state: &'s QueryState, + fetch: ::Fetch<'w>, + query_state: S, + marker: PhantomData<&'s ()>, } -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> QuerySortedIter<'w, 's, D, F, I> +impl<'w, 's, D: QueryData, S: ConsumableQueryState<'s, Data = D>, I: Iterator> + QuerySortedIter<'w, 's, S, I> where I: Iterator, { @@ -1019,14 +984,15 @@ where /// - `entity_list` must only contain unique entities or be empty. pub(crate) unsafe fn new>( world: UnsafeWorldCell<'w>, - query_state: &'s QueryState, + query_state: S, entity_list: EntityList, last_run: Tick, this_run: Tick, - ) -> QuerySortedIter<'w, 's, D, F, I> { + ) -> Self { let fetch = D::init_fetch(world, &query_state.fetch_state, last_run, this_run); QuerySortedIter { query_state, + marker: PhantomData, entities: world.entities(), archetypes: world.archetypes(), // SAFETY: We only access table data that has been registered in `query_state`. @@ -1070,18 +1036,14 @@ where // - 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( - &self.query_state.fetch_state, - &mut self.fetch, - entity, - location.table_row, - ) + self.query_state + .fetch(&mut self.fetch, entity, location.table_row) } } } -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Iterator - for QuerySortedIter<'w, 's, D, F, I> +impl<'w, 's, D: QueryData, S: ConsumableQueryState<'s, Data = D>, I: Iterator> Iterator + for QuerySortedIter<'w, 's, S, I> where I: Iterator, { @@ -1099,8 +1061,8 @@ where } } -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> DoubleEndedIterator - for QuerySortedIter<'w, 's, D, F, I> +impl<'w, 's, S: ConsumableQueryState<'s>, I: Iterator> DoubleEndedIterator + for QuerySortedIter<'w, 's, S, I> where I: DoubleEndedIterator, { @@ -1112,16 +1074,16 @@ where } } -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> ExactSizeIterator - for QuerySortedIter<'w, 's, D, F, I> +impl<'w, 's, S: ConsumableQueryState<'s>, I: Iterator> ExactSizeIterator + for QuerySortedIter<'w, 's, S, I> where I: ExactSizeIterator, { } // This is correct as [`QuerySortedIter`] returns `None` once exhausted if `entity_iter` does. -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> FusedIterator - for QuerySortedIter<'w, 's, D, F, I> +impl<'w, 's, S: ConsumableQueryState<'s>, I: Iterator> FusedIterator + for QuerySortedIter<'w, 's, S, I> where I: FusedIterator, { @@ -1129,14 +1091,18 @@ where // SAFETY: // `I` stems from a collected and sorted `EntitySetIterator` ([`QueryIter`]). -// Fetching unique entities maintains uniqueness. -unsafe impl<'w, 's, F: QueryFilter, I: Iterator> EntitySetIterator - for QuerySortedIter<'w, 's, Entity, F, I> +// [`EntityEquivalentQueryData`] ensures that unique items return unique `Entity`s. +unsafe impl< + 'w, + 's, + S: ConsumableQueryState<'s, Data: EntityEquivalentQueryData>, + I: Iterator, + > EntitySetIterator for QuerySortedIter<'w, 's, S, I> { } -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Debug - for QuerySortedIter<'w, 's, D, F, I> +impl<'w, 's, S: ConsumableQueryState<'s>, I: Iterator> Debug + for QuerySortedIter<'w, 's, S, I> { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { f.debug_struct("QuerySortedIter").finish() @@ -1149,36 +1115,43 @@ 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, S: ConsumableQueryState<'s>, I: Iterator> { world: UnsafeWorldCell<'w>, entity_iter: I, entities: &'w Entities, tables: &'w Tables, archetypes: &'w Archetypes, - fetch: D::Fetch<'w>, - filter: F::Fetch<'w>, - query_state: &'s QueryState, + fetch: ::Fetch<'w>, + filter: ::Fetch<'w>, + query_state: S, + marker: PhantomData<&'s ()>, } -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> - QueryManyIter<'w, 's, D, F, I> +impl< + 'w, + 's, + D: QueryData, + F: QueryFilter, + S: ConsumableQueryState<'s, Data = D, Filter = F>, + I: Iterator, + > QueryManyIter<'w, 's, S, 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`. pub(crate) unsafe fn new>( world: UnsafeWorldCell<'w>, - query_state: &'s QueryState, + query_state: S, entity_list: EntityList, last_run: Tick, this_run: Tick, - ) -> QueryManyIter<'w, 's, D, F, I> { + ) -> Self { 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, + marker: PhantomData, entities: world.entities(), archetypes: world.archetypes(), // SAFETY: We only access table data that has been registered in `query_state`. @@ -1206,7 +1179,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> archetypes: &'w Archetypes, fetch: &mut D::Fetch<'w>, filter: &mut F::Fetch<'w>, - query_state: &'s QueryState, + query_state: &S, ) -> Option> { for entity_borrow in entity_iter { let entity = entity_borrow.entity(); @@ -1248,9 +1221,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> // 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(&query_state.fetch_state, fetch, entity, location.table_row) - }); + return Some(unsafe { query_state.fetch(fetch, entity, location.table_row) }); } } None @@ -1272,7 +1243,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> self.archetypes, &mut self.fetch, &mut self.filter, - self.query_state, + &self.query_state, ) .map(D::shrink) } @@ -1377,8 +1348,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> ) -> QuerySortedManyIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > where @@ -1435,8 +1405,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> ) -> QuerySortedManyIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > where @@ -1501,8 +1470,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> ) -> QuerySortedManyIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > { self.sort_impl::(move |keyed_query| { @@ -1532,8 +1500,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> ) -> QuerySortedManyIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > { self.sort_impl::(move |keyed_query| { @@ -1626,8 +1593,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> ) -> QuerySortedManyIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > where @@ -1658,8 +1624,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> ) -> QuerySortedManyIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > where @@ -1692,8 +1657,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> ) -> QuerySortedManyIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > where @@ -1721,8 +1685,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> ) -> QuerySortedManyIter< 'w, 's, - D, - F, + S, impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, > { let world = self.world; @@ -1762,8 +1725,13 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> } } -impl<'w, 's, D: QueryData, F: QueryFilter, I: DoubleEndedIterator> - QueryManyIter<'w, 's, D, F, I> +impl< + 'w, + 's, + D: QueryData, + S: ConsumableQueryState<'s, Data = D>, + I: DoubleEndedIterator, + > QueryManyIter<'w, 's, S, I> { /// Get next result from the back of the query #[inline(always)] @@ -1781,15 +1749,20 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: DoubleEndedIterator> Iterator - for QueryManyIter<'w, 's, D, F, I> +impl< + 'w, + 's, + D: ReadOnlyQueryData, + S: ConsumableQueryState<'s, Data = D>, + I: Iterator, + > Iterator for QueryManyIter<'w, 's, S, I> { type Item = D::Item<'w, 's>; @@ -1806,7 +1779,7 @@ impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator, I: DoubleEndedIterator, - > DoubleEndedIterator for QueryManyIter<'w, 's, D, F, I> + > DoubleEndedIterator for QueryManyIter<'w, 's, S, I> { #[inline(always)] fn next_back(&mut self) -> Option { @@ -1838,26 +1811,34 @@ impl< self.archetypes, &mut self.fetch, &mut self.filter, - self.query_state, + &self.query_state, ) } } } // This is correct as [`QueryManyIter`] always returns `None` once exhausted. -impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator> - FusedIterator for QueryManyIter<'w, 's, D, F, I> +impl< + 'w, + 's, + S: ConsumableQueryState<'s, Data: ReadOnlyQueryData>, + I: Iterator, + > FusedIterator for QueryManyIter<'w, 's, S, I> { } -// SAFETY: Fetching unique entities maintains uniqueness. -unsafe impl<'w, 's, F: QueryFilter, I: EntitySetIterator> EntitySetIterator - for QueryManyIter<'w, 's, Entity, F, I> +// SAFETY: [`EntityEquivalentQueryData`] ensures that unique items return unique `Entity`s. +unsafe impl< + 'w, + 's, + S: ConsumableQueryState<'s, Data: EntityEquivalentQueryData + ReadOnlyQueryData>, + I: EntitySetIterator, + > EntitySetIterator for QueryManyIter<'w, 's, S, I> { } -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Debug - for QueryManyIter<'w, 's, D, F, I> +impl<'w, 's, S: ConsumableQueryState<'s>, I: Iterator> Debug + for QueryManyIter<'w, 's, S, I> { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { f.debug_struct("QueryManyIter").finish() @@ -1877,23 +1858,23 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> /// [`iter_many_unique`]: crate::system::Query::iter_many /// [`iter_many_unique_mut`]: crate::system::Query::iter_many_mut /// [`Query`]: crate::system::Query -pub struct QueryManyUniqueIter<'w, 's, D: QueryData, F: QueryFilter, I: EntitySetIterator>( - QueryManyIter<'w, 's, D, F, I>, +pub struct QueryManyUniqueIter<'w, 's, S: ConsumableQueryState<'s>, I: EntitySetIterator>( + QueryManyIter<'w, 's, S, I>, ); -impl<'w, 's, D: QueryData, F: QueryFilter, I: EntitySetIterator> - QueryManyUniqueIter<'w, 's, D, F, I> +impl<'w, 's, D: QueryData, S: ConsumableQueryState<'s, Data = D>, I: EntitySetIterator> + QueryManyUniqueIter<'w, 's, S, 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`. pub(crate) unsafe fn new>( world: UnsafeWorldCell<'w>, - query_state: &'s QueryState, + query_state: S, entity_list: EntityList, last_run: Tick, this_run: Tick, - ) -> QueryManyUniqueIter<'w, 's, D, F, I> { + ) -> Self { QueryManyUniqueIter(QueryManyIter::new( world, query_state, @@ -1904,8 +1885,8 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: EntitySetIterator> } } -impl<'w, 's, D: QueryData, F: QueryFilter, I: EntitySetIterator> Iterator - for QueryManyUniqueIter<'w, 's, D, F, I> +impl<'w, 's, D: QueryData, S: ConsumableQueryState<'s, Data = D>, I: EntitySetIterator> Iterator + for QueryManyUniqueIter<'w, 's, S, I> { type Item = D::Item<'w, 's>; @@ -1913,14 +1894,14 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: EntitySetIterator> Iterator fn next(&mut self) -> Option { // SAFETY: Entities are guaranteed to be unique, thus do not alias. unsafe { - QueryManyIter::<'w, 's, D, F, I>::fetch_next_aliased_unchecked( + QueryManyIter::<'w, 's, S, I>::fetch_next_aliased_unchecked( &mut self.0.entity_iter, self.0.entities, self.0.tables, self.0.archetypes, &mut self.0.fetch, &mut self.0.filter, - self.0.query_state, + &self.0.query_state, ) } } @@ -1932,19 +1913,23 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: EntitySetIterator> Iterator } // This is correct as [`QueryManyIter`] always returns `None` once exhausted. -impl<'w, 's, D: QueryData, F: QueryFilter, I: EntitySetIterator> FusedIterator - for QueryManyUniqueIter<'w, 's, D, F, I> +impl<'w, 's, S: ConsumableQueryState<'s>, I: EntitySetIterator> FusedIterator + for QueryManyUniqueIter<'w, 's, S, I> { } -// SAFETY: Fetching unique entities maintains uniqueness. -unsafe impl<'w, 's, F: QueryFilter, I: EntitySetIterator> EntitySetIterator - for QueryManyUniqueIter<'w, 's, Entity, F, I> +// SAFETY: [`EntityEquivalentQueryData`] ensures that unique items return unique `Entity`s. +unsafe impl< + 'w, + 's, + S: ConsumableQueryState<'s, Data: EntityEquivalentQueryData>, + I: EntitySetIterator, + > EntitySetIterator for QueryManyUniqueIter<'w, 's, S, I> { } -impl<'w, 's, D: QueryData, F: QueryFilter, I: EntitySetIterator> Debug - for QueryManyUniqueIter<'w, 's, D, F, I> +impl<'w, 's, S: ConsumableQueryState<'s>, I: EntitySetIterator> Debug + for QueryManyUniqueIter<'w, 's, S, I> { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { f.debug_struct("QueryManyUniqueIter").finish() @@ -1956,17 +1941,18 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: EntitySetIterator> Debug /// 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> { +pub struct QuerySortedManyIter<'w, 's, S: ConsumableQueryState<'s>, I: Iterator> { entity_iter: I, entities: &'w Entities, tables: &'w Tables, archetypes: &'w Archetypes, - fetch: D::Fetch<'w>, - query_state: &'s QueryState, + fetch: ::Fetch<'w>, + query_state: S, + marker: PhantomData<&'s ()>, } -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> - QuerySortedManyIter<'w, 's, D, F, I> +impl<'w, 's, D: QueryData, S: ConsumableQueryState<'s, Data = D>, I: Iterator> + QuerySortedManyIter<'w, 's, S, I> { /// # Safety /// - `world` must have permission to access any of the components registered in `query_state`. @@ -1974,14 +1960,15 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> /// - `entity_list` must only contain unique entities or be empty. pub(crate) unsafe fn new>( world: UnsafeWorldCell<'w>, - query_state: &'s QueryState, + query_state: S, entity_list: EntityList, last_run: Tick, this_run: Tick, - ) -> QuerySortedManyIter<'w, 's, D, F, I> { + ) -> QuerySortedManyIter<'w, 's, S, I> { let fetch = D::init_fetch(world, &query_state.fetch_state, last_run, this_run); QuerySortedManyIter { query_state, + marker: PhantomData, entities: world.entities(), archetypes: world.archetypes(), // SAFETY: We only access table data that has been registered in `query_state`. @@ -2030,12 +2017,8 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> // - 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( - &self.query_state.fetch_state, - &mut self.fetch, - entity, - location.table_row, - ) + self.query_state + .fetch(&mut self.fetch, entity, location.table_row) } } @@ -2055,8 +2038,13 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> } } -impl<'w, 's, D: QueryData, F: QueryFilter, I: DoubleEndedIterator> - QuerySortedManyIter<'w, 's, D, F, I> +impl< + 'w, + 's, + D: QueryData, + S: ConsumableQueryState<'s, Data = D>, + I: DoubleEndedIterator, + > QuerySortedManyIter<'w, 's, S, I> { /// Get next result from the query #[inline(always)] @@ -2074,8 +2062,13 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: DoubleEndedIterator } } -impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator> Iterator - for QuerySortedManyIter<'w, 's, D, F, I> +impl< + 'w, + 's, + D: ReadOnlyQueryData, + S: ConsumableQueryState<'s, Data = D>, + I: Iterator, + > Iterator for QuerySortedManyIter<'w, 's, S, I> { type Item = D::Item<'w, 's>; @@ -2093,8 +2086,12 @@ impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator> I } } -impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: DoubleEndedIterator> - DoubleEndedIterator for QuerySortedManyIter<'w, 's, D, F, I> +impl< + 'w, + 's, + S: ConsumableQueryState<'s, Data: ReadOnlyQueryData>, + I: DoubleEndedIterator, + > DoubleEndedIterator for QuerySortedManyIter<'w, 's, S, I> { #[inline(always)] fn next_back(&mut self) -> Option { @@ -2106,13 +2103,17 @@ impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: DoubleEndedIterator> - ExactSizeIterator for QuerySortedManyIter<'w, 's, D, F, I> +impl< + 'w, + 's, + S: ConsumableQueryState<'s, Data: ReadOnlyQueryData>, + I: ExactSizeIterator, + > ExactSizeIterator for QuerySortedManyIter<'w, 's, S, I> { } -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Debug - for QuerySortedManyIter<'w, 's, D, F, I> +impl<'w, 's, S: ConsumableQueryState<'s>, I: Iterator> Debug + for QuerySortedManyIter<'w, 's, S, I> { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { f.debug_struct("QuerySortedManyIter").finish() @@ -2182,20 +2183,22 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Debug /// [`Query`]: crate::system::Query /// [`Query::iter_combinations`]: crate::system::Query::iter_combinations /// [`Query::iter_combinations_mut`]: crate::system::Query::iter_combinations_mut -pub struct QueryCombinationIter<'w, 's, D: QueryData, F: QueryFilter, const K: usize> { +pub struct QueryCombinationIter<'w, 's, S: ConsumableQueryState<'s>, const K: usize> { tables: &'w Tables, archetypes: &'w Archetypes, - query_state: &'s QueryState, - cursors: [QueryIterationCursor<'w, 's, D, F>; K], + query_state: S, + cursors: [QueryIterationCursor<'w, 's, S>; K], } -impl<'w, 's, D: QueryData, F: QueryFilter, const K: usize> QueryCombinationIter<'w, 's, D, F, K> { +impl<'w, 's, D: QueryData, S: ConsumableQueryState<'s, Data = D>, const K: usize> + QueryCombinationIter<'w, 's, S, K> +{ /// # 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`. pub(crate) unsafe fn new( world: UnsafeWorldCell<'w>, - query_state: &'s QueryState, + query_state: S, last_run: Tick, this_run: Tick, ) -> Self { @@ -2203,20 +2206,18 @@ impl<'w, 's, D: QueryData, F: QueryFilter, const K: usize> QueryCombinationIter< // Initialize array with cursors. // There is no FromIterator on arrays, so instead initialize it manually with MaybeUninit - let mut array: MaybeUninit<[QueryIterationCursor<'w, 's, D, F>; K]> = MaybeUninit::uninit(); - let ptr = array - .as_mut_ptr() - .cast::>(); + let mut array: MaybeUninit<[QueryIterationCursor<'w, 's, S>; K]> = MaybeUninit::uninit(); + let ptr = array.as_mut_ptr().cast::>(); ptr.write(QueryIterationCursor::init( world, - query_state, + &query_state, last_run, this_run, )); for slot in (1..K).map(|offset| ptr.add(offset)) { slot.write(QueryIterationCursor::init_empty( world, - query_state, + &query_state, last_run, this_run, )); @@ -2247,11 +2248,12 @@ impl<'w, 's, D: QueryData, F: QueryFilter, const K: usize> QueryCombinationIter< // Make cursor in index `j` for all `j` in `[i, K)` a copy of `c` advanced `j-i+1` times. // If no such `c` exists, return `None` 'outer: for i in (0..K).rev() { - match self.cursors[i].next(self.tables, self.archetypes, self.query_state) { + match self.cursors[i].next(self.tables, self.archetypes, &self.query_state) { Some(_) => { for j in (i + 1)..K { self.cursors[j] = self.cursors[j - 1].clone(); - match self.cursors[j].next(self.tables, self.archetypes, self.query_state) { + match self.cursors[j].next(self.tables, self.archetypes, &self.query_state) + { Some(_) => {} None if i > 0 => continue 'outer, None => return None, @@ -2269,7 +2271,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, const K: usize> QueryCombinationIter< let ptr = values.as_mut_ptr().cast::>(); for (offset, cursor) in self.cursors.iter_mut().enumerate() { ptr.add(offset) - .write(cursor.peek_last(self.query_state).unwrap()); + .write(cursor.peek_last(&self.query_state).unwrap()); } Some(values.assume_init()) @@ -2291,8 +2293,14 @@ impl<'w, 's, D: QueryData, F: QueryFilter, const K: usize> QueryCombinationIter< // Iterator type is intentionally implemented only for read-only access. // Doing so for mutable references would be unsound, because calling `next` // multiple times would allow multiple owned references to the same data to exist. -impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, const K: usize> Iterator - for QueryCombinationIter<'w, 's, D, F, K> +impl< + 'w, + 's, + D: ReadOnlyQueryData, + F: QueryFilter, + S: ConsumableQueryState<'s, Data = D, Filter = F>, + const K: usize, + > Iterator for QueryCombinationIter<'w, 's, S, K> { type Item = [D::Item<'w, 's>; K]; @@ -2334,9 +2342,9 @@ impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, const K: usize> Iterator } } -impl<'w, 's, D: QueryData, F: QueryFilter> ExactSizeIterator for QueryIter<'w, 's, D, F> +impl<'w, 's, S: ConsumableQueryState<'s>> ExactSizeIterator for QueryIter<'w, 's, S> where - F: ArchetypeFilter, + S::Filter: ArchetypeFilter, { fn len(&self) -> usize { self.size_hint().0 @@ -2344,34 +2352,34 @@ where } // This is correct as [`QueryCombinationIter`] always returns `None` once exhausted. -impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, const K: usize> FusedIterator - for QueryCombinationIter<'w, 's, D, F, K> +impl<'w, 's, S: ConsumableQueryState<'s, Data: ReadOnlyQueryData>, const K: usize> FusedIterator + for QueryCombinationIter<'w, 's, S, K> { } -impl<'w, 's, D: QueryData, F: QueryFilter, const K: usize> Debug - for QueryCombinationIter<'w, 's, D, F, K> +impl<'w, 's, S: ConsumableQueryState<'s>, const K: usize> Debug + for QueryCombinationIter<'w, 's, S, K> { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { f.debug_struct("QueryCombinationIter").finish() } } -struct QueryIterationCursor<'w, 's, D: QueryData, F: QueryFilter> { +struct QueryIterationCursor<'w, 's, S: ConsumableQueryState<'s>> { // whether the query iteration is dense or not. Mirrors QueryState's `is_dense` field. is_dense: bool, - storage_id_iter: core::slice::Iter<'s, StorageId>, + storage_id_iter: S::StorageIter, table_entities: &'w [Entity], archetype_entities: &'w [ArchetypeEntity], - fetch: D::Fetch<'w>, - filter: F::Fetch<'w>, + fetch: ::Fetch<'w>, + filter: ::Fetch<'w>, // length of the table or length of the archetype, depending on whether both `D`'s and `F`'s fetches are dense current_len: u32, // either table row or archetype index, depending on whether both `D`'s and `F`'s fetches are dense current_row: u32, } -impl Clone for QueryIterationCursor<'_, '_, D, F> { +impl<'s, S: ConsumableQueryState<'s>> Clone for QueryIterationCursor<'_, 's, S> { fn clone(&self) -> Self { Self { is_dense: self.is_dense, @@ -2386,18 +2394,20 @@ impl Clone for QueryIterationCursor<'_, '_, D, F> } } -impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { +impl<'w, 's, D: QueryData, F: QueryFilter, S: ConsumableQueryState<'s, Data = D, Filter = F>> + QueryIterationCursor<'w, 's, S> +{ /// # 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`. unsafe fn init_empty( world: UnsafeWorldCell<'w>, - query_state: &'s QueryState, + query_state: &S, last_run: Tick, this_run: Tick, ) -> Self { QueryIterationCursor { - storage_id_iter: [].iter(), + storage_id_iter: S::StorageIter::default(), ..Self::init(world, query_state, last_run, this_run) } } @@ -2407,7 +2417,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { /// - `world` must be the same one used to initialize `query_state`. unsafe fn init( world: UnsafeWorldCell<'w>, - query_state: &'s QueryState, + query_state: &S, last_run: Tick, this_run: Tick, ) -> Self { @@ -2418,21 +2428,21 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { filter, table_entities: &[], archetype_entities: &[], - storage_id_iter: query_state.matched_storage_ids.iter(), + storage_id_iter: query_state.storage_ids(), is_dense: query_state.is_dense, current_len: 0, current_row: 0, } } - fn reborrow(&mut self) -> QueryIterationCursor<'_, 's, D, F> { + fn reborrow(&mut self) -> QueryIterationCursor<'_, '_, &'_ QueryState> { QueryIterationCursor { is_dense: self.is_dense, fetch: D::shrink_fetch(self.fetch.clone()), filter: F::shrink_fetch(self.filter.clone()), table_entities: self.table_entities, archetype_entities: self.archetype_entities, - storage_id_iter: self.storage_id_iter.clone(), + storage_id_iter: S::reborrow_storage_ids(&self.storage_id_iter), current_len: self.current_len, current_row: self.current_row, } @@ -2444,7 +2454,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { /// The result of `next` and any previous calls to `peek_last` with this row must have been /// dropped to prevent aliasing mutable references. #[inline] - unsafe fn peek_last(&mut self, query_state: &'s QueryState) -> Option> { + unsafe fn peek_last(&mut self, query_state: &S) -> Option> { if self.current_row > 0 { let index = self.current_row - 1; if self.is_dense { @@ -2454,8 +2464,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { // - `set_table` must have been called previously either in `next` or before it. // - `*entity` and `index` are in the current table. unsafe { - Some(D::fetch( - &query_state.fetch_state, + Some(query_state.fetch( &mut self.fetch, *entity, // SAFETY: This is from an exclusive range, so it can't be max. @@ -2470,8 +2479,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { // - `set_archetype` must have been called previously either in `next` or before it. // - `archetype_entity.id()` and `archetype_entity.table_row()` are in the current archetype. unsafe { - Some(D::fetch( - &query_state.fetch_state, + Some(query_state.fetch( &mut self.fetch, archetype_entity.id(), archetype_entity.table_row(), @@ -2512,7 +2520,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { &mut self, tables: &'w Tables, archetypes: &'w Archetypes, - query_state: &'s QueryState, + query_state: &S, ) -> Option> { if self.is_dense { loop { @@ -2550,8 +2558,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { // - `current_row` must be a table row in range of the current table, // because if it was not, then the above would have been executed. // - fetch is only called once for each `entity`. - let item = - unsafe { D::fetch(&query_state.fetch_state, &mut self.fetch, *entity, row) }; + let item = unsafe { query_state.fetch(&mut self.fetch, *entity, row) }; self.current_row += 1; return Some(item); @@ -2608,8 +2615,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> { // because if it was not, then the if above would have been executed. // - fetch is only called once for each `archetype_entity`. let item = unsafe { - D::fetch( - &query_state.fetch_state, + query_state.fetch( &mut self.fetch, archetype_entity.id(), archetype_entity.table_row(), diff --git a/crates/bevy_ecs/src/query/mod.rs b/crates/bevy_ecs/src/query/mod.rs index 7c1487fde4e72..1d0e74bb97130 100644 --- a/crates/bevy_ecs/src/query/mod.rs +++ b/crates/bevy_ecs/src/query/mod.rs @@ -167,7 +167,7 @@ mod tests { F: ArchetypeFilter, { let mut query = world.query_filtered::(); - let query_type = type_name::>(); + let query_type = type_name::, K>>(); let iter = query.iter_combinations::(world); assert_all_sizes_iterator_equal(iter, expected_size, 0, query_type); let iter = query.iter_combinations::(world); diff --git a/crates/bevy_ecs/src/query/par_iter.rs b/crates/bevy_ecs/src/query/par_iter.rs index b8d8618fa5bf1..a2cdf32196b4f 100644 --- a/crates/bevy_ecs/src/query/par_iter.rs +++ b/crates/bevy_ecs/src/query/par_iter.rs @@ -1,3 +1,5 @@ +use core::marker::PhantomData; + use crate::{ batching::BatchingStrategy, component::Tick, @@ -5,7 +7,7 @@ use crate::{ world::unsafe_world_cell::UnsafeWorldCell, }; -use super::{QueryData, QueryFilter, QueryItem, QueryState, ReadOnlyQueryData}; +use super::{ConsumableQueryState, QueryData, QueryItem, ReadOnlyQueryData}; use alloc::vec::Vec; @@ -13,15 +15,16 @@ use alloc::vec::Vec; /// /// This struct is created by the [`Query::par_iter`](crate::system::Query::par_iter) and /// [`Query::par_iter_mut`](crate::system::Query::par_iter_mut) methods. -pub struct QueryParIter<'w, 's, D: QueryData, F: QueryFilter> { +pub struct QueryParIter<'w, 's, S: ConsumableQueryState<'s>> { pub(crate) world: UnsafeWorldCell<'w>, - pub(crate) state: &'s QueryState, + pub(crate) state: S, + pub(crate) marker: PhantomData<&'s ()>, pub(crate) last_run: Tick, pub(crate) this_run: Tick, pub(crate) batching_strategy: BatchingStrategy, } -impl<'w, 's, D: QueryData, F: QueryFilter> QueryParIter<'w, 's, D, F> { +impl<'w, 's, D: QueryData, S: ConsumableQueryState<'s, Data = D>> QueryParIter<'w, 's, S> { /// Changes the batching strategy used when iterating. /// /// For more information on how this affects the resultant iteration, see @@ -39,7 +42,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryParIter<'w, 's, D, F> { /// /// [`ComputeTaskPool`]: bevy_tasks::ComputeTaskPool #[inline] - pub fn for_each) + Send + Sync + Clone>(self, func: FN) { + pub fn for_each) + Send + Sync + Clone>(self, func: FN) { self.for_each_init(|| {}, |_, item| func(item)); } @@ -76,7 +79,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryParIter<'w, 's, D, F> { #[inline] pub fn for_each_init(self, init: INIT, func: FN) where - FN: Fn(&mut T, QueryItem<'w, 's, D>) + Send + Sync + Clone, + FN: Fn(&mut T, QueryItem<'w, '_, D>) + Send + Sync + Clone, INIT: Fn() -> T + Sync + Send + Clone, { let func = |mut init, item| { @@ -132,7 +135,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryParIter<'w, 's, D, F> { #[cfg(all(not(target_arch = "wasm32"), feature = "multi_threaded"))] fn get_batch_size(&self, thread_count: usize) -> u32 { let max_items = || { - let id_iter = self.state.matched_storage_ids.iter(); + let id_iter = self.state.storage_ids(); if self.state.is_dense { // SAFETY: We only access table metadata. let tables = unsafe { &self.world.world_metadata().storages().tables }; @@ -161,17 +164,23 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryParIter<'w, 's, D, F> { /// /// [`Entity`]: crate::entity::Entity /// [`Query::par_iter_many`]: crate::system::Query::par_iter_many -pub struct QueryParManyIter<'w, 's, D: QueryData, F: QueryFilter, E: EntityEquivalent> { +pub struct QueryParManyIter<'w, 's, S: ConsumableQueryState<'s>, E: EntityEquivalent> { pub(crate) world: UnsafeWorldCell<'w>, - pub(crate) state: &'s QueryState, + pub(crate) state: S, + pub(crate) marker: PhantomData<&'s ()>, pub(crate) entity_list: Vec, pub(crate) last_run: Tick, pub(crate) this_run: Tick, pub(crate) batching_strategy: BatchingStrategy, } -impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, E: EntityEquivalent + Sync> - QueryParManyIter<'w, 's, D, F, E> +impl< + 'w, + 's, + D: ReadOnlyQueryData, + S: ConsumableQueryState<'s, Data = D>, + E: EntityEquivalent + Sync, + > QueryParManyIter<'w, 's, S, E> { /// Changes the batching strategy used when iterating. /// @@ -190,7 +199,7 @@ impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, E: EntityEquivalent + Sync> /// /// [`ComputeTaskPool`]: bevy_tasks::ComputeTaskPool #[inline] - pub fn for_each) + Send + Sync + Clone>(self, func: FN) { + pub fn for_each) + Send + Sync + Clone>(self, func: FN) { self.for_each_init(|| {}, |_, item| func(item)); } @@ -247,7 +256,7 @@ impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, E: EntityEquivalent + Sync> #[inline] pub fn for_each_init(self, init: INIT, func: FN) where - FN: Fn(&mut T, QueryItem<'w, 's, D>) + Send + Sync + Clone, + FN: Fn(&mut T, QueryItem<'w, '_, D>) + Send + Sync + Clone, INIT: Fn() -> T + Sync + Send + Clone, { let func = |mut init, item| { @@ -315,18 +324,18 @@ impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, E: EntityEquivalent + Sync> /// [`EntitySet`]: crate::entity::EntitySet /// [`Query::par_iter_many_unique`]: crate::system::Query::par_iter_many_unique /// [`Query::par_iter_many_unique_mut`]: crate::system::Query::par_iter_many_unique_mut -pub struct QueryParManyUniqueIter<'w, 's, D: QueryData, F: QueryFilter, E: EntityEquivalent + Sync> -{ +pub struct QueryParManyUniqueIter<'w, 's, S: ConsumableQueryState<'s>, E: EntityEquivalent + Sync> { pub(crate) world: UnsafeWorldCell<'w>, - pub(crate) state: &'s QueryState, + pub(crate) state: S, + pub(crate) marker: PhantomData<&'s ()>, pub(crate) entity_list: UniqueEntityEquivalentVec, pub(crate) last_run: Tick, pub(crate) this_run: Tick, pub(crate) batching_strategy: BatchingStrategy, } -impl<'w, 's, D: QueryData, F: QueryFilter, E: EntityEquivalent + Sync> - QueryParManyUniqueIter<'w, 's, D, F, E> +impl<'w, 's, D: QueryData, S: ConsumableQueryState<'s, Data = D>, E: EntityEquivalent + Sync> + QueryParManyUniqueIter<'w, 's, S, E> { /// Changes the batching strategy used when iterating. /// @@ -345,7 +354,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, E: EntityEquivalent + Sync> /// /// [`ComputeTaskPool`]: bevy_tasks::ComputeTaskPool #[inline] - pub fn for_each) + Send + Sync + Clone>(self, func: FN) { + pub fn for_each) + Send + Sync + Clone>(self, func: FN) { self.for_each_init(|| {}, |_, item| func(item)); } @@ -402,7 +411,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, E: EntityEquivalent + Sync> #[inline] pub fn for_each_init(self, init: INIT, func: FN) where - FN: Fn(&mut T, QueryItem<'w, 's, D>) + Send + Sync + Clone, + FN: Fn(&mut T, QueryItem<'w, '_, D>) + Send + Sync + Clone, INIT: Fn() -> T + Sync + Send + Clone, { let func = |mut init, item| { diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index 63ae1134a5745..221a3cfa77c45 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -5,7 +5,7 @@ use crate::{ entity_disabling::DefaultQueryFilters, prelude::FromWorld, query::{FilteredAccess, QueryCombinationIter, QueryIter, QueryParIter, WorldQuery}, - storage::{SparseSetIndex, TableId}, + storage::{SparseSetIndex, TableId, TableRow}, system::Query, world::{unsafe_world_cell::UnsafeWorldCell, World, WorldId}, }; @@ -13,9 +13,12 @@ use crate::{ #[cfg(all(not(target_arch = "wasm32"), feature = "multi_threaded"))] use crate::entity::UniqueEntityEquivalentSlice; -use alloc::vec::Vec; +use alloc::{ + boxed::Box, + vec::{self, Vec}, +}; use bevy_utils::prelude::DebugName; -use core::{fmt, ptr}; +use core::{fmt, iter, ops::Deref, ptr, slice}; use fixedbitset::FixedBitSet; use log::warn; #[cfg(feature = "trace")] @@ -23,13 +26,13 @@ use tracing::Span; use super::{ NopWorldQuery, QueryBuilder, QueryData, QueryEntityError, QueryFilter, QueryManyIter, - QueryManyUniqueIter, QuerySingleError, ROQueryItem, ReadOnlyQueryData, + QueryManyUniqueIter, QuerySingleError, ROQueryItem, ReadOnlyQueryData, ReleaseStateQueryData, }; /// An ID for either a table or an archetype. Used for Query iteration. /// /// Query iteration is exclusively dense (over tables) or archetypal (over archetypes) based on whether -/// the query filters are dense or not. This is represented by the [`QueryState::is_dense`] field. +/// the query filters are dense or not. This is represented by the [`QueryState::is_dense`](QueryState) field. /// /// Note that `D::IS_DENSE` and `F::IS_DENSE` have no relationship with `QueryState::is_dense` and /// any combination of their values can happen. @@ -43,7 +46,7 @@ use super::{ /// Must be initialized and accessed as a [`TableId`], if both generic parameters to the query are dense. /// Must be initialized and accessed as an [`ArchetypeId`] otherwise. #[derive(Clone, Copy)] -pub(super) union StorageId { +pub union StorageId { pub(super) table_id: TableId, pub(super) archetype_id: ArchetypeId, } @@ -88,6 +91,126 @@ pub struct QueryState { par_iter_span: Span, } +/// An owned or borrowed [`QueryState`] that can be consumed to query. +/// +/// A borrowed `&QueryState` can always be consumed, +/// but an owned `Box` cannot be used +/// if the query data needs to borrow from the state. +pub trait ConsumableQueryState<'state>: + Deref> +{ + /// The [`QueryData`] for this `QueryState`. + type Data: QueryData; + + /// The [`QueryFilter`] for this `QueryState`. + type Filter: QueryFilter; + + /// The type returned by [`Self::storage_ids`]. + type StorageIter: Iterator + Clone + Default; + + /// A read-only version of the state. + type ReadOnly: ConsumableQueryState< + 'state, + Data = ::ReadOnly, + Filter = Self::Filter, + >; + + /// Iterates the storage ids that this [`QueryState`] matches. + fn storage_ids(&self) -> Self::StorageIter; + + /// Borrows the remainder of the [`Self::StorageIter`] as an iterator + /// usable with `&QueryState`. + fn reborrow_storage_ids( + storage_iter: &Self::StorageIter, + ) -> iter::Copied>; + + /// Converts this state to a read-only version. + fn into_readonly(self) -> Self::ReadOnly; + + /// Calls [`QueryData::fetch`] with the provided parameters, + /// adjusting the `'state` lifetime as necessary. + /// + /// # Safety + /// + /// It must be safe to call [`QueryData::fetch`], meaning: + /// - Must always be called _after_ [`WorldQuery::set_table`] or [`WorldQuery::set_archetype`]. `entity` and + /// `table_row` must be in the range of the current table and archetype. + /// - There must not be simultaneous conflicting component access registered in `update_component_access`. + unsafe fn fetch<'w>( + &self, + fetch: &mut ::Fetch<'w>, + entity: Entity, + table_row: TableRow, + ) -> ::Item<'w, 'state>; +} + +impl, F: QueryFilter> + ConsumableQueryState<'static> for Box> +{ + type Data = D; + type Filter = F; + type StorageIter = vec::IntoIter; + type ReadOnly = Box>; + + fn storage_ids(&self) -> Self::StorageIter { + // Query iteration holds both the state and the storage iterator, + // so the iterator cannot borrow from the state, + // which requires us to `clone` the `Vec` here. + self.matched_storage_ids.clone().into_iter() + } + + fn reborrow_storage_ids( + storage_iter: &Self::StorageIter, + ) -> iter::Copied> { + storage_iter.as_slice().iter().copied() + } + + fn into_readonly(self) -> Self::ReadOnly { + self.into_readonly() + } + + unsafe fn fetch<'w>( + &self, + fetch: &mut ::Fetch<'w>, + entity: Entity, + table_row: TableRow, + ) -> ::Item<'w, 'static> { + // SAFETY: Caller ensures it is safe to call `QueryData::Fetch` + D::release_state(unsafe { D::fetch(&self.fetch_state, fetch, entity, table_row) }) + } +} + +impl<'s, D: QueryData, F: QueryFilter> ConsumableQueryState<'s> for &'s QueryState { + type Data = D; + type Filter = F; + type StorageIter = iter::Copied>; + type ReadOnly = &'s QueryState; + + fn storage_ids(&self) -> Self::StorageIter { + self.matched_storage_ids.iter().copied() + } + + fn reborrow_storage_ids( + storage_iter: &Self::StorageIter, + ) -> iter::Copied> { + storage_iter.clone() + } + + fn into_readonly(self) -> Self::ReadOnly { + self.as_readonly() + } + + unsafe fn fetch<'w>( + &self, + fetch: &mut ::Fetch<'w>, + entity: Entity, + table_row: TableRow, + ) -> ::Item<'w, 's> { + // SAFETY: Caller ensures it is safe to call `QueryData::Fetch` + unsafe { D::fetch(&self.fetch_state, fetch, entity, table_row) } + } +} + impl fmt::Debug for QueryState { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("QueryState") @@ -144,6 +267,31 @@ impl QueryState { &*ptr::from_ref(self).cast::>() } + /// Converts this `QueryState` reference to a `QueryState` that does not access anything mutably. + pub fn into_readonly(self: Box) -> Box> { + // SAFETY: invariant on `WorldQuery` trait upholds that `D::ReadOnly` and `F::ReadOnly` + // have a subset of the access, and match the exact same archetypes/tables as `D`/`F` respectively. + unsafe { self.into_transmuted_state::() } + } + + /// Converts this `QueryState` reference to any other `QueryState` with + /// the same `WorldQuery::State` associated types. + /// + /// Consider using `into_readonly` instead which is a safe function. + /// + /// # Safety + /// + /// `NewD` must have a subset of the access that `D` does and match the exact same archetypes/tables + /// `NewF` must have a subset of the access that `F` does and match the exact same archetypes/tables + pub(crate) unsafe fn into_transmuted_state< + NewD: ReadOnlyQueryData, + NewF: QueryFilter, + >( + self: Box, + ) -> Box> { + Box::from_raw(Box::into_raw(self).cast::>()) + } + /// Returns the components accessed by this query. pub fn component_access(&self) -> &FilteredAccess { &self.component_access @@ -1089,7 +1237,10 @@ impl QueryState { /// If you need to iterate multiple times at once but get borrowing errors, /// consider using [`Self::update_archetypes`] followed by multiple [`Self::iter_manual`] calls. #[inline] - pub fn iter<'w, 's>(&'s mut self, world: &'w World) -> QueryIter<'w, 's, D::ReadOnly, F> { + pub fn iter<'w, 's>( + &'s mut self, + world: &'w World, + ) -> QueryIter<'w, 's, &'s QueryState> { self.query(world).into_iter() } @@ -1098,7 +1249,7 @@ impl QueryState { /// This iterator is always guaranteed to return results from each matching entity once and only once. /// Iteration order is not guaranteed. #[inline] - pub fn iter_mut<'w, 's>(&'s mut self, world: &'w mut World) -> QueryIter<'w, 's, D, F> { + pub fn iter_mut<'w, 's>(&'s mut self, world: &'w mut World) -> QueryIter<'w, 's, &'s Self> { self.query_mut(world).into_iter() } @@ -1110,7 +1261,10 @@ impl QueryState { /// /// This can only be called for read-only queries. #[inline] - pub fn iter_manual<'w, 's>(&'s self, world: &'w World) -> QueryIter<'w, 's, D::ReadOnly, F> { + pub fn iter_manual<'w, 's>( + &'s self, + world: &'w World, + ) -> QueryIter<'w, 's, &'s QueryState> { self.query_manual(world).into_iter() } @@ -1142,7 +1296,7 @@ impl QueryState { pub fn iter_combinations<'w, 's, const K: usize>( &'s mut self, world: &'w World, - ) -> QueryCombinationIter<'w, 's, D::ReadOnly, F, K> { + ) -> QueryCombinationIter<'w, 's, &'s QueryState, K> { self.query(world).iter_combinations_inner() } @@ -1167,7 +1321,7 @@ impl QueryState { pub fn iter_combinations_mut<'w, 's, const K: usize>( &'s mut self, world: &'w mut World, - ) -> QueryCombinationIter<'w, 's, D, F, K> { + ) -> QueryCombinationIter<'w, 's, &'s Self, K> { self.query_mut(world).iter_combinations_inner() } @@ -1187,7 +1341,7 @@ impl QueryState { &'s mut self, world: &'w World, entities: EntityList, - ) -> QueryManyIter<'w, 's, D::ReadOnly, F, EntityList::IntoIter> { + ) -> QueryManyIter<'w, 's, &'s QueryState, EntityList::IntoIter> { self.query(world).iter_many_inner(entities) } @@ -1210,7 +1364,7 @@ impl QueryState { &'s self, world: &'w World, entities: EntityList, - ) -> QueryManyIter<'w, 's, D::ReadOnly, F, EntityList::IntoIter> { + ) -> QueryManyIter<'w, 's, &'s QueryState, EntityList::IntoIter> { self.query_manual(world).iter_many_inner(entities) } @@ -1223,7 +1377,7 @@ impl QueryState { &'s mut self, world: &'w mut World, entities: EntityList, - ) -> QueryManyIter<'w, 's, D, F, EntityList::IntoIter> { + ) -> QueryManyIter<'w, 's, &'s Self, EntityList::IntoIter> { self.query_mut(world).iter_many_inner(entities) } @@ -1240,7 +1394,7 @@ impl QueryState { &'s mut self, world: &'w World, entities: EntityList, - ) -> QueryManyUniqueIter<'w, 's, D::ReadOnly, F, EntityList::IntoIter> { + ) -> QueryManyUniqueIter<'w, 's, &'s QueryState, EntityList::IntoIter> { self.query(world).iter_many_unique_inner(entities) } @@ -1264,7 +1418,7 @@ impl QueryState { &'s self, world: &'w World, entities: EntityList, - ) -> QueryManyUniqueIter<'w, 's, D::ReadOnly, F, EntityList::IntoIter> { + ) -> QueryManyUniqueIter<'w, 's, &'s QueryState, EntityList::IntoIter> { self.query_manual(world).iter_many_unique_inner(entities) } @@ -1277,7 +1431,7 @@ impl QueryState { &'s mut self, world: &'w mut World, entities: EntityList, - ) -> QueryManyUniqueIter<'w, 's, D, F, EntityList::IntoIter> { + ) -> QueryManyUniqueIter<'w, 's, &'s Self, EntityList::IntoIter> { self.query_mut(world).iter_many_unique_inner(entities) } /// Returns an [`Iterator`] over the query results for the given [`World`]. @@ -1293,7 +1447,7 @@ impl QueryState { pub unsafe fn iter_unchecked<'w, 's>( &'s mut self, world: UnsafeWorldCell<'w>, - ) -> QueryIter<'w, 's, D, F> { + ) -> QueryIter<'w, 's, &'s Self> { self.query_unchecked(world).into_iter() } @@ -1312,7 +1466,7 @@ impl QueryState { pub unsafe fn iter_combinations_unchecked<'w, 's, const K: usize>( &'s mut self, world: UnsafeWorldCell<'w>, - ) -> QueryCombinationIter<'w, 's, D, F, K> { + ) -> QueryCombinationIter<'w, 's, &'s Self, K> { self.query_unchecked(world).iter_combinations_inner() } @@ -1328,7 +1482,7 @@ impl QueryState { pub fn par_iter<'w, 's>( &'s mut self, world: &'w World, - ) -> QueryParIter<'w, 's, D::ReadOnly, F> { + ) -> QueryParIter<'w, 's, &'s QueryState> { self.query(world).par_iter_inner() } @@ -1377,7 +1531,10 @@ impl QueryState { /// [`par_iter`]: Self::par_iter /// [`ComputeTaskPool`]: bevy_tasks::ComputeTaskPool #[inline] - pub fn par_iter_mut<'w, 's>(&'s mut self, world: &'w mut World) -> QueryParIter<'w, 's, D, F> { + pub fn par_iter_mut<'w, 's>( + &'s mut self, + world: &'w mut World, + ) -> QueryParIter<'w, 's, &'s Self> { self.query_mut(world).par_iter_inner() } @@ -2032,7 +2189,7 @@ mod tests { .run_system_once(|query: Query<&mut A>| { let mut readonly = query.as_readonly(); let mut lens: QueryLens<&mut A> = readonly.transmute_lens(); - bad(lens.query(), query.as_readonly()); + bad(lens.reborrow(), query.as_readonly()); }) .unwrap(); } @@ -2167,7 +2324,7 @@ mod tests { .run_system_once(|query_a: Query<&mut A>, mut query_b: Query<&mut B>| { let mut readonly = query_a.as_readonly(); let mut lens: QueryLens<(&mut A, &mut B)> = readonly.join(&mut query_b); - bad(lens.query(), query_a.as_readonly()); + bad(lens.reborrow(), query_a.as_readonly()); }) .unwrap(); } diff --git a/crates/bevy_ecs/src/relationship/relationship_query.rs b/crates/bevy_ecs/src/relationship/relationship_query.rs index a7acea7de0732..d6827b4c94688 100644 --- a/crates/bevy_ecs/src/relationship/relationship_query.rs +++ b/crates/bevy_ecs/src/relationship/relationship_query.rs @@ -1,6 +1,8 @@ +use core::ops::Deref; + use crate::{ entity::Entity, - query::{QueryData, QueryFilter}, + query::{QueryData, QueryFilter, QueryState}, relationship::{Relationship, RelationshipTarget}, system::Query, }; @@ -9,24 +11,26 @@ use smallvec::SmallVec; use super::SourceIter; -impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { +impl<'w, 's, D: QueryData, F: QueryFilter, S: Deref>> + Query<'w, 's, D, F, S> +{ /// If the given `entity` contains the `R` [`Relationship`] component, returns the /// target entity of that relationship. pub fn related(&'w self, entity: Entity) -> Option where - ::ReadOnly: QueryData = &'w R>, + D::ReadOnly: QueryData = &'w R>, { self.get(entity).map(R::get).ok() } /// If the given `entity` contains the `S` [`RelationshipTarget`] component, returns the /// source entities stored on that component. - pub fn relationship_sources( + pub fn relationship_sources( &'w self, entity: Entity, ) -> impl Iterator + 'w where - ::ReadOnly: QueryData = &'w S>, + D::ReadOnly: QueryData = &'w R>, { self.get(entity) .into_iter() @@ -42,7 +46,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// If your relationship is not a tree (like Bevy's hierarchy), be sure to stop if you encounter a duplicate entity. pub fn root_ancestor(&'w self, entity: Entity) -> Entity where - ::ReadOnly: QueryData = &'w R>, + D::ReadOnly: QueryData = &'w R>, { // Recursively search up the tree until we're out of parents match self.get(entity) { @@ -57,13 +61,13 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// For relationship graphs that contain loops, this could loop infinitely. /// If your relationship is not a tree (like Bevy's hierarchy), be sure to stop if you encounter a duplicate entity. - pub fn iter_leaves( + pub fn iter_leaves( &'w self, entity: Entity, - ) -> impl Iterator + use<'w, 's, S, D, F> + ) -> impl Iterator + use<'w, 's, D, F, S, R> where - ::ReadOnly: QueryData = &'w S>, - SourceIter<'w, S>: DoubleEndedIterator, + D::ReadOnly: QueryData = &'w R>, + SourceIter<'w, R>: DoubleEndedIterator, { self.iter_descendants_depth_first(entity).filter(|entity| { self.get(*entity) @@ -80,7 +84,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { entity: Entity, ) -> impl Iterator + 'w where - D::ReadOnly: QueryData = (Option<&'w R>, Option<&'w R::RelationshipTarget>)>, + D::ReadOnly: QueryData = (Option<&'w R>, Option<&'w R::RelationshipTarget>)>, { self.get(entity) .ok() @@ -98,12 +102,12 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// For relationship graphs that contain loops, this could loop infinitely. /// If your relationship is not a tree (like Bevy's hierarchy), be sure to stop if you encounter a duplicate entity. - pub fn iter_descendants( + pub fn iter_descendants( &'w self, entity: Entity, - ) -> DescendantIter<'w, 's, D, F, S> + ) -> DescendantIter<'w, 's, D, F, S, R> where - D::ReadOnly: QueryData = &'w S>, + D::ReadOnly: QueryData = &'w R>, { DescendantIter::new(self, entity) } @@ -115,13 +119,13 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// For relationship graphs that contain loops, this could loop infinitely. /// If your relationship is not a tree (like Bevy's hierarchy), be sure to stop if you encounter a duplicate entity. - pub fn iter_descendants_depth_first( + pub fn iter_descendants_depth_first( &'w self, entity: Entity, - ) -> DescendantDepthFirstIter<'w, 's, D, F, S> + ) -> DescendantDepthFirstIter<'w, 's, D, F, S, R> where - D::ReadOnly: QueryData = &'w S>, - SourceIter<'w, S>: DoubleEndedIterator, + D::ReadOnly: QueryData = &'w R>, + SourceIter<'w, R>: DoubleEndedIterator, { DescendantDepthFirstIter::new(self, entity) } @@ -135,9 +139,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn iter_ancestors( &'w self, entity: Entity, - ) -> AncestorIter<'w, 's, D, F, R> + ) -> AncestorIter<'w, 's, D, F, S, R> where - D::ReadOnly: QueryData = &'w R>, + D::ReadOnly: QueryData = &'w R>, { AncestorIter::new(self, entity) } @@ -146,20 +150,33 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// An [`Iterator`] of [`Entity`]s over the descendants of an [`Entity`]. /// /// Traverses the hierarchy breadth-first. -pub struct DescendantIter<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget> -where - D::ReadOnly: QueryData = &'w S>, +pub struct DescendantIter< + 'w, + 's, + D: QueryData, + F: QueryFilter, + S: Deref>, + R: RelationshipTarget, +> where + D::ReadOnly: QueryData = &'w R>, { - children_query: &'w Query<'w, 's, D, F>, + children_query: &'w Query<'w, 's, D, F, S>, vecdeque: VecDeque, } -impl<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget> DescendantIter<'w, 's, D, F, S> +impl< + 'w, + 's, + D: QueryData, + F: QueryFilter, + S: Deref>, + R: RelationshipTarget, + > DescendantIter<'w, 's, D, F, S, R> where - D::ReadOnly: QueryData = &'w S>, + D::ReadOnly: QueryData = &'w R>, { /// Returns a new [`DescendantIter`]. - pub fn new(children_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self { + pub fn new(children_query: &'w Query<'w, 's, D, F, S>, entity: Entity) -> Self { DescendantIter { children_query, vecdeque: children_query @@ -171,10 +188,16 @@ where } } -impl<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget> Iterator - for DescendantIter<'w, 's, D, F, S> +impl< + 'w, + 's, + D: QueryData, + F: QueryFilter, + S: Deref>, + R: RelationshipTarget, + > Iterator for DescendantIter<'w, 's, D, F, S, R> where - D::ReadOnly: QueryData = &'w S>, + D::ReadOnly: QueryData = &'w R>, { type Item = Entity; @@ -192,22 +215,34 @@ where /// An [`Iterator`] of [`Entity`]s over the descendants of an [`Entity`]. /// /// Traverses the hierarchy depth-first. -pub struct DescendantDepthFirstIter<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget> -where - D::ReadOnly: QueryData = &'w S>, +pub struct DescendantDepthFirstIter< + 'w, + 's, + D: QueryData, + F: QueryFilter, + S: Deref>, + R: RelationshipTarget, +> where + D::ReadOnly: QueryData = &'w R>, { - children_query: &'w Query<'w, 's, D, F>, + children_query: &'w Query<'w, 's, D, F, S>, stack: SmallVec<[Entity; 8]>, } -impl<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget> - DescendantDepthFirstIter<'w, 's, D, F, S> +impl< + 'w, + 's, + D: QueryData, + F: QueryFilter, + S: Deref>, + R: RelationshipTarget, + > DescendantDepthFirstIter<'w, 's, D, F, S, R> where - D::ReadOnly: QueryData = &'w S>, - SourceIter<'w, S>: DoubleEndedIterator, + D::ReadOnly: QueryData = &'w R>, + SourceIter<'w, R>: DoubleEndedIterator, { /// Returns a new [`DescendantDepthFirstIter`]. - pub fn new(children_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self { + pub fn new(children_query: &'w Query<'w, 's, D, F, S>, entity: Entity) -> Self { DescendantDepthFirstIter { children_query, stack: children_query @@ -217,11 +252,17 @@ where } } -impl<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget> Iterator - for DescendantDepthFirstIter<'w, 's, D, F, S> +impl< + 'w, + 's, + D: QueryData, + F: QueryFilter, + S: Deref>, + R: RelationshipTarget, + > Iterator for DescendantDepthFirstIter<'w, 's, D, F, S, R> where - D::ReadOnly: QueryData = &'w S>, - SourceIter<'w, S>: DoubleEndedIterator, + D::ReadOnly: QueryData = &'w R>, + SourceIter<'w, R>: DoubleEndedIterator, { type Item = Entity; @@ -237,20 +278,33 @@ where } /// An [`Iterator`] of [`Entity`]s over the ancestors of an [`Entity`]. -pub struct AncestorIter<'w, 's, D: QueryData, F: QueryFilter, R: Relationship> -where - D::ReadOnly: QueryData = &'w R>, +pub struct AncestorIter< + 'w, + 's, + D: QueryData, + F: QueryFilter, + S: Deref>, + R: Relationship, +> where + D::ReadOnly: QueryData = &'w R>, { - parent_query: &'w Query<'w, 's, D, F>, + parent_query: &'w Query<'w, 's, D, F, S>, next: Option, } -impl<'w, 's, D: QueryData, F: QueryFilter, R: Relationship> AncestorIter<'w, 's, D, F, R> +impl< + 'w, + 's, + D: QueryData, + F: QueryFilter, + S: Deref>, + R: Relationship, + > AncestorIter<'w, 's, D, F, S, R> where - D::ReadOnly: QueryData = &'w R>, + D::ReadOnly: QueryData = &'w R>, { /// Returns a new [`AncestorIter`]. - pub fn new(parent_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self { + pub fn new(parent_query: &'w Query<'w, 's, D, F, S>, entity: Entity) -> Self { AncestorIter { parent_query, next: Some(entity), @@ -258,10 +312,16 @@ where } } -impl<'w, 's, D: QueryData, F: QueryFilter, R: Relationship> Iterator - for AncestorIter<'w, 's, D, F, R> +impl< + 'w, + 's, + D: QueryData, + F: QueryFilter, + S: Deref>, + R: Relationship, + > Iterator for AncestorIter<'w, 's, D, F, S, R> where - D::ReadOnly: QueryData = &'w R>, + D::ReadOnly: QueryData = &'w R>, { type Item = Entity; diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index 6e44301b18cb1..96777fbc9a51a 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -5,12 +5,14 @@ use crate::{ component::Tick, entity::{Entity, EntityDoesNotExistError, EntityEquivalent, EntitySet, UniqueEntityArray}, query::{ - DebugCheckedUnwrap, NopWorldQuery, QueryCombinationIter, QueryData, QueryEntityError, - QueryFilter, QueryIter, QueryManyIter, QueryManyUniqueIter, QueryParIter, QueryParManyIter, - QueryParManyUniqueIter, QuerySingleError, QueryState, ROQueryItem, ReadOnlyQueryData, + ConsumableQueryState, DebugCheckedUnwrap, NopWorldQuery, QueryCombinationIter, QueryData, + QueryEntityError, QueryFilter, QueryIter, QueryManyIter, QueryManyUniqueIter, QueryParIter, + QueryParManyIter, QueryParManyUniqueIter, QuerySingleError, QueryState, ROQueryItem, + ReadOnlyQueryData, }, world::unsafe_world_cell::UnsafeWorldCell, }; +use bevy_platform::prelude::Box; use core::{ marker::PhantomData, mem::MaybeUninit, @@ -482,12 +484,23 @@ use core::{ /// ``` /// /// [autovectorization]: https://en.wikipedia.org/wiki/Automatic_vectorization -pub struct Query<'world, 'state, D: QueryData, F: QueryFilter = ()> { +pub struct Query< + 'world, + // Note that `'state` is only used in the default type for `S`. + // That's necessary to make `Query` be a valid `SystemParam`. + 'state, + D: QueryData, + F: QueryFilter = (), + S: Deref> = &'state QueryState, +> { // SAFETY: Must have access to the components registered in `state`. world: UnsafeWorldCell<'world>, - state: &'state QueryState, + state: S, last_run: Tick, this_run: Tick, + // Note that `&'state QueryState` won't work because + // `QueryLens` uses `'static`, but `D` and `F` might not be `'static`. + marker: PhantomData<(&'state (), fn() -> QueryState)>, } impl Clone for Query<'_, '_, D, F> { @@ -510,7 +523,9 @@ impl core::fmt::Debug for Query<'_, '_, D, F> { } } -impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { +impl<'w, 's, D: QueryData, F: QueryFilter, S: Deref>> + Query<'w, 's, D, F, S> +{ /// Creates a new query. /// /// # Safety @@ -521,7 +536,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { #[inline] pub(crate) unsafe fn new( world: UnsafeWorldCell<'w>, - state: &'s QueryState, + state: S, last_run: Tick, this_run: Tick, ) -> Self { @@ -530,6 +545,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { state, last_run, this_run, + marker: PhantomData, } } @@ -542,14 +558,14 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// # See also /// /// [`into_readonly`](Self::into_readonly) for a version that consumes the `Query` to return one with the full `'world` lifetime. - pub fn as_readonly(&self) -> Query<'_, 's, D::ReadOnly, F> { + pub fn as_readonly(&self) -> Query<'_, '_, D::ReadOnly, F> { // SAFETY: The reborrowed query is converted to read-only, so it cannot perform mutable access, // and the original query is held with a shared borrow, so it cannot perform mutable access either. unsafe { self.reborrow_unsafe() }.into_readonly() } /// Returns another `Query` from this does not return any data, which can be faster. - fn as_nop(&self) -> Query<'_, 's, NopWorldQuery, F> { + fn as_nop(&self) -> Query<'_, '_, NopWorldQuery, F> { let new_state = self.state.as_nop(); // SAFETY: // - The reborrowed query is converted to read-only, so it cannot perform mutable access, @@ -569,8 +585,11 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// # See also /// /// [`as_readonly`](Self::as_readonly) for a version that borrows the `Query` instead of consuming it. - pub fn into_readonly(self) -> Query<'w, 's, D::ReadOnly, F> { - let new_state = self.state.as_readonly(); + pub fn into_readonly(self) -> Query<'w, 's, D::ReadOnly, F, S::ReadOnly> + where + S: ConsumableQueryState<'s, Data = D, Filter = F>, + { + let new_state = self.state.into_readonly(); // SAFETY: // - This is memory safe because it turns the query immutable. // - The world matches because it was the same one used to construct self. @@ -601,7 +620,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// } /// } /// ``` - pub fn reborrow(&mut self) -> Query<'_, 's, D, F> { + pub fn reborrow(&mut self) -> Query<'_, '_, D, F> { // SAFETY: this query is exclusively borrowed while the new one exists, so // no overlapping access can occur. unsafe { self.reborrow_unsafe() } @@ -618,29 +637,11 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// # See also /// /// - [`reborrow`](Self::reborrow) for the safe versions. - pub unsafe fn reborrow_unsafe(&self) -> Query<'_, 's, D, F> { - // SAFETY: - // - This is memory safe because the caller ensures that there are no conflicting references. - // - The world matches because it was the same one used to construct self. - unsafe { self.copy_unsafe() } - } - - /// Returns a new `Query` copying the access from this one. - /// The current query will still be usable while the new one exists, but must not be used in a way that violates aliasing. - /// - /// # Safety - /// - /// This function makes it possible to violate Rust's aliasing guarantees. - /// You must make sure this call does not result in a mutable or shared reference to a component with a mutable reference. - /// - /// # See also - /// - /// - [`reborrow_unsafe`](Self::reborrow_unsafe) for a safer version that constrains the returned `'w` lifetime to the length of the borrow. - unsafe fn copy_unsafe(&self) -> Query<'w, 's, D, F> { + pub unsafe fn reborrow_unsafe(&self) -> Query<'_, '_, D, F> { // SAFETY: // - This is memory safe because the caller ensures that there are no conflicting references. // - The world matches because it was the same one used to construct self. - unsafe { Query::new(self.world, self.state, self.last_run, self.this_run) } + unsafe { Query::new(self.world, &self.state, self.last_run, self.this_run) } } /// Returns an [`Iterator`] over the read-only query items. @@ -670,7 +671,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// [`iter_mut`](Self::iter_mut) for mutable query items. #[inline] - pub fn iter(&self) -> QueryIter<'_, 's, D::ReadOnly, F> { + pub fn iter(&self) -> QueryIter<'_, '_, &QueryState> { self.as_readonly().into_iter() } @@ -701,7 +702,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// [`iter`](Self::iter) for read-only query items. #[inline] - pub fn iter_mut(&mut self) -> QueryIter<'_, 's, D, F> { + pub fn iter_mut(&mut self) -> QueryIter<'_, '_, &QueryState> { self.reborrow().into_iter() } @@ -731,7 +732,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { #[inline] pub fn iter_combinations( &self, - ) -> QueryCombinationIter<'_, 's, D::ReadOnly, F, K> { + ) -> QueryCombinationIter<'_, '_, &QueryState, K> { self.as_readonly().iter_combinations_inner() } @@ -761,7 +762,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { #[inline] pub fn iter_combinations_mut( &mut self, - ) -> QueryCombinationIter<'_, 's, D, F, K> { + ) -> QueryCombinationIter<'_, '_, &QueryState, K> { self.reborrow().iter_combinations_inner() } @@ -790,7 +791,10 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// - [`iter_combinations`](Self::iter_combinations) for read-only query item combinations. /// - [`iter_combinations_mut`](Self::iter_combinations_mut) for mutable query item combinations. #[inline] - pub fn iter_combinations_inner(self) -> QueryCombinationIter<'w, 's, D, F, K> { + pub fn iter_combinations_inner(self) -> QueryCombinationIter<'w, 's, S, K> + where + S: ConsumableQueryState<'s, Data = D, Filter = F>, + { // SAFETY: `self.world` has permission to access the required components. unsafe { QueryCombinationIter::new(self.world, self.state, self.last_run, self.this_run) } } @@ -836,7 +840,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn iter_many>( &self, entities: EntityList, - ) -> QueryManyIter<'_, 's, D::ReadOnly, F, EntityList::IntoIter> { + ) -> QueryManyIter<'_, '_, &QueryState, EntityList::IntoIter> { self.as_readonly().iter_many_inner(entities) } @@ -881,7 +885,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn iter_many_mut>( &mut self, entities: EntityList, - ) -> QueryManyIter<'_, 's, D, F, EntityList::IntoIter> { + ) -> QueryManyIter<'_, '_, &QueryState, EntityList::IntoIter> { self.reborrow().iter_many_inner(entities) } @@ -899,7 +903,10 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn iter_many_inner>( self, entities: EntityList, - ) -> QueryManyIter<'w, 's, D, F, EntityList::IntoIter> { + ) -> QueryManyIter<'w, 's, S, EntityList::IntoIter> + where + S: ConsumableQueryState<'s, Data = D, Filter = F>, + { // SAFETY: `self.world` has permission to access the required components. unsafe { QueryManyIter::new( @@ -964,7 +971,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn iter_many_unique( &self, entities: EntityList, - ) -> QueryManyUniqueIter<'_, 's, D::ReadOnly, F, EntityList::IntoIter> { + ) -> QueryManyUniqueIter<'_, '_, &QueryState, EntityList::IntoIter> { self.as_readonly().iter_many_unique_inner(entities) } @@ -1019,7 +1026,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn iter_many_unique_mut( &mut self, entities: EntityList, - ) -> QueryManyUniqueIter<'_, 's, D, F, EntityList::IntoIter> { + ) -> QueryManyUniqueIter<'_, '_, &QueryState, EntityList::IntoIter> { self.reborrow().iter_many_unique_inner(entities) } @@ -1074,7 +1081,10 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn iter_many_unique_inner( self, entities: EntityList, - ) -> QueryManyUniqueIter<'w, 's, D, F, EntityList::IntoIter> { + ) -> QueryManyUniqueIter<'w, 's, S, EntityList::IntoIter> + where + S: ConsumableQueryState<'s, Data = D, Filter = F>, + { // SAFETY: `self.world` has permission to access the required components. unsafe { QueryManyUniqueIter::new( @@ -1101,7 +1111,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// - [`iter`](Self::iter) and [`iter_mut`](Self::iter_mut) for the safe versions. #[inline] - pub unsafe fn iter_unsafe(&self) -> QueryIter<'_, 's, D, F> { + pub unsafe fn iter_unsafe(&self) -> QueryIter<'_, '_, &QueryState> { // SAFETY: The caller promises that this will not result in multiple mutable references. unsafe { self.reborrow_unsafe() }.into_iter() } @@ -1122,7 +1132,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { #[inline] pub unsafe fn iter_combinations_unsafe( &self, - ) -> QueryCombinationIter<'_, 's, D, F, K> { + ) -> QueryCombinationIter<'_, '_, &QueryState, K> { // SAFETY: The caller promises that this will not result in multiple mutable references. unsafe { self.reborrow_unsafe() }.iter_combinations_inner() } @@ -1144,7 +1154,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub unsafe fn iter_many_unsafe>( &self, entities: EntityList, - ) -> QueryManyIter<'_, 's, D, F, EntityList::IntoIter> { + ) -> QueryManyIter<'_, '_, &QueryState, EntityList::IntoIter> { // SAFETY: The caller promises that this will not result in multiple mutable references. unsafe { self.reborrow_unsafe() }.iter_many_inner(entities) } @@ -1166,7 +1176,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub unsafe fn iter_many_unique_unsafe( &self, entities: EntityList, - ) -> QueryManyUniqueIter<'_, 's, D, F, EntityList::IntoIter> { + ) -> QueryManyUniqueIter<'_, '_, &QueryState, EntityList::IntoIter> { // SAFETY: The caller promises that this will not result in multiple mutable references. unsafe { self.reborrow_unsafe() }.iter_many_unique_inner(entities) } @@ -1187,7 +1197,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// [`par_iter_mut`]: Self::par_iter_mut /// [`World`]: crate::world::World #[inline] - pub fn par_iter(&self) -> QueryParIter<'_, 's, D::ReadOnly, F> { + pub fn par_iter(&self) -> QueryParIter<'_, '_, &QueryState> { self.as_readonly().par_iter_inner() } @@ -1222,7 +1232,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// [`par_iter`]: Self::par_iter /// [`World`]: crate::world::World #[inline] - pub fn par_iter_mut(&mut self) -> QueryParIter<'_, 's, D, F> { + pub fn par_iter_mut(&mut self) -> QueryParIter<'_, '_, &QueryState> { self.reborrow().par_iter_inner() } @@ -1253,10 +1263,14 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// # bevy_ecs::system::assert_is_system(gravity_system); /// ``` #[inline] - pub fn par_iter_inner(self) -> QueryParIter<'w, 's, D, F> { + pub fn par_iter_inner(self) -> QueryParIter<'w, 's, S> + where + S: ConsumableQueryState<'s, Data = D, Filter = F>, + { QueryParIter { world: self.world, state: self.state, + marker: PhantomData, last_run: self.last_run, this_run: self.this_run, batching_strategy: BatchingStrategy::new(), @@ -1282,10 +1296,11 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn par_iter_many>( &self, entities: EntityList, - ) -> QueryParManyIter<'_, 's, D::ReadOnly, F, EntityList::Item> { + ) -> QueryParManyIter<'_, '_, &QueryState, EntityList::Item> { QueryParManyIter { world: self.world, state: self.state.as_readonly(), + marker: PhantomData, entity_list: entities.into_iter().collect(), last_run: self.last_run, this_run: self.this_run, @@ -1311,10 +1326,11 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn par_iter_many_unique>( &self, entities: EntityList, - ) -> QueryParManyUniqueIter<'_, 's, D::ReadOnly, F, EntityList::Item> { + ) -> QueryParManyUniqueIter<'_, '_, &QueryState, EntityList::Item> { QueryParManyUniqueIter { world: self.world, state: self.state.as_readonly(), + marker: PhantomData, entity_list: entities.into_iter().collect(), last_run: self.last_run, this_run: self.this_run, @@ -1340,10 +1356,11 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn par_iter_many_unique_mut>( &mut self, entities: EntityList, - ) -> QueryParManyUniqueIter<'_, 's, D, F, EntityList::Item> { + ) -> QueryParManyUniqueIter<'_, '_, &QueryState, EntityList::Item> { QueryParManyUniqueIter { world: self.world, - state: self.state, + state: &self.state, + marker: PhantomData, entity_list: entities.into_iter().collect(), last_run: self.last_run, this_run: self.this_run, @@ -1385,7 +1402,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// - [`get_mut`](Self::get_mut) to get a mutable query item. #[inline] - pub fn get(&self, entity: Entity) -> Result, QueryEntityError> { + pub fn get(&self, entity: Entity) -> Result, QueryEntityError> { self.as_readonly().get_inner(entity) } @@ -1436,7 +1453,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn get_many( &self, entities: [Entity; N], - ) -> Result<[ROQueryItem<'_, 's, D>; N], QueryEntityError> { + ) -> Result<[ROQueryItem<'_, '_, D>; N], QueryEntityError> { // Note that we call a separate `*_inner` method from `get_many_mut` // because we don't need to check for duplicates. self.as_readonly().get_many_inner(entities) @@ -1487,7 +1504,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn get_many_unique( &self, entities: UniqueEntityArray, - ) -> Result<[ROQueryItem<'_, 's, D>; N], QueryEntityError> { + ) -> Result<[ROQueryItem<'_, '_, D>; N], QueryEntityError> { self.as_readonly().get_many_unique_inner(entities) } @@ -1521,7 +1538,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// - [`get`](Self::get) to get a read-only query item. #[inline] - pub fn get_mut(&mut self, entity: Entity) -> Result, QueryEntityError> { + pub fn get_mut(&mut self, entity: Entity) -> Result, QueryEntityError> { self.reborrow().get_inner(entity) } @@ -1536,7 +1553,31 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// - [`get_mut`](Self::get_mut) to get the item using a mutable borrow of the [`Query`]. #[inline] - pub fn get_inner(self, entity: Entity) -> Result, QueryEntityError> { + pub fn get_inner(self, entity: Entity) -> Result, QueryEntityError> + where + S: ConsumableQueryState<'s, Data = D, Filter = F>, + { + // SAFETY: This query has access to this item, + // and we consume the query so it is never used again. + unsafe { self.get_unchecked_inner(entity) } + } + + /// Returns the query item for the given [`Entity`]. + /// + /// This is similar to [`Self::get_unchecked`], but returns results with the actual "inner" world lifetime. + /// + /// # Safety + /// + /// This function makes it possible to violate Rust's aliasing guarantees. + /// You must make sure this call does not result in multiple mutable references to the same component. + #[inline] + unsafe fn get_unchecked_inner( + &self, + entity: Entity, + ) -> Result, QueryEntityError> + where + S: ConsumableQueryState<'s, Data = D, Filter = F>, + { // SAFETY: system runs without conflicts with other systems. // same-system queries have runtime borrow checks when they conflict unsafe { @@ -1588,12 +1629,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { entity, location.table_row, ) { - Ok(D::fetch( - &self.state.fetch_state, - &mut fetch, - entity, - location.table_row, - )) + Ok(self.state.fetch(&mut fetch, entity, location.table_row)) } else { Err(QueryEntityError::QueryDoesNotMatch( entity, @@ -1674,7 +1710,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn get_many_mut( &mut self, entities: [Entity; N], - ) -> Result<[D::Item<'_, 's>; N], QueryEntityError> { + ) -> Result<[D::Item<'_, '_>; N], QueryEntityError> { self.reborrow().get_many_mut_inner(entities) } @@ -1742,7 +1778,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn get_many_unique_mut( &mut self, entities: UniqueEntityArray, - ) -> Result<[D::Item<'_, 's>; N], QueryEntityError> { + ) -> Result<[D::Item<'_, '_>; N], QueryEntityError> { self.reborrow().get_many_unique_inner(entities) } @@ -1761,7 +1797,10 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn get_many_mut_inner( self, entities: [Entity; N], - ) -> Result<[D::Item<'w, 's>; N], QueryEntityError> { + ) -> Result<[D::Item<'w, 's>; N], QueryEntityError> + where + S: ConsumableQueryState<'s, Data = D, Filter = F>, + { // Verify that all entities are unique for i in 0..N { for j in 0..i { @@ -1792,6 +1831,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { ) -> Result<[D::Item<'w, 's>; N], QueryEntityError> where D: ReadOnlyQueryData, + S: ConsumableQueryState<'s, Data = D, Filter = F>, { // SAFETY: The query results are read-only, so they don't conflict if there are duplicate entities. unsafe { self.get_many_impl(entities) } @@ -1811,7 +1851,10 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub fn get_many_unique_inner( self, entities: UniqueEntityArray, - ) -> Result<[D::Item<'w, 's>; N], QueryEntityError> { + ) -> Result<[D::Item<'w, 's>; N], QueryEntityError> + where + S: ConsumableQueryState<'s, Data = D, Filter = F>, + { // SAFETY: All entities are unique, so the results don't alias. unsafe { self.get_many_impl(entities.into_inner()) } } @@ -1826,12 +1869,15 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { unsafe fn get_many_impl( self, entities: [Entity; N], - ) -> Result<[D::Item<'w, 's>; N], QueryEntityError> { + ) -> Result<[D::Item<'w, 's>; N], QueryEntityError> + where + S: ConsumableQueryState<'s, Data = D, Filter = F>, + { let mut values = [(); N].map(|_| MaybeUninit::uninit()); for (value, entity) in core::iter::zip(&mut values, entities) { // SAFETY: The caller asserts that the results don't alias - let item = unsafe { self.copy_unsafe() }.get_inner(entity)?; + let item = unsafe { self.get_unchecked_inner(entity)? }; *value = MaybeUninit::new(item); } @@ -1857,7 +1903,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { pub unsafe fn get_unchecked( &self, entity: Entity, - ) -> Result, QueryEntityError> { + ) -> Result, QueryEntityError> { // SAFETY: The caller promises that this will not result in multiple mutable references. unsafe { self.reborrow_unsafe() }.get_inner(entity) } @@ -1893,7 +1939,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// - [`single_mut`](Self::single_mut) to get the mutable query item. #[inline] - pub fn single(&self) -> Result, QuerySingleError> { + pub fn single(&self) -> Result, QuerySingleError> { self.as_readonly().single_inner() } @@ -1922,7 +1968,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// - [`single`](Self::single) to get the read-only query item. #[inline] - pub fn single_mut(&mut self) -> Result, QuerySingleError> { + pub fn single_mut(&mut self) -> Result, QuerySingleError> { self.reborrow().single_inner() } @@ -1954,7 +2000,10 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// - [`single_mut`](Self::single_mut) to get the mutable query item. /// - [`single_inner`](Self::single_inner) for the panicking version. #[inline] - pub fn single_inner(self) -> Result, QuerySingleError> { + pub fn single_inner(self) -> Result, QuerySingleError> + where + S: ConsumableQueryState<'s, Data = D, Filter = F>, + { let mut query = self.into_iter(); let first = query.next(); let extra = query.next().is_some(); @@ -2110,7 +2159,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// # world.spawn((A(10), B(5))); /// # /// fn reusable_function(lens: &mut QueryLens<&A>) { - /// assert_eq!(lens.query().single().unwrap().0, 10); + /// assert_eq!(lens.single().unwrap().0, 10); /// } /// /// // We can use the function in a system that takes the exact query. @@ -2239,7 +2288,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// # world.spawn((A(10), B(5))); /// # /// fn reusable_function(mut lens: QueryLens<&A>) { - /// assert_eq!(lens.query().single().unwrap().0, 10); + /// assert_eq!(lens.single().unwrap().0, 10); /// } /// /// // We can use the function in a system that takes the exact query. @@ -2302,12 +2351,10 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { self, ) -> QueryLens<'w, NewD, NewF> { let state = self.state.transmute_filtered::(self.world); - QueryLens { - world: self.world, - state, - last_run: self.last_run, - this_run: self.this_run, - } + // SAFETY: + // - This is memory safe because the original query had compatible access and was consumed. + // - The world matches because it was the same one used to construct self. + unsafe { Query::new(self.world, Box::new(state), self.last_run, self.this_run) } } /// Gets a [`QueryLens`] with the same accesses as the existing query @@ -2356,12 +2403,12 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// mut enemies: Query<&Enemy> /// ) { /// let mut players_transforms: QueryLens<(&Transform, &Player)> = transforms.join(&mut players); - /// for (transform, player) in &players_transforms.query() { + /// for (transform, player) in &players_transforms { /// // do something with a and b /// } /// /// let mut enemies_transforms: QueryLens<(&Transform, &Enemy)> = transforms.join(&mut enemies); - /// for (transform, enemy) in &enemies_transforms.query() { + /// for (transform, enemy) in &enemies_transforms { /// // do something with a and b /// } /// } @@ -2456,18 +2503,18 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { let state = self .state .join_filtered::(self.world, other.state); - QueryLens { - world: self.world, - state, - last_run: self.last_run, - this_run: self.this_run, - } + // SAFETY: + // - This is memory safe because the original query had compatible access and was consumed. + // - The world matches because it was the same one used to construct self. + unsafe { Query::new(self.world, Box::new(state), self.last_run, self.this_run) } } } -impl<'w, 's, D: QueryData, F: QueryFilter> IntoIterator for Query<'w, 's, D, F> { +impl<'w, 's, D: QueryData, F: QueryFilter, S: ConsumableQueryState<'s, Data = D, Filter = F>> + IntoIterator for Query<'w, 's, D, F, S> +{ type Item = D::Item<'w, 's>; - type IntoIter = QueryIter<'w, 's, D, F>; + type IntoIter = QueryIter<'w, 's, S>; fn into_iter(self) -> Self::IntoIter { // SAFETY: @@ -2478,18 +2525,22 @@ impl<'w, 's, D: QueryData, F: QueryFilter> IntoIterator for Query<'w, 's, D, F> } } -impl<'w, 's, D: QueryData, F: QueryFilter> IntoIterator for &'w Query<'_, 's, D, F> { - type Item = ROQueryItem<'w, 's, D>; - type IntoIter = QueryIter<'w, 's, D::ReadOnly, F>; +impl<'w, D: QueryData, F: QueryFilter, S: Deref>> IntoIterator + for &'w Query<'_, '_, D, F, S> +{ + type Item = ROQueryItem<'w, 'w, D>; + type IntoIter = QueryIter<'w, 'w, &'w QueryState>; fn into_iter(self) -> Self::IntoIter { self.iter() } } -impl<'w, 's, D: QueryData, F: QueryFilter> IntoIterator for &'w mut Query<'_, 's, D, F> { - type Item = D::Item<'w, 's>; - type IntoIter = QueryIter<'w, 's, D, F>; +impl<'w, D: QueryData, F: QueryFilter, S: Deref>> IntoIterator + for &'w mut Query<'_, '_, D, F, S> +{ + type Item = D::Item<'w, 'w>; + type IntoIter = QueryIter<'w, 'w, &'w QueryState>; fn into_iter(self) -> Self::IntoIter { self.iter_mut() @@ -2521,30 +2572,24 @@ impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter> Query<'w, 's, D, F> { /// # bevy_ecs::system::assert_is_system(report_names_system); /// ``` #[inline] - pub fn iter_inner(&self) -> QueryIter<'w, 's, D::ReadOnly, F> { + pub fn iter_inner(&self) -> QueryIter<'w, 's, &'s QueryState> { (*self).into_iter() } } -/// Type returned from [`Query::transmute_lens`] containing the new [`QueryState`]. +/// A [`Query`] with an owned [`QueryState`]. /// -/// Call [`query`](QueryLens::query) or [`into`](Into::into) to construct the resulting [`Query`] -pub struct QueryLens<'w, Q: QueryData, F: QueryFilter = ()> { - world: UnsafeWorldCell<'w>, - state: QueryState, - last_run: Tick, - this_run: Tick, -} +/// This is returned from methods like [`Query::transmute_lens`] that construct a fresh [`QueryState`]. +pub type QueryLens<'w, Q, F = ()> = Query<'w, 'static, Q, F, Box>>; impl<'w, Q: QueryData, F: QueryFilter> QueryLens<'w, Q, F> { /// Create a [`Query`] from the underlying [`QueryState`]. + #[deprecated( + since = "0.16.0", + note = "This can usually be removed, since `QueryLens` supports all methods from `Query`. If you need to consume the resulting `Query`, call `reborrow()` instead." + )] pub fn query(&mut self) -> Query<'_, '_, Q, F> { - Query { - world: self.world, - state: &self.state, - last_run: self.last_run, - this_run: self.this_run, - } + self.reborrow() } } @@ -2558,6 +2603,7 @@ impl<'w, Q: ReadOnlyQueryData, F: QueryFilter> QueryLens<'w, Q, F> { state: &self.state, last_run: self.last_run, this_run: self.this_run, + marker: PhantomData, } } } @@ -2566,7 +2612,7 @@ impl<'w, 's, Q: QueryData, F: QueryFilter> From<&'s mut QueryLens<'w, Q, F>> for Query<'s, 's, Q, F> { fn from(value: &'s mut QueryLens<'w, Q, F>) -> Query<'s, 's, Q, F> { - value.query() + value.reborrow() } } @@ -2741,4 +2787,14 @@ mod tests { QueryEntityError::AliasedMutability(entities[9]) ); } + + #[test] + fn query_lens_can_be_consumed() { + let mut world = World::new(); + let lens = world + .query::() + .query(&world) + .transmute_lens_inner::(); + lens.single_inner().unwrap_err(); + } } diff --git a/release-content/migration-guides/combine_query_and_query_lens.md b/release-content/migration-guides/combine_query_and_query_lens.md new file mode 100644 index 0000000000000..ddf22d41aae2a --- /dev/null +++ b/release-content/migration-guides/combine_query_and_query_lens.md @@ -0,0 +1,20 @@ +--- +title: Combine `Query` and `QueryLens` +pull_requests: [18162] +--- + +The `QueryLens::query()` method has been deprecated. +The `QueryLens` type has now been combined with `Query`, so most methods can be called directly on the `QueryLens` and the call can simply be removed. +If that doesn't work and you do need a fresh `Query`, the call to `.query()` can be replaced with `.reborrow()`. + +```rust +fn with_query(query: Query<&T>) {} +fn with_lens(lens: QueryLens<&T>) -> Result { + // 0.16 + for item in lens.query().iter() {} + with_query(lens.query()); + // 0.17 + for item in lens.iter() {} + with_query(lens.reborrow()); +} +```