diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index f675de4d89a1d..8368fa1bb97d7 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -190,6 +190,141 @@ where } } + /// Returns the read-only query results for the provided Array of [`Entity`]s. + /// + /// These values follow the order of your input array. + /// In case of a nonexisting entity, + /// a [`QueryEntityError`] is returned instead. + /// + /// # Example + /// ```rust + /// # use bevy_ecs::prelude::*; + /// #[derive(Component, PartialEq, Debug)] + /// struct A(u64); + /// + /// let mut world = World::new(); + /// let entity_1 = world.spawn().insert(A(1)).id(); + /// let entity_2 = world.spawn().insert(A(2)).id(); + /// let entity_3 = world.spawn().insert(A(3)).id(); + /// + /// let query_state = world.query::<&A>(); + /// let mut a_iterator = a_query.get_multiple([entity_3, entity_2, entity_1]).map(|i|i.unwrap()); + /// assert_eq!(*a_iterator.next().unwrap(), A(3)); + /// assert_eq!(*a_iterator.next().unwrap(), A(2)); + /// assert_eq!(*a_iterator.next().unwrap(), A(1)); + /// ``` + #[inline] + pub fn get_multiple<'w, 's, const N: usize>( + &'s self, + world: &'w World, + entities: [Entity; N], + ) -> Result<[>::Item; N], QueryEntityError> { + let array_of_results = entities.map(move |e| self.get(world, e)); + + // If any of the entities were missing, return the first error + for result in array_of_results { + result?; + } + + // At this point, we're guaranteed that all results are okay + Ok(array_of_results.map(move |result| result.unwrap())) + } + + /// Returns the read-only query items for the provided Array of [`Entity`]s. + /// + /// # Panics + /// + /// Panics if any entities do not exist. + #[inline] + pub fn multiple<'w, 's, const N: usize>( + &'s self, + world: &'w World, + entities: [Entity; N], + ) -> [>::Item; N] { + self.get_multiple(world, entities).unwrap() + } + + /// Returns the query results for the provided Array of [`Entity`]s. + /// + /// These values follow the order of your input array. + /// In case of nonunique or nonexisting entities or a mismatched component, + /// a [`QueryEntityError`] is returned instead. + /// + /// If you absolutely cannot afford the overhead of verifying uniqueness in this way, + /// you can (carefully) call the unsafe [`get_unchecked`](Self::get_unchecked) method repeatedly instead. + /// + /// # Example + /// ```rust + /// # use bevy_ecs::prelude::*; + /// #[derive(Component, PartialEq, Debug)] + /// struct A(u64); + /// + /// let mut world = World::new(); + /// let entity_1 = world.spawn().insert(A(1)).id(); + /// let entity_2 = world.spawn().insert(A(2)).id(); + /// let entity_3 = world.spawn().insert(A(3)).id(); + /// + /// let query_state = world.query::<&mut A>(); + /// let mut a_iterator = a_query.get_multiple_mut([entity_1, entity_3]).map(|i|i.unwrap()); + /// let mut a_1 = a_iterator.next().unwrap(); + /// let mut a_3 = a_iterator.next().unwrap(); + /// + /// *a_1 = A(11); + /// *a_3 = A(33); + /// + /// // Manually drop references so we can access the `World` again + /// std::mem::drop(a_iterator); + /// std::mem::drop(a_query); + /// + /// assert_eq!(*world.get::(entity_1).unwrap(), A(11)); + /// assert_eq!(*world.get::(entity_2).unwrap(), A(2)); + /// assert_eq!(*world.get::(entity_3).unwrap(), A(33)); + /// ``` + #[inline] + pub fn get_multiple_mut<'w, 's, const N: usize>( + &'s mut self, + world: &'w mut World, + entities: [Entity; N], + ) -> Result<[>::Item; N], QueryEntityError> { + // Brute force verification of uniqueness + for entity_i in entities { + for entity_j in entities { + if entity_i == entity_j { + return Err(QueryEntityError::AliasedMutability(entity_i)); + } + } + } + + // SAFE: the entities are checked for uniqueness above + // No other references to the query can be live, as this method takes &mut self + // The World cannot be modified in other ways as we take &mut World + unsafe { + let array_of_results = entities.map(move |e| self.get_unchecked(world, e)); + + // If any of the entities were missing, return the first error + for result in array_of_results { + result?; + } + + // At this point, we're guaranteed that all results are okay + Ok(array_of_results.map(move |result| result.unwrap())) + } + } + + /// Returns the query items for the provided Array of [`Entity`]s. + /// + /// # Panics + /// + /// Panics if any entities do not exist, or any entities are repeated. + #[inline] + pub fn multiple_mut<'w, 's, const N: usize>( + &'s mut self, + world: &'w mut World, + entities: [Entity; N], + ) -> [>::Item; N] { + self.get_multiple_mut(world, entities).unwrap() + } + /// Gets the query result for the given [`World`] and [`Entity`]. /// /// # Safety @@ -228,7 +363,7 @@ where let location = world .entities .get(entity) - .ok_or(QueryEntityError::NoSuchEntity)?; + .ok_or(QueryEntityError::NoSuchEntity(entity))?; if !self .matched_archetypes .contains(location.archetype_id.index()) @@ -733,10 +868,12 @@ where } /// An error that occurs when retrieving a specific [`Entity`]'s query result. -#[derive(Error, Debug)] +#[derive(Error, Debug, Clone, Copy)] pub enum QueryEntityError { #[error("The given entity does not have the requested component.")] QueryDoesNotMatch, #[error("The requested entity does not exist.")] - NoSuchEntity, + NoSuchEntity(Entity), + #[error("The same entity was accessed mutably more than once.")] + AliasedMutability(Entity), } diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index df1098fd71513..1021d6020332b 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -233,6 +233,9 @@ use thiserror::Error; /// If you have an [`Entity`] ID, you can use the [`get`](Self::get) or /// [`get_mut`](Self::get_mut) methods to access the query result for that particular entity. /// +/// If you require access to the data of multiple entities at once, +/// you can use the [`get_multiple`](Self::get_multiple) or [`get_multiple_mut`](Self::get_multiple_mut) methods. +/// /// ## Getting a single query result /// /// While it's possible to get a single result from a query by using `iter.next()`, a more @@ -551,7 +554,7 @@ where }; } - /// Returns the query result for the given [`Entity`]. + /// Returns the read-only query result for the given [`Entity`]. /// /// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is /// returned instead. @@ -584,7 +587,7 @@ where /// ``` #[inline] pub fn get( - &'s self, + &self, entity: Entity, ) -> Result<>::Item, QueryEntityError> { // SAFE: system runs without conflicts with other systems. @@ -640,6 +643,60 @@ where } } + /// Returns the read-only query results for the provided Array of [`Entity`]s. + /// + /// These values follow the order of your input array. + /// In case of a nonexisting entity, + /// a [`QueryEntityError`] is returned instead. + #[inline] + pub fn get_multiple( + &self, + entities: [Entity; N], + ) -> Result<[>::Item; N], QueryEntityError> { + self.state.get_multiple(self.world, entities) + } + + /// Returns the read-only query items for the provided Array of [`Entity`]s. + /// + /// # Panics + /// + /// Panics if any entities do not exist. + #[inline] + pub fn multiple( + &self, + entities: [Entity; N], + ) -> [>::Item; N] { + self.state.multiple(self.world, entities) + } + + /// Returns the query results for the provided Array of [`Entity`]s. + /// + /// These values follow the order of your input array. + /// In case of nonunique or nonexisting entities or a mismatched component, + /// a [`QueryEntityError`] is returned instead. + /// + /// If you absolutely cannot afford the overhead of verifying uniqueness in this way, + /// you can (carefully) call the unsafe [`get_unchecked`](Self::get_unchecked) method repeatedly instead. + #[inline] + pub fn get_multiple_mut( + &mut self, + entities: [Entity; N], + ) -> Result<[>::Item; N], QueryEntityError> { + self.state.get_multiple_mut(self.world, entities) + } + + /// Returns the query items for the provided Array of [`Entity`]s. + /// + /// # Panics + /// Panics if any entities do not exist, or any entities are repeated. + #[inline] + pub fn multiple_mut( + &mut self, + entities: [Entity; N], + ) -> [>::Item; N] { + self.state.multiple_mut(self.world, entities) + } + /// Returns the query result for the given [`Entity`]. /// /// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is @@ -651,7 +708,7 @@ where /// this call does not result in multiple mutable references to the same component #[inline] pub unsafe fn get_unchecked( - &'s self, + &self, entity: Entity, ) -> Result<>::Item, QueryEntityError> { // SEMI-SAFE: system runs without conflicts with other systems. diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_state_uniqueness_safety.rs b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_state_uniqueness_safety.rs new file mode 100644 index 0000000000000..1155a9f224874 --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_state_uniqueness_safety.rs @@ -0,0 +1,19 @@ +use bevy_ecs::prelude::*; + +#[derive(Component, Eq, PartialEq, Debug)] +struct A(usize); + +fn main() { + let mut world = World::default(); + world.spawn().insert(A(1)); + + let first_query_state: QueryState<&mut A> = world.query(); + let second_query_state: QueryState<&mut A> = world.query(); + + let mut first_query = Query::from_state(&mut world, &first_query_state); + // This should fail to compile, as another query is already active + let mut second_query = Query::from_state(&mut world, &second_query_state); + + // This is a clear violation of no-aliased mutability + assert_eq!(*first_query.single_mut(), *second_query.single_mut()); +} diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_state_uniqueness_safety.stderr b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_state_uniqueness_safety.stderr new file mode 100644 index 0000000000000..b94eccc6dd07a --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_state_uniqueness_safety.stderr @@ -0,0 +1,11 @@ +error[E0499]: cannot borrow `world` as mutable more than once at a time + --> tests/ui/query_state_uniqueness_safety.rs:15:46 + | +13 | let mut first_query = Query::from_state(&mut world, &first_query_state); + | ---------- first mutable borrow occurs here +14 | // This should fail to compile, as another query is already active +15 | let mut second_query = Query::from_state(&mut world, &second_query_state); + | ^^^^^^^^^^ second mutable borrow occurs here +... +18 | assert_eq!(*first_query.single_mut(), *second_query.single_mut()); + | ------------------------ first borrow later used here