From b7d52547629e3a8ff610f446c19e42fe34d9adc0 Mon Sep 17 00:00:00 2001 From: Vic <59878206+Victoronz@users.noreply.github.com> Date: Sun, 16 Mar 2025 22:12:26 +0100 Subject: [PATCH] implement get_many_unique (#18315) # Objective Continuation to #16547 and #17954. The `get_many` family are the last methods on `Query`/`QueryState` for which we're still missing a `unique` version. ## Solution Offer `get_many_unique`/`get_many_unique_mut` and `get_many_unique_inner`! Their implementation is the same as `get_many`, the difference lies in their guaranteed-to-be unique inputs, meaning we never do any aliasing checks. To reduce confusion, we also rename `get_many_readonly` into `get_many_inner` and the current `get_many_inner` into `get_many_mut_inner` to clarify their purposes. ## Testing Doc examples. ## Migration Guide `get_many_inner` is now called `get_many_mut_inner`. `get_many_readonly` is now called `get_many_inner`. --- crates/bevy_ecs/src/query/state.rs | 87 +++++++++++++- crates/bevy_ecs/src/system/query.rs | 169 +++++++++++++++++++++++++--- 2 files changed, 238 insertions(+), 18 deletions(-) diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index f9323e8076..393971c4ac 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -1,7 +1,7 @@ use crate::{ archetype::{Archetype, ArchetypeComponentId, ArchetypeGeneration, ArchetypeId}, component::{ComponentId, Tick}, - entity::{Entity, EntityBorrow, EntitySet}, + entity::{unique_array::UniqueEntityArray, Entity, EntityBorrow, EntitySet}, entity_disabling::DefaultQueryFilters, prelude::FromWorld, query::{Access, FilteredAccess, QueryCombinationIter, QueryIter, QueryParIter, WorldQuery}, @@ -997,6 +997,44 @@ impl QueryState { self.query(world).get_many_inner(entities) } + /// Returns the read-only query results for the given [`UniqueEntityArray`]. + /// + /// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is + /// returned instead. + /// + /// # Examples + /// + /// ``` + /// use bevy_ecs::{prelude::*, query::QueryEntityError, entity::{EntitySetIterator, unique_array::UniqueEntityArray, unique_vec::UniqueEntityVec}}; + /// + /// #[derive(Component, PartialEq, Debug)] + /// struct A(usize); + /// + /// let mut world = World::new(); + /// let entity_set: UniqueEntityVec = world.spawn_batch((0..3).map(A)).collect_set(); + /// let entity_set: UniqueEntityArray<3> = entity_set.try_into().unwrap(); + /// + /// world.spawn(A(73)); + /// + /// let mut query_state = world.query::<&A>(); + /// + /// let component_values = query_state.get_many_unique(&world, entity_set).unwrap(); + /// + /// assert_eq!(component_values, [&A(0), &A(1), &A(2)]); + /// + /// let wrong_entity = Entity::from_raw(365); + /// + /// assert_eq!(match query_state.get_many_unique(&mut world, UniqueEntityArray::from([wrong_entity])).unwrap_err() {QueryEntityError::EntityDoesNotExist(error) => error.entity, _ => panic!()}, wrong_entity); + /// ``` + #[inline] + pub fn get_many_unique<'w, const N: usize>( + &mut self, + world: &'w World, + entities: UniqueEntityArray, + ) -> Result<[ROQueryItem<'w, D>; N], QueryEntityError> { + self.query(world).get_many_unique_inner(entities) + } + /// Gets the query result for the given [`World`] and [`Entity`]. /// /// This is always guaranteed to run in `O(1)` time. @@ -1053,7 +1091,52 @@ impl QueryState { world: &'w mut World, entities: [Entity; N], ) -> Result<[D::Item<'w>; N], QueryEntityError> { - self.query_mut(world).get_many_inner(entities) + self.query_mut(world).get_many_mut_inner(entities) + } + + /// Returns the query results for the given [`UniqueEntityArray`]. + /// + /// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is + /// returned instead. + /// + /// ``` + /// use bevy_ecs::{prelude::*, query::QueryEntityError, entity::{EntitySetIterator, unique_array::UniqueEntityArray, unique_vec::UniqueEntityVec}}; + /// + /// #[derive(Component, PartialEq, Debug)] + /// struct A(usize); + /// + /// let mut world = World::new(); + /// + /// let entity_set: UniqueEntityVec = world.spawn_batch((0..3).map(A)).collect_set(); + /// let entity_set: UniqueEntityArray<3> = entity_set.try_into().unwrap(); + /// + /// world.spawn(A(73)); + /// + /// let mut query_state = world.query::<&mut A>(); + /// + /// let mut mutable_component_values = query_state.get_many_unique_mut(&mut world, entity_set).unwrap(); + /// + /// for mut a in &mut mutable_component_values { + /// a.0 += 5; + /// } + /// + /// let component_values = query_state.get_many_unique(&world, entity_set).unwrap(); + /// + /// assert_eq!(component_values, [&A(5), &A(6), &A(7)]); + /// + /// let wrong_entity = Entity::from_raw(57); + /// let invalid_entity = world.spawn_empty().id(); + /// + /// assert_eq!(match query_state.get_many_unique(&mut world, UniqueEntityArray::from([wrong_entity])).unwrap_err() {QueryEntityError::EntityDoesNotExist(error) => error.entity, _ => panic!()}, wrong_entity); + /// assert_eq!(match query_state.get_many_unique_mut(&mut world, UniqueEntityArray::from([invalid_entity])).unwrap_err() {QueryEntityError::QueryDoesNotMatch(entity, _) => entity, _ => panic!()}, invalid_entity); + /// ``` + #[inline] + pub fn get_many_unique_mut<'w, const N: usize>( + &mut self, + world: &'w mut World, + entities: UniqueEntityArray, + ) -> Result<[D::Item<'w>; N], QueryEntityError> { + self.query_mut(world).get_many_unique_inner(entities) } /// Gets the query result for the given [`World`] and [`Entity`]. diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index cc97433b76..2f21e4dec9 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -1,7 +1,9 @@ use crate::{ batching::BatchingStrategy, component::Tick, - entity::{Entity, EntityBorrow, EntityDoesNotExistError, EntitySet}, + entity::{ + unique_array::UniqueEntityArray, Entity, EntityBorrow, EntityDoesNotExistError, EntitySet, + }, query::{ DebugCheckedUnwrap, NopWorldQuery, QueryCombinationIter, QueryData, QueryEntityError, QueryFilter, QueryIter, QueryManyIter, QueryManyUniqueIter, QueryParIter, QueryParManyIter, @@ -1323,15 +1325,65 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// # See also /// /// - [`get_many_mut`](Self::get_many_mut) to get mutable query items. + /// - [`get_many_unique`](Self::get_many_unique) to only handle unique inputs. /// - [`many`](Self::many) for the panicking version. #[inline] pub fn get_many( &self, entities: [Entity; N], ) -> Result<[ROQueryItem<'_, D>; N], QueryEntityError> { - // Note that this calls `get_many_readonly` instead of `get_many_inner` - // since we don't need to check for duplicates. - self.as_readonly().get_many_readonly(entities) + // 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) + } + + /// Returns the read-only query items for the given [`UniqueEntityArray`]. + /// + /// The returned query items are in the same order as the input. + /// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is returned instead. + /// + /// # Examples + /// + /// ``` + /// use bevy_ecs::{prelude::*, query::QueryEntityError, entity::{EntitySetIterator, unique_array::UniqueEntityArray, unique_vec::UniqueEntityVec}}; + /// + /// #[derive(Component, PartialEq, Debug)] + /// struct A(usize); + /// + /// let mut world = World::new(); + /// let entity_set: UniqueEntityVec = world.spawn_batch((0..3).map(A)).collect_set(); + /// let entity_set: UniqueEntityArray<3> = entity_set.try_into().unwrap(); + /// + /// world.spawn(A(73)); + /// + /// let mut query_state = world.query::<&A>(); + /// let query = query_state.query(&world); + /// + /// let component_values = query.get_many_unique(entity_set).unwrap(); + /// + /// assert_eq!(component_values, [&A(0), &A(1), &A(2)]); + /// + /// let wrong_entity = Entity::from_raw(365); + /// + /// assert_eq!( + /// match query.get_many_unique(UniqueEntityArray::from([wrong_entity])).unwrap_err() { + /// QueryEntityError::EntityDoesNotExist(error) => error.entity, + /// _ => panic!(), + /// }, + /// wrong_entity + /// ); + /// ``` + /// + /// # See also + /// + /// - [`get_many_unique_mut`](Self::get_many_mut) to get mutable query items. + /// - [`get_many`](Self::get_many) to handle inputs with duplicates. + #[inline] + pub fn get_many_unique( + &self, + entities: UniqueEntityArray, + ) -> Result<[ROQueryItem<'_, D>; N], QueryEntityError> { + self.as_readonly().get_many_unique_inner(entities) } /// Returns the read-only query items for the given array of [`Entity`]. @@ -1560,7 +1612,75 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { &mut self, entities: [Entity; N], ) -> Result<[D::Item<'_>; N], QueryEntityError> { - self.reborrow().get_many_inner(entities) + self.reborrow().get_many_mut_inner(entities) + } + + /// Returns the query items for the given [`UniqueEntityArray`]. + /// + /// The returned query items are in the same order as the input. + /// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is returned instead. + /// + /// # Examples + /// + /// ``` + /// use bevy_ecs::{prelude::*, query::QueryEntityError, entity::{EntitySetIterator, unique_array::UniqueEntityArray, unique_vec::UniqueEntityVec}}; + /// + /// #[derive(Component, PartialEq, Debug)] + /// struct A(usize); + /// + /// let mut world = World::new(); + /// + /// let entity_set: UniqueEntityVec<_> = world.spawn_batch((0..3).map(A)).collect_set(); + /// let entity_set: UniqueEntityArray<3> = entity_set.try_into().unwrap(); + /// + /// world.spawn(A(73)); + /// let wrong_entity = Entity::from_raw(57); + /// let invalid_entity = world.spawn_empty().id(); + /// + /// + /// let mut query_state = world.query::<&mut A>(); + /// let mut query = query_state.query_mut(&mut world); + /// + /// let mut mutable_component_values = query.get_many_unique_mut(entity_set).unwrap(); + /// + /// for mut a in &mut mutable_component_values { + /// a.0 += 5; + /// } + /// + /// let component_values = query.get_many_unique(entity_set).unwrap(); + /// + /// assert_eq!(component_values, [&A(5), &A(6), &A(7)]); + /// + /// assert_eq!( + /// match query + /// .get_many_unique_mut(UniqueEntityArray::from([wrong_entity])) + /// .unwrap_err() + /// { + /// QueryEntityError::EntityDoesNotExist(error) => error.entity, + /// _ => panic!(), + /// }, + /// wrong_entity + /// ); + /// assert_eq!( + /// match query + /// .get_many_unique_mut(UniqueEntityArray::from([invalid_entity])) + /// .unwrap_err() + /// { + /// QueryEntityError::QueryDoesNotMatch(entity, _) => entity, + /// _ => panic!(), + /// }, + /// invalid_entity + /// ); + /// ``` + /// # See also + /// + /// - [`get_many_unique`](Self::get_many) to get read-only query items. + #[inline] + pub fn get_many_unique_mut( + &mut self, + entities: UniqueEntityArray, + ) -> Result<[D::Item<'_>; N], QueryEntityError> { + self.reborrow().get_many_unique_inner(entities) } /// Returns the query items for the given array of [`Entity`]. @@ -1573,10 +1693,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// - [`get_many`](Self::get_many) to get read-only query items without checking for duplicate entities. /// - [`get_many_mut`](Self::get_many_mut) to get items using a mutable reference. - /// - [`get_many_readonly`](Self::get_many_readonly) to get read-only query items without checking for duplicate entities - /// with the actual "inner" world lifetime. + /// - [`get_many_inner`](Self::get_many_mut_inner) to get read-only query items with the actual "inner" world lifetime. #[inline] - pub fn get_many_inner( + pub fn get_many_mut_inner( self, entities: [Entity; N], ) -> Result<[D::Item<'w>; N], QueryEntityError> { @@ -1588,7 +1707,6 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { } } } - // SAFETY: All entities are unique, so the results don't alias. unsafe { self.get_many_impl(entities) } } @@ -1603,9 +1721,9 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { /// /// - [`get_many`](Self::get_many) to get read-only query items without checking for duplicate entities. /// - [`get_many_mut`](Self::get_many_mut) to get items using a mutable reference. - /// - [`get_many_inner`](Self::get_many_readonly) to get mutable query items with the actual "inner" world lifetime. + /// - [`get_many_mut_inner`](Self::get_many_mut_inner) to get mutable query items with the actual "inner" world lifetime. #[inline] - pub fn get_many_readonly( + pub fn get_many_inner( self, entities: [Entity; N], ) -> Result<[D::Item<'w>; N], QueryEntityError> @@ -1616,6 +1734,25 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { unsafe { self.get_many_impl(entities) } } + /// Returns the query items for the given [`UniqueEntityArray`]. + /// This consumes the [`Query`] to return results with the actual "inner" world lifetime. + /// + /// The returned query items are in the same order as the input. + /// In case of a nonexisting entity, duplicate entities or mismatched component, a [`QueryEntityError`] is returned instead. + /// + /// # See also + /// + /// - [`get_many_unique`](Self::get_many_unique) to get read-only query items without checking for duplicate entities. + /// - [`get_many_unique_mut`](Self::get_many_unique_mut) to get items using a mutable reference. + #[inline] + pub fn get_many_unique_inner( + self, + entities: UniqueEntityArray, + ) -> Result<[D::Item<'w>; N], QueryEntityError> { + // SAFETY: All entities are unique, so the results don't alias. + unsafe { self.get_many_impl(entities.into_inner()) } + } + /// Returns the query items for the given array of [`Entity`]. /// This consumes the [`Query`] to return results with the actual "inner" world lifetime. /// @@ -2518,19 +2655,19 @@ mod tests { let mut query_state = world.query::(); - // It's best to test get_many_inner directly, as it is shared + // It's best to test get_many_mut_inner directly, as it is shared // We don't care about aliased mutability for the read-only equivalent // SAFETY: Query does not access world data. assert!(query_state .query_mut(&mut world) - .get_many_inner::<10>(entities.clone().try_into().unwrap()) + .get_many_mut_inner::<10>(entities.clone().try_into().unwrap()) .is_ok()); assert_eq!( query_state .query_mut(&mut world) - .get_many_inner([entities[0], entities[0]]) + .get_many_mut_inner([entities[0], entities[0]]) .unwrap_err(), QueryEntityError::AliasedMutability(entities[0]) ); @@ -2538,7 +2675,7 @@ mod tests { assert_eq!( query_state .query_mut(&mut world) - .get_many_inner([entities[0], entities[1], entities[0]]) + .get_many_mut_inner([entities[0], entities[1], entities[0]]) .unwrap_err(), QueryEntityError::AliasedMutability(entities[0]) ); @@ -2546,7 +2683,7 @@ mod tests { assert_eq!( query_state .query_mut(&mut world) - .get_many_inner([entities[9], entities[9]]) + .get_many_mut_inner([entities[9], entities[9]]) .unwrap_err(), QueryEntityError::AliasedMutability(entities[9]) );