diff --git a/crates/bevy_ecs/src/entity/entity_set.rs b/crates/bevy_ecs/src/entity/entity_set.rs new file mode 100644 index 0000000000..cc81fa4607 --- /dev/null +++ b/crates/bevy_ecs/src/entity/entity_set.rs @@ -0,0 +1,444 @@ +use alloc::{ + boxed::Box, + collections::{btree_map, btree_set}, + rc::Rc, + sync::Arc, +}; + +use core::{ + array, + fmt::{Debug, Formatter}, + iter::{self, FusedIterator}, + option, result, +}; + +use super::Entity; + +/// A trait for entity borrows. +/// +/// This trait can be thought of as `Borrow`, but yielding `Entity` directly. +pub trait EntityBorrow { + /// Returns the borrowed entity. + fn entity(&self) -> Entity; +} + +/// A trait for [`Entity`] borrows with trustworthy comparison behavior. +/// +/// Comparison trait behavior between a [`TrustedEntityBorrow`] type and its underlying entity will match. +/// This property includes [`PartialEq`], [`Eq`], [`PartialOrd`], [`Ord`] and [`Hash`], +/// and remains even after [`Clone`] and/or [`Borrow`] calls. +/// +/// # Safety +/// Any [`PartialEq`], [`Eq`], [`PartialOrd`], [`Ord`], and [`Hash`] impls must be +/// equivalent for `Self` and its underlying entity: +/// `x.entity() == y.entity()` should give the same result as `x == y`. +/// The above equivalence must also hold through and between calls to any [`Clone`] +/// and [`Borrow`]/[`BorrowMut`] impls in place of [`entity()`]. +/// +/// The result of [`entity()`] must be unaffected by any interior mutability. +/// +/// [`Hash`]: core::hash::Hash +/// [`Borrow`]: core::borrow::Borrow +/// [`BorrowMut`]: core::borrow::BorrowMut +/// [`entity()`]: EntityBorrow::entity +pub unsafe trait TrustedEntityBorrow: EntityBorrow + Eq {} + +impl EntityBorrow for Entity { + fn entity(&self) -> Entity { + *self + } +} + +// SAFETY: +// The trait implementations of Entity are correct and deterministic. +unsafe impl TrustedEntityBorrow for Entity {} + +impl EntityBorrow for &T { + fn entity(&self) -> Entity { + (**self).entity() + } +} + +// SAFETY: +// `&T` delegates `PartialEq`, `Eq`, `PartialOrd`, `Ord`, and `Hash` to T. +// `Clone` and `Borrow` maintain equality. +// `&T` is `Freeze`. +unsafe impl TrustedEntityBorrow for &T {} + +impl EntityBorrow for &mut T { + fn entity(&self) -> Entity { + (**self).entity() + } +} + +// SAFETY: +// `&mut T` delegates `PartialEq`, `Eq`, `PartialOrd`, `Ord`, and `Hash` to T. +// `Borrow` and `BorrowMut` maintain equality. +// `&mut T` is `Freeze`. +unsafe impl TrustedEntityBorrow for &mut T {} + +impl EntityBorrow for Box { + fn entity(&self) -> Entity { + (**self).entity() + } +} + +// SAFETY: +// `Box` delegates `PartialEq`, `Eq`, `PartialOrd`, `Ord`, and `Hash` to T. +// `Clone`, `Borrow` and `BorrowMut` maintain equality. +// `Box` is `Freeze`. +unsafe impl TrustedEntityBorrow for Box {} + +impl EntityBorrow for Rc { + fn entity(&self) -> Entity { + (**self).entity() + } +} + +// SAFETY: +// `Rc` delegates `PartialEq`, `Eq`, `PartialOrd`, `Ord`, and `Hash` to T. +// `Clone`, `Borrow` and `BorrowMut` maintain equality. +// `Rc` is `Freeze`. +unsafe impl TrustedEntityBorrow for Rc {} + +impl EntityBorrow for Arc { + fn entity(&self) -> Entity { + (**self).entity() + } +} + +// SAFETY: +// `Arc` delegates `PartialEq`, `Eq`, `PartialOrd`, `Ord`, and `Hash` to T. +// `Clone`, `Borrow` and `BorrowMut` maintain equality. +// `Arc` is `Freeze`. +unsafe impl TrustedEntityBorrow for Arc {} + +/// A set of unique entities. +/// +/// Any element returned by [`Self::IntoIter`] will compare non-equal to every other element in the iterator. +/// As a consequence, [`into_iter()`] on `EntitySet` will always produce another `EntitySet`. +/// +/// Implementing this trait allows for unique query iteration over a list of entities. +/// See [`iter_many_unique`] and [`iter_many_unique_mut`] +/// +/// Note that there is no guarantee of the [`IntoIterator`] impl being deterministic, +/// it might return different iterators when called multiple times. +/// Neither is there a guarantee that the comparison trait impls of `EntitySet` match that +/// of the respective [`EntitySetIterator`] (or of a [`Vec`] collected from its elements) +/// +/// [`Self::IntoIter`]: IntoIterator::IntoIter +/// [`into_iter()`]: IntoIterator::into_iter +/// [`iter_many_unique`]: crate::system::Query::iter_many_unique +/// [`iter_many_unique_mut`]: crate::system::Query::iter_many_unique_mut +pub trait EntitySet: IntoIterator {} + +impl> EntitySet for T {} + +/// An iterator over a set of unique entities. +/// +/// Every `EntitySetIterator` is also [`EntitySet`]. +/// +/// # Safety +/// +/// `x != y` must hold for any 2 elements returned by the iterator. +/// This is always true for iterators that cannot return more than one element. +pub unsafe trait EntitySetIterator: Iterator {} + +// SAFETY: +// A correct `BTreeMap` contains only unique keys. +// TrustedEntityBorrow guarantees a trustworthy Ord impl for T, and thus a correct `BTreeMap`. +unsafe impl EntitySetIterator for btree_map::Keys<'_, K, V> {} + +// SAFETY: +// A correct `BTreeMap` contains only unique keys. +// TrustedEntityBorrow guarantees a trustworthy Ord impl for T, and thus a correct `BTreeMap`. +unsafe impl EntitySetIterator for btree_map::IntoKeys {} + +// SAFETY: +// A correct `BTreeSet` contains only unique elements. +// TrustedEntityBorrow guarantees a trustworthy Ord impl for T, and thus a correct `BTreeSet`. +// The sub-range maintains uniqueness. +unsafe impl EntitySetIterator for btree_set::Range<'_, T> {} + +// SAFETY: +// A correct `BTreeSet` contains only unique elements. +// TrustedEntityBorrow guarantees a trustworthy Ord impl for T, and thus a correct `BTreeSet`. +// The "intersection" operation maintains uniqueness. +unsafe impl EntitySetIterator for btree_set::Intersection<'_, T> {} + +// SAFETY: +// A correct `BTreeSet` contains only unique elements. +// TrustedEntityBorrow guarantees a trustworthy Ord impl for T, and thus a correct `BTreeSet`. +// The "union" operation maintains uniqueness. +unsafe impl EntitySetIterator for btree_set::Union<'_, T> {} + +// SAFETY: +// A correct `BTreeSet` contains only unique elements. +// TrustedEntityBorrow guarantees a trustworthy Ord impl for T, and thus a correct `BTreeSet`. +// The "difference" operation maintains uniqueness. +unsafe impl EntitySetIterator for btree_set::Difference<'_, T> {} + +// SAFETY: +// A correct `BTreeSet` contains only unique elements. +// TrustedEntityBorrow guarantees a trustworthy Ord impl for T, and thus a correct `BTreeSet`. +// The "symmetric difference" operation maintains uniqueness. +unsafe impl EntitySetIterator + for btree_set::SymmetricDifference<'_, T> +{ +} + +// SAFETY: +// A correct `BTreeSet` contains only unique elements. +// TrustedEntityBorrow guarantees a trustworthy Ord impl for T, and thus a correct `BTreeSet`. +unsafe impl EntitySetIterator for btree_set::Iter<'_, T> {} + +// SAFETY: +// A correct `BTreeSet` contains only unique elements. +// TrustedEntityBorrow guarantees a trustworthy Ord impl for T, and thus a correct `BTreeSet`. +unsafe impl EntitySetIterator for btree_set::IntoIter {} + +// SAFETY: This iterator only returns one element. +unsafe impl EntitySetIterator for option::Iter<'_, T> {} + +// SAFETY: This iterator only returns one element. +// unsafe impl EntitySetIterator for option::IterMut<'_, T> {} + +// SAFETY: This iterator only returns one element. +unsafe impl EntitySetIterator for option::IntoIter {} + +// SAFETY: This iterator only returns one element. +unsafe impl EntitySetIterator for result::Iter<'_, T> {} + +// SAFETY: This iterator only returns one element. +// unsafe impl EntitySetIterator for result::IterMut<'_, T> {} + +// SAFETY: This iterator only returns one element. +unsafe impl EntitySetIterator for result::IntoIter {} + +// SAFETY: This iterator only returns one element. +unsafe impl EntitySetIterator for array::IntoIter {} + +// SAFETY: This iterator does not return any elements. +unsafe impl EntitySetIterator for array::IntoIter {} + +// SAFETY: This iterator only returns one element. +unsafe impl T> EntitySetIterator for iter::OnceWith {} + +// SAFETY: This iterator only returns one element. +unsafe impl EntitySetIterator for iter::Once {} + +// SAFETY: This iterator does not return any elements. +unsafe impl EntitySetIterator for iter::Empty {} + +// SAFETY: Taking a mutable reference of an iterator has no effect on its elements. +unsafe impl EntitySetIterator for &mut I {} + +// SAFETY: Boxing an iterator has no effect on its elements. +unsafe impl EntitySetIterator for Box {} + +// SAFETY: TrustedEntityBorrow ensures that Copy does not affect equality, via its restrictions on Clone. +unsafe impl<'a, T: 'a + TrustedEntityBorrow + Copy, I: EntitySetIterator> + EntitySetIterator for iter::Copied +{ +} + +// SAFETY: TrustedEntityBorrow ensures that Clone does not affect equality. +unsafe impl<'a, T: 'a + TrustedEntityBorrow + Clone, I: EntitySetIterator> + EntitySetIterator for iter::Cloned +{ +} + +// SAFETY: Discarding elements maintains uniqueness. +unsafe impl::Item) -> bool> EntitySetIterator + for iter::Filter +{ +} + +// SAFETY: Yielding only `None` after yielding it once can only remove elements, which maintains uniqueness. +unsafe impl EntitySetIterator for iter::Fuse {} + +// SAFETY: +// Obtaining immutable references the elements of an iterator does not affect uniqueness. +// TrustedEntityBorrow ensures the lack of interior mutability. +unsafe impl::Item)> EntitySetIterator + for iter::Inspect +{ +} + +// SAFETY: Reversing an iterator does not affect uniqueness. +unsafe impl EntitySetIterator for iter::Rev {} + +// SAFETY: Discarding elements maintains uniqueness. +unsafe impl EntitySetIterator for iter::Skip {} + +// SAFETY: Discarding elements maintains uniqueness. +unsafe impl::Item) -> bool> EntitySetIterator + for iter::SkipWhile +{ +} + +// SAFETY: Discarding elements maintains uniqueness. +unsafe impl EntitySetIterator for iter::Take {} + +// SAFETY: Discarding elements maintains uniqueness. +unsafe impl::Item) -> bool> EntitySetIterator + for iter::TakeWhile +{ +} + +// SAFETY: Discarding elements maintains uniqueness. +unsafe impl EntitySetIterator for iter::StepBy {} + +/// An iterator that yields unique entities. +/// +/// This wrapper can provide an [`EntitySetIterator`] implementation when an instance of `I` is known to uphold uniqueness. +pub struct UniqueEntityIter> { + iter: I, +} + +impl UniqueEntityIter { + /// Constructs a `UniqueEntityIter` from an [`EntitySetIterator`]. + pub fn from_entity_set_iterator(iter: I) -> Self { + Self { iter } + } +} +impl> UniqueEntityIter { + /// Constructs a [`UniqueEntityIter`] from an iterator unsafely. + /// + /// # Safety + /// `iter` must only yield unique elements. + /// As in, the resulting iterator must adhere to the safety contract of [`EntitySetIterator`]. + pub unsafe fn from_iterator_unchecked(iter: I) -> Self { + Self { iter } + } +} + +impl> Iterator for UniqueEntityIter { + type Item = I::Item; + + fn next(&mut self) -> Option { + self.iter.next() + } + + fn size_hint(&self) -> (usize, Option) { + self.iter.size_hint() + } +} + +impl> ExactSizeIterator for UniqueEntityIter {} + +impl> DoubleEndedIterator + for UniqueEntityIter +{ + fn next_back(&mut self) -> Option { + self.iter.next_back() + } +} + +impl> FusedIterator for UniqueEntityIter {} + +// SAFETY: The underlying iterator is ensured to only return unique elements by its construction. +unsafe impl> EntitySetIterator for UniqueEntityIter {} + +impl + AsRef<[T]>> AsRef<[T]> for UniqueEntityIter { + fn as_ref(&self) -> &[T] { + self.iter.as_ref() + } +} + +// Default does not guarantee uniqueness, meaning `I` needs to be EntitySetIterator. +impl Default for UniqueEntityIter { + fn default() -> Self { + Self { + iter: Default::default(), + } + } +} + +// Clone does not guarantee to maintain uniqueness, meaning `I` needs to be EntitySetIterator. +impl Clone for UniqueEntityIter { + fn clone(&self) -> Self { + Self { + iter: self.iter.clone(), + } + } +} + +impl + Debug> Debug for UniqueEntityIter { + fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result { + f.debug_struct("UniqueEntityIter") + .field("iter", &self.iter) + .finish() + } +} + +#[cfg(test)] +mod tests { + #[allow(unused_imports)] + use crate::prelude::{Schedule, World}; + + #[allow(unused_imports)] + use crate::component::Component; + use crate::query::{QueryState, With}; + use crate::system::Query; + use crate::world::Mut; + #[allow(unused_imports)] + use crate::{self as bevy_ecs}; + #[allow(unused_imports)] + use crate::{entity::Entity, world::unsafe_world_cell}; + + use super::UniqueEntityIter; + + #[derive(Component, Clone)] + pub struct Thing; + + #[allow(clippy::iter_skip_zero)] + #[test] + fn preserving_uniqueness() { + let mut world = World::new(); + + let mut query = QueryState::<&mut Thing>::new(&mut world); + + let spawn_batch: Vec = world.spawn_batch(vec![Thing; 1000]).collect(); + + // SAFETY: SpawnBatchIter is `EntitySetIterator`, + let mut unique_entity_iter = + unsafe { UniqueEntityIter::from_iterator_unchecked(spawn_batch.iter()) }; + + let entity_set = unique_entity_iter + .by_ref() + .filter(|_| true) + .fuse() + .inspect(|_| ()) + .rev() + .skip(0) + .skip_while(|_| false) + .take(1000) + .take_while(|_| true) + .step_by(2) + .cloned(); + + // With `iter_many_mut` collecting is not possible, because you need to drop each `Mut`/`&mut` before the next is retrieved. + let _results: Vec> = + query.iter_many_unique_mut(&mut world, entity_set).collect(); + } + + #[test] + fn nesting_queries() { + let mut world = World::new(); + + world.spawn_batch(vec![Thing; 1000]); + + pub fn system( + mut thing_entities: Query>, + mut things: Query<&mut Thing>, + ) { + things.iter_many_unique(thing_entities.iter()); + things.iter_many_unique_mut(thing_entities.iter_mut()); + } + + let mut schedule = Schedule::default(); + schedule.add_systems(system); + schedule.run(&mut world); + } +} diff --git a/crates/bevy_ecs/src/entity/mod.rs b/crates/bevy_ecs/src/entity/mod.rs index 51bfa59983..4932c5fc11 100644 --- a/crates/bevy_ecs/src/entity/mod.rs +++ b/crates/bevy_ecs/src/entity/mod.rs @@ -35,14 +35,18 @@ //! [`World::despawn`]: crate::world::World::despawn //! [`EntityWorldMut::insert`]: crate::world::EntityWorldMut::insert //! [`EntityWorldMut::remove`]: crate::world::EntityWorldMut::remove + mod clone_entities; +mod entity_set; mod map_entities; mod visit_entities; #[cfg(feature = "bevy_reflect")] use bevy_reflect::Reflect; #[cfg(all(feature = "bevy_reflect", feature = "serialize"))] use bevy_reflect::{ReflectDeserialize, ReflectSerialize}; + pub use clone_entities::*; +pub use entity_set::*; pub use map_entities::*; pub use visit_entities::*; @@ -508,6 +512,9 @@ impl<'a> Iterator for ReserveEntitiesIterator<'a> { impl<'a> ExactSizeIterator for ReserveEntitiesIterator<'a> {} impl<'a> core::iter::FusedIterator for ReserveEntitiesIterator<'a> {} +// SAFETY: Newly reserved entity values are unique. +unsafe impl EntitySetIterator for ReserveEntitiesIterator<'_> {} + /// A [`World`]'s internal metadata store on all of its entities. /// /// Contains metadata on: diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 052658ffb2..816c92e76c 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -2,14 +2,13 @@ use super::{QueryData, QueryFilter, ReadOnlyQueryData}; use crate::{ archetype::{Archetype, ArchetypeEntity, Archetypes}, component::Tick, - entity::{Entities, Entity}, + entity::{Entities, Entity, EntityBorrow, EntitySet, EntitySetIterator}, query::{ArchetypeFilter, DebugCheckedUnwrap, QueryState, StorageId}, storage::{Table, TableRow, Tables}, world::unsafe_world_cell::UnsafeWorldCell, }; use alloc::vec::Vec; use core::{ - borrow::Borrow, cmp::Ordering, fmt::{self, Debug, Formatter}, iter::FusedIterator, @@ -1103,6 +1102,9 @@ 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> {} + impl<'w, 's, D: QueryData, F: QueryFilter> Debug for QueryIter<'w, 's, D, F> { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { f.debug_struct("QueryIter").finish() @@ -1237,6 +1239,14 @@ 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> +{ +} + impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Debug for QuerySortedIter<'w, 's, D, F, I> { @@ -1251,10 +1261,7 @@ 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> -where - I::Item: Borrow, -{ +pub struct QueryManyIter<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> { world: UnsafeWorldCell<'w>, entity_iter: I, entities: &'w Entities, @@ -1265,7 +1272,7 @@ where query_state: &'s QueryState, } -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> +impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> QueryManyIter<'w, 's, D, F, I> { /// # Safety @@ -1304,7 +1311,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> /// It is always safe for shared access. #[inline(always)] unsafe fn fetch_next_aliased_unchecked( - entity_iter: impl Iterator>, + entity_iter: impl Iterator, entities: &'w Entities, tables: &'w Tables, archetypes: &'w Archetypes, @@ -1312,8 +1319,8 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> filter: &mut F::Fetch<'w>, query_state: &'s QueryState, ) -> Option> { - for entity in entity_iter { - let entity = *entity.borrow(); + for entity_borrow in entity_iter { + let entity = entity_borrow.entity(); let Some(location) = entities.get(entity) else { continue; }; @@ -2048,7 +2055,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> } } -impl<'w, 's, D: QueryData, F: QueryFilter, I: DoubleEndedIterator>> +impl<'w, 's, D: QueryData, F: QueryFilter, I: DoubleEndedIterator> QueryManyIter<'w, 's, D, F, I> { /// Get next result from the back of the query @@ -2074,7 +2081,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: DoubleEndedIterator>> Iterator +impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator> Iterator for QueryManyIter<'w, 's, D, F, I> { type Item = D::Item<'w>; @@ -2103,13 +2110,8 @@ impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator>, - > DoubleEndedIterator for QueryManyIter<'w, 's, D, F, I> +impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: DoubleEndedIterator> + DoubleEndedIterator for QueryManyIter<'w, 's, D, F, I> { #[inline(always)] fn next_back(&mut self) -> Option { @@ -2131,12 +2133,18 @@ impl< } // This is correct as [`QueryManyIter`] always returns `None` once exhausted. -impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator>> FusedIterator +impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter, I: Iterator> FusedIterator for QueryManyIter<'w, 's, D, F, I> { } -impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> Debug +// SAFETY: Fetching unique entities maintains uniqueness. +unsafe impl<'w, 's, F: QueryFilter, I: EntitySetIterator> EntitySetIterator + for QueryManyIter<'w, 's, Entity, F, I> +{ +} + +impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Debug for QueryManyIter<'w, 's, D, F, I> { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { @@ -2144,6 +2152,93 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator>> De } } +/// An [`Iterator`] over the query items generated from an iterator of unique [`Entity`]s. +/// +/// Items are returned in the order of the provided iterator. +/// Entities that don't match the query are skipped. +/// +/// In contrast with [`QueryManyIter`], this allows for mutable iteration without a [`fetch_next`] method. +/// +/// This struct is created by the [`iter_many_unique`] and [`iter_many_unique_mut`] methods on [`Query`]. +/// +/// [`fetch_next`]: QueryManyIter::fetch_next +/// [`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>, +); + +impl<'w, 's, D: QueryData, F: QueryFilter, I: EntitySetIterator> + QueryManyUniqueIter<'w, 's, D, F, I> +{ + /// # Safety + /// - `world` must have permission to access any of the components registered in `query_state`. + /// - `world` must be the same one used to initialize `query_state`. + pub(crate) unsafe fn new>( + world: UnsafeWorldCell<'w>, + query_state: &'s QueryState, + entity_list: EntityList, + last_run: Tick, + this_run: Tick, + ) -> QueryManyUniqueIter<'w, 's, D, F, I> { + QueryManyUniqueIter(QueryManyIter::new( + world, + query_state, + entity_list, + last_run, + this_run, + )) + } +} + +impl<'w, 's, D: QueryData, F: QueryFilter, I: EntitySetIterator> Iterator + for QueryManyUniqueIter<'w, 's, D, F, I> +{ + type Item = D::Item<'w>; + + #[inline(always)] + 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( + &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, + ) + } + } + + fn size_hint(&self) -> (usize, Option) { + let (_, max_size) = self.0.entity_iter.size_hint(); + (0, max_size) + } +} + +// 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> +{ +} + +// SAFETY: Fetching unique entities maintains uniqueness. +unsafe impl<'w, 's, F: QueryFilter, I: EntitySetIterator> EntitySetIterator + for QueryManyUniqueIter<'w, 's, Entity, F, I> +{ +} + +impl<'w, 's, D: QueryData, F: QueryFilter, I: EntitySetIterator> Debug + for QueryManyUniqueIter<'w, 's, D, F, I> +{ + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("QueryManyUniqueIter").finish() + } +} + /// An [`Iterator`] over sorted query results of a [`QueryManyIter`]. /// /// This struct is created by the [`sort`](QueryManyIter), [`sort_unstable`](QueryManyIter), diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index c018ff29c1..36d82772d4 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -2,7 +2,7 @@ use crate::{ archetype::{Archetype, ArchetypeComponentId, ArchetypeGeneration, ArchetypeId}, batching::BatchingStrategy, component::{ComponentId, Tick}, - entity::Entity, + entity::{Entity, EntityBorrow, EntitySet}, prelude::FromWorld, query::{ Access, DebugCheckedUnwrap, FilteredAccess, QueryCombinationIter, QueryIter, QueryParIter, @@ -11,16 +11,17 @@ use crate::{ storage::{SparseSetIndex, TableId}, world::{unsafe_world_cell::UnsafeWorldCell, World, WorldId}, }; + use alloc::vec::Vec; -use core::{borrow::Borrow, fmt, mem::MaybeUninit, ptr}; +#[cfg(feature = "trace")] +use bevy_utils::tracing::Span; +use core::{fmt, mem::MaybeUninit, ptr}; use fixedbitset::FixedBitSet; use log::warn; -#[cfg(feature = "trace")] -use tracing::Span; use super::{ NopWorldQuery, QueryBuilder, QueryData, QueryEntityError, QueryFilter, QueryManyIter, - QuerySingleError, ROQueryItem, + QueryManyUniqueIter, QuerySingleError, ROQueryItem, }; /// An ID for either a table or an archetype. Used for Query iteration. @@ -1253,7 +1254,7 @@ impl QueryState { /// /// - [`iter_many_mut`](Self::iter_many_mut) to get mutable query items. #[inline] - pub fn iter_many<'w, 's, EntityList: IntoIterator>>( + pub fn iter_many<'w, 's, EntityList: IntoIterator>( &'s mut self, world: &'w World, entities: EntityList, @@ -1285,7 +1286,7 @@ impl QueryState { /// - [`iter_many`](Self::iter_many) to update archetypes. /// - [`iter_manual`](Self::iter_manual) to iterate over all query items. #[inline] - pub fn iter_many_manual<'w, 's, EntityList: IntoIterator>>( + pub fn iter_many_manual<'w, 's, EntityList: IntoIterator>( &'s self, world: &'w World, entities: EntityList, @@ -1307,7 +1308,7 @@ impl QueryState { /// Items are returned in the order of the list of entities. /// Entities that don't match the query are skipped. #[inline] - pub fn iter_many_mut<'w, 's, EntityList: IntoIterator>>( + pub fn iter_many_mut<'w, 's, EntityList: IntoIterator>( &'s mut self, world: &'w mut World, entities: EntityList, @@ -1326,6 +1327,88 @@ impl QueryState { } } + /// Returns an [`Iterator`] over the unique read-only query items generated from an [`EntitySet`]. + /// + /// Items are returned in the order of the list of entities. + /// Entities that don't match the query are skipped. + /// + /// # See also + /// + /// - [`iter_many_unique_mut`](Self::iter_many_unique_mut) to get mutable query items. + #[inline] + pub fn iter_many_unique<'w, 's, EntityList: EntitySet>( + &'s mut self, + world: &'w World, + entities: EntityList, + ) -> QueryManyUniqueIter<'w, 's, D::ReadOnly, F, EntityList::IntoIter> { + self.update_archetypes(world); + // SAFETY: query is read only + unsafe { + self.as_readonly().iter_many_unique_unchecked_manual( + entities, + world.as_unsafe_world_cell_readonly(), + world.last_change_tick(), + world.read_change_tick(), + ) + } + } + + /// Returns an [`Iterator`] over the unique read-only query items generated from an [`EntitySet`]. + /// + /// Items are returned in the order of the list of entities. + /// Entities that don't match the query are skipped. + /// + /// If `world` archetypes changed since [`Self::update_archetypes`] was last called, + /// this will skip entities contained in new archetypes. + /// + /// This can only be called for read-only queries. + /// + /// # See also + /// + /// - [`iter_many_unique`](Self::iter_many) to update archetypes. + /// - [`iter_many`](Self::iter_many) to iterate over a non-unique entity list. + /// - [`iter_manual`](Self::iter_manual) to iterate over all query items. + #[inline] + pub fn iter_many_unique_manual<'w, 's, EntityList: EntitySet>( + &'s self, + world: &'w World, + entities: EntityList, + ) -> QueryManyUniqueIter<'w, 's, D::ReadOnly, F, EntityList::IntoIter> { + self.validate_world(world.id()); + // SAFETY: query is read only, world id is validated + unsafe { + self.as_readonly().iter_many_unique_unchecked_manual( + entities, + world.as_unsafe_world_cell_readonly(), + world.last_change_tick(), + world.read_change_tick(), + ) + } + } + + /// Returns an iterator over the unique query items generated from an [`EntitySet`]. + /// + /// Items are returned in the order of the list of entities. + /// Entities that don't match the query are skipped. + #[inline] + pub fn iter_many_unique_mut<'w, 's, EntityList: EntitySet>( + &'s mut self, + world: &'w mut World, + entities: EntityList, + ) -> QueryManyUniqueIter<'w, 's, D, F, EntityList::IntoIter> { + self.update_archetypes(world); + let last_change_tick = world.last_change_tick(); + let change_tick = world.change_tick(); + // SAFETY: Query has unique world access. + unsafe { + self.iter_many_unique_unchecked_manual( + entities, + world.as_unsafe_world_cell(), + last_change_tick, + change_tick, + ) + } + } /// Returns an [`Iterator`] over the query results for the given [`World`]. /// /// This iterator is always guaranteed to return results from each matching entity once and only once. @@ -1412,11 +1495,34 @@ impl QueryState { this_run: Tick, ) -> QueryManyIter<'w, 's, D, F, EntityList::IntoIter> where - EntityList: IntoIterator>, + EntityList: IntoIterator, { QueryManyIter::new(world, self, entities, last_run, this_run) } + /// Returns an [`Iterator`] for the given [`World`] and an [`EntitySet`], where the last change and + /// the current change tick are given. + /// + /// Items are returned in the order of the list of entities. + /// Entities that don't match the query are skipped. + /// + /// # Safety + /// + /// This does not check for mutable query correctness. To be safe, make sure mutable queries + /// have unique access to the components they query. + /// This does not validate that `world.id()` matches `self.world_id`. Calling this on a `world` + /// with a mismatched [`WorldId`] is unsound. + #[inline] + pub(crate) unsafe fn iter_many_unique_unchecked_manual<'w, 's, EntityList: EntitySet>( + &'s self, + entities: EntityList, + world: UnsafeWorldCell<'w>, + last_run: Tick, + this_run: Tick, + ) -> QueryManyUniqueIter<'w, 's, D, F, EntityList::IntoIter> { + QueryManyUniqueIter::new(world, self, entities, last_run, this_run) + } + /// Returns an [`Iterator`] over all possible combinations of `K` query results for the /// given [`World`] without repetition. /// This can only be called for read-only queries. diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index b36be3fee7..6fc2a4caa5 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -1,15 +1,15 @@ use crate::{ batching::BatchingStrategy, component::Tick, - entity::Entity, + entity::{Entity, EntityBorrow, EntitySet}, query::{ QueryCombinationIter, QueryData, QueryEntityError, QueryFilter, QueryIter, QueryManyIter, - QueryParIter, QuerySingleError, QueryState, ROQueryItem, ReadOnlyQueryData, + QueryManyUniqueIter, QueryParIter, QuerySingleError, QueryState, ROQueryItem, + ReadOnlyQueryData, }, world::unsafe_world_cell::UnsafeWorldCell, }; use core::{ - borrow::Borrow, marker::PhantomData, ops::{Deref, DerefMut}, }; @@ -631,7 +631,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// - [`iter_many_mut`](Self::iter_many_mut) to get mutable query items. #[inline] - pub fn iter_many>>( + pub fn iter_many>( &self, entities: EntityList, ) -> QueryManyIter<'_, 's, D::ReadOnly, F, EntityList::IntoIter> { @@ -682,7 +682,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// # bevy_ecs::system::assert_is_system(system); /// ``` #[inline] - pub fn iter_many_mut>>( + pub fn iter_many_mut>( &mut self, entities: EntityList, ) -> QueryManyIter<'_, 's, D, F, EntityList::IntoIter> { @@ -697,6 +697,130 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { } } + /// Returns an [`Iterator`] over the unique read-only query items generated from an [`EntitySet`]. + /// + /// Items are returned in the order of the list of entities. Entities that don't match the query are skipped. + /// + /// # Example + /// + /// ``` + /// # use bevy_ecs::{prelude::*, entity::{EntitySet, UniqueEntityIter}}; + /// # use core::slice; + /// # #[derive(Component)] + /// # struct Counter { + /// # value: i32 + /// # } + /// # + /// // `Friends` ensures that it only lists unique entities. + /// #[derive(Component)] + /// struct Friends { + /// unique_list: Vec, + /// } + /// + /// impl<'a> IntoIterator for &'a Friends { + /// + /// type Item = &'a Entity; + /// type IntoIter = UniqueEntityIter>; + /// + /// fn into_iter(self) -> Self::IntoIter { + /// // SAFETY: `Friends` ensures that it unique_list contains only unique entities. + /// unsafe { UniqueEntityIter::from_iterator_unchecked(self.unique_list.iter()) } + /// } + /// } + /// + /// fn system( + /// friends_query: Query<&Friends>, + /// counter_query: Query<&Counter>, + /// ) { + /// for friends in &friends_query { + /// for counter in counter_query.iter_many_unique(friends) { + /// println!("Friend's counter: {:?}", counter.value); + /// } + /// } + /// } + /// # bevy_ecs::system::assert_is_system(system); + /// ``` + /// + /// # See also + /// + /// - [`iter_many_unique_mut`](Self::iter_many_unique_mut) to get mutable query items. + #[inline] + pub fn iter_many_unique( + &self, + entities: EntityList, + ) -> QueryManyUniqueIter<'_, 's, D::ReadOnly, F, EntityList::IntoIter> { + // SAFETY: + // - `self.world` has permission to access the required components. + // - The query is read-only, so it can be aliased even if it was originally mutable. + unsafe { + self.state.as_readonly().iter_many_unique_unchecked_manual( + entities, + self.world, + self.last_run, + self.this_run, + ) + } + } + + /// Returns an iterator over the unique query items generated from an [`EntitySet`]. + /// + /// Items are returned in the order of the list of entities. Entities that don't match the query are skipped. + /// + /// # Examples + /// + /// ``` + /// # use bevy_ecs::{prelude::*, entity::{EntitySet, UniqueEntityIter}}; + /// # use core::slice; + /// #[derive(Component)] + /// struct Counter { + /// value: i32 + /// } + /// + /// // `Friends` ensures that it only lists unique entities. + /// #[derive(Component)] + /// struct Friends { + /// unique_list: Vec, + /// } + /// + /// impl<'a> IntoIterator for &'a Friends { + /// type Item = &'a Entity; + /// type IntoIter = UniqueEntityIter>; + /// + /// fn into_iter(self) -> Self::IntoIter { + /// // SAFETY: `Friends` ensures that it unique_list contains only unique entities. + /// unsafe { UniqueEntityIter::from_iterator_unchecked(self.unique_list.iter()) } + /// } + /// } + /// + /// fn system( + /// friends_query: Query<&Friends>, + /// mut counter_query: Query<&mut Counter>, + /// ) { + /// for friends in &friends_query { + /// for mut counter in counter_query.iter_many_unique_mut(friends) { + /// println!("Friend's counter: {:?}", counter.value); + /// counter.value += 1; + /// } + /// } + /// } + /// # bevy_ecs::system::assert_is_system(system); + /// ``` + #[inline] + pub fn iter_many_unique_mut( + &mut self, + entities: EntityList, + ) -> QueryManyUniqueIter<'_, 's, D, F, EntityList::IntoIter> { + // SAFETY: `self.world` has permission to access the required components. + unsafe { + self.state.iter_many_unique_unchecked_manual( + entities, + self.world, + self.last_run, + self.this_run, + ) + } + } + /// Returns an [`Iterator`] over the query items. /// /// This iterator is always guaranteed to return results from each matching entity once and only once. @@ -761,7 +885,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// # See also /// /// - [`iter_many_mut`](Self::iter_many_mut) to safely access the query items. - pub unsafe fn iter_many_unsafe>>( + pub unsafe fn iter_many_unsafe>( &self, entities: EntityList, ) -> QueryManyIter<'_, 's, D, F, EntityList::IntoIter> { @@ -778,6 +902,35 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { } } + /// Returns an [`Iterator`] over the unique query items generated from an [`Entity`] list. + /// + /// Items are returned in the order of the list of entities. Entities that don't match the query are skipped. + /// + /// # Safety + /// + /// This allows aliased mutability. + /// You must make sure this call does not result in multiple mutable references to the same component. + /// + /// # See also + /// + /// - [`iter_many_mut`](Self::iter_many_mut) to safely access the query items. + pub unsafe fn iter_many_unique_unsafe( + &self, + entities: EntityList, + ) -> QueryManyUniqueIter<'_, 's, D, F, EntityList::IntoIter> { + // SAFETY: + // - `self.world` has permission to access the required components. + // - The caller ensures that this operation will not result in any aliased mutable accesses. + unsafe { + self.state.iter_many_unique_unchecked_manual( + entities, + self.world, + self.last_run, + self.this_run, + ) + } + } + /// Returns a parallel iterator over the query results for the given [`World`]. /// /// This parallel iterator is always guaranteed to return results from each matching entity once and diff --git a/crates/bevy_ecs/src/world/spawn_batch.rs b/crates/bevy_ecs/src/world/spawn_batch.rs index dbdbc1cfac..6be8613695 100644 --- a/crates/bevy_ecs/src/world/spawn_batch.rs +++ b/crates/bevy_ecs/src/world/spawn_batch.rs @@ -1,6 +1,6 @@ use crate::{ bundle::{Bundle, BundleSpawner}, - entity::Entity, + entity::{Entity, EntitySetIterator}, world::World, }; use core::iter::FusedIterator; @@ -110,3 +110,11 @@ where T: Bundle, { } + +// SAFETY: Newly spawned entities are unique. +unsafe impl EntitySetIterator for SpawnBatchIter<'_, I> +where + I: FusedIterator, + T: Bundle, +{ +}