From a4b89d0d5e98f14f1a81f4963a63d788c0528141 Mon Sep 17 00:00:00 2001 From: Vic <59878206+Victoronz@users.noreply.github.com> Date: Wed, 18 Dec 2024 01:49:01 +0100 Subject: [PATCH] implement EntitySet and iter_many_unique methods (#16547) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Objective In current Bevy, it is very inconvenient to mutably retrieve a user-provided list of entities more than one element at a time. If the list contains any duplicate entities, we risk mutable aliasing. Users of `Query::iter_many_mut` do not have access to `Iterator` trait, and thus miss out on common functionality, for instance collecting their `QueryManyIter`. We can circumvent this issue with validation, however that entails checking every entity against all others for inequality, or utilizing an `EntityHashSet`. Even if an entity list remains unchanged, this validation is/would have to be redone every time we wish to fetch with the list. This presents a lot of wasted work, as we often trivially know an entity list to be unique f.e.: `QueryIter` will fetch every `Entity` once and only once. As more things become entities – assets, components, queries – this issue will become more pronounced. `get_many`/`many`/`iter_many`/`par_iter_many`-like functionality is all affected. ## Solution The solution this PR proposes is to introduce functionality built around a new trait: `EntitySet`. The goal is to preserve the property of "uniqueness" in a list wherever possible, and then rely on it as a bound within new `*_many_unique` methods to avoid the need for validation. This is achieved using `Iterator`: `EntitySet` is blanket implemented for any `T` that implements `IntoIterator`. `EntitySetIterator` is the unsafe trait that actually guarantees an iterator to be "unique" via its safety contract. We define an "Iterator over unique entities" as: "No two entities returned by the iterator may compare equal." For iterators that cannot return more than 1 element, this is trivially true. Whether an iterator can satisfy this is up to the `EntitySetIterator` implementor to ensure, hence the unsafe. However, this is not yet a complete solution. Looking at the signature of `iter_many`, we find that `IntoIterator::Item` is not `Entity`, but is instead bounded by the `Borrow` trait. That is because iteration without consuming the collection will often yield us references, not owned items. `Borrow` presents an issue: The `Borrow` docs state that `x = y` should equal `x.borrow() = y.borrow()`, but unsafe cannot rely on this for soundness. We run into similar problems with other trait implementations of any `Borrow` type: `PartialEq`, `Eq`, `PartialOrd`, `Ord`, `Hash`, `Clone`, `Borrow`, and `BorrowMut`. This PR solves this with the unsafe `TrustedEntityBorrow` trait: Any implementor promises that the behavior of the aforementioned traits matches that of the underlying entity. While `Borrow` was the inspiration, we use our own counterpart trait `EntityBorrow` as the supertrait to `TrustedEntityBorrow`, so we can circumvent the limitations of the existing `Borrow` blanket impls. All together, these traits allow us to implement `*_many_unique` functionality with a lone `EntitySet` bound. `EntitySetIterator` is implemented for all the std iterators and iterator adapters that guarantee or preserve uniqueness, so we can filter, skip, take, step, reverse, ... our unique entity iterators without worry! Sadly, current `HashSet` iterators do not carry the necessary type information with them to determine whether the source `HashSet` produces logic errors; A malicious `Hasher` could compromise a `HashSet`. `HashSet` iteration is generally discouraged in the first place, so we also exclude the set operation iterators, even though they do carry the `Hasher` type parameter. `BTreeSet` implements `EntitySet` without any problems. If an iterator type cannot guarantee uniqueness at compile time, then a user can still attach `EntitySetIterator` to an individual instance of that type via `UniqueEntityIter::from_iterator_unchecked`. With this, custom types can use `UniqueEntityIter` as their `IntoIterator::IntoIter` type, if necessary. This PR is focused on the base concept, and expansions on it are left for follow-up PRs. See "Potential Future Work" below. ## Testing Doctests on `iter_many_unique`/`iter_many_unique_mut` + 2 tests in entity_set.rs. ## Showcase ```rust // Before: fn system(player_list: Res, players: Query<&mut Player>) { let value = 0; while let Some(player) = players.iter_many_mut(player_list).fetch_next() { value += mem::take(player.value_mut()) } } // After: fn system(player_list: Res, players: Query<&mut Player>) { let value = players .iter_many_unique_mut(player_list) .map(|player| mem::take(player.value_mut())) .sum(); } ``` ## Changelog - added `EntityBorrow`, `TrustedEntityBorrow`, `EntitySet` and `EntitySetIterator` traits - added `iter_many_unique`, `iter_many_unique_mut`, `iter_many_unique_unsafe` methods on `Query` - added `iter_many_unique`, `iter_many_unique_mut`, `iter_many_unique_manual` and `iter_many_unique_unchecked_manual` methods on `QueryState` - added corresponding `QueryManyUniqueIter` - added `UniqueEntityIter` ## Migration Guide Any custom type used as a `Borrow` entity list item for an `iter_many` method now has to implement `EntityBorrow` instead. Any type that implements `Borrow` can trivially implement `EntityBorrow`. ## Potential Future Work - `ToEntitySet` trait for converting any entity iterator into an `EntitySetIterator` - `EntityIndexSet/Map` to tie in hashing with `EntitySet` - add `EntityIndexSetSlice/MapSlice` - requires: `EntityIndexSet/Map` - Implementing `par_iter_many_unique_mut` for parallel mutable iteration - requires: `par_iter_many` - allow collecting into `UniqueEntityVec` to store entity sets - add `UniqueEntitySlice`s - Doesn't require, but should be done after: `UniqueEntityVec` - add `UniqueEntityArray`s - Doesn't require, but should be done after: `UniqueEntitySlice` - `get_many_unique`/`many_unique` methods - requires: `UniqueEntityArray` - `World::entity_unique` to match `World::entity` methods - Doesn't require, but makes sense after: `get_many_unique`/`many_unique` - implement `TrustedEntityBorrow` for the `EntityRef` family - Doesn't require, but makes sense after: `UniqueEntityVec` --- crates/bevy_ecs/src/entity/entity_set.rs | 444 +++++++++++++++++++++++ crates/bevy_ecs/src/entity/mod.rs | 7 + crates/bevy_ecs/src/query/iter.rs | 137 +++++-- crates/bevy_ecs/src/query/state.rs | 124 ++++++- crates/bevy_ecs/src/system/query.rs | 165 ++++++++- crates/bevy_ecs/src/world/spawn_batch.rs | 10 +- 6 files changed, 850 insertions(+), 37 deletions(-) create mode 100644 crates/bevy_ecs/src/entity/entity_set.rs 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 0000000000000..cc81fa4607cd0 --- /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 51bfa59983309..4932c5fc110bf 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 052658ffb2163..816c92e76ca14 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 c018ff29c1dd7..36d82772d4783 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 b36be3fee7597..6fc2a4caa5d47 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 dbdbc1cfac31b..6be86136953c3 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, +{ +}