diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index f675de4d89a1d..8327d82648ff0 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -153,6 +153,57 @@ where } } + /// Returns the read-only query results for the given array of [`Entity`]. + /// + /// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is + /// returned instead. + /// + /// Note that the unlike [`QueryState::get_multiple_mut`], the entities passed in do not need to be unique. + /// + /// # Examples + /// + /// ```rust + /// use bevy_ecs::prelude::*; + /// use bevy_ecs::query::QueryEntityError; + /// + /// #[derive(Component, PartialEq, Debug)] + /// struct A(usize); + /// + /// let mut world = World::new(); + /// let entity_vec: Vec = (0..3).map(|i|world.spawn().insert(A(i)).id()).collect(); + /// let entities: [Entity; 3] = entity_vec.try_into().unwrap(); + /// + /// world.spawn().insert(A(73)); + /// + /// let mut query_state = world.query::<&A>(); + /// + /// let component_values = query_state.get_multiple(&world, entities).unwrap(); + /// + /// assert_eq!(component_values, [&A(0), &A(1), &A(2)]); + /// + /// let wrong_entity = Entity::from_raw(365); + /// + /// assert_eq!(query_state.get_multiple(&world, [wrong_entity]), Err(QueryEntityError::NoSuchEntity(wrong_entity))); + /// ``` + #[inline] + pub fn get_multiple<'w, 's, const N: usize>( + &'s mut self, + world: &'w World, + entities: [Entity; N], + ) -> Result<[>::Item; N], QueryEntityError> { + self.update_archetypes(world); + + // SAFE: update_archetypes validates the `World` matches + unsafe { + self.get_multiple_read_only_manual( + world, + entities, + world.last_change_tick(), + world.read_change_tick(), + ) + } + } + /// Gets the query result for the given [`World`] and [`Entity`]. #[inline] pub fn get_mut<'w, 's>( @@ -172,6 +223,64 @@ where } } + /// Returns the query results for the given array of [`Entity`]. + /// + /// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is + /// returned instead. + /// + /// ```rust + /// use bevy_ecs::prelude::*; + /// use bevy_ecs::query::QueryEntityError; + /// + /// #[derive(Component, PartialEq, Debug)] + /// struct A(usize); + /// + /// let mut world = World::new(); + /// + /// let entities: Vec = (0..3).map(|i|world.spawn().insert(A(i)).id()).collect(); + /// let entities: [Entity; 3] = entities.try_into().unwrap(); + /// + /// world.spawn().insert(A(73)); + /// + /// let mut query_state = world.query::<&mut A>(); + /// + /// let mut mutable_component_values = query_state.get_multiple_mut(&mut world, entities).unwrap(); + /// + /// for mut a in mutable_component_values.iter_mut(){ + /// a.0 += 5; + /// } + /// + /// let component_values = query_state.get_multiple(&world, entities).unwrap(); + /// + /// assert_eq!(component_values, [&A(5), &A(6), &A(7)]); + /// + /// let wrong_entity = Entity::from_raw(57); + /// let invalid_entity = world.spawn().id(); + /// + /// assert_eq!(query_state.get_multiple_mut(&mut world, [wrong_entity]).unwrap_err(), QueryEntityError::NoSuchEntity(wrong_entity)); + /// assert_eq!(query_state.get_multiple_mut(&mut world, [invalid_entity]).unwrap_err(), QueryEntityError::QueryDoesNotMatch(invalid_entity)); + /// assert_eq!(query_state.get_multiple_mut(&mut world, [entities[0], entities[0]]).unwrap_err(), QueryEntityError::AliasedMutability(entities[0])); + /// ``` + #[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> { + self.update_archetypes(world); + + // SAFE: method requires exclusive world access + // and world has been validated via update_archetypes + unsafe { + self.get_multiple_unchecked_manual( + world, + entities, + world.last_change_tick(), + world.read_change_tick(), + ) + } + } + #[inline] pub fn get_manual<'w, 's>( &'s self, @@ -218,6 +327,9 @@ where /// /// This does not check for mutable query correctness. To be safe, make sure mutable queries /// have unique access to the components they query. + /// + /// This must be called on the same `World` that the `Query` was generated from: + /// use `QueryState::validate_world` to verify this. pub(crate) unsafe fn get_unchecked_manual<'w, 's, QF: Fetch<'w, 's, State = Q::State>>( &'s self, world: &'w World, @@ -228,12 +340,12 @@ 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()) { - return Err(QueryEntityError::QueryDoesNotMatch); + return Err(QueryEntityError::QueryDoesNotMatch(entity)); } let archetype = &world.archetypes[location.archetype_id]; let mut fetch = QF::init(world, &self.fetch_state, last_change_tick, change_tick); @@ -245,10 +357,90 @@ where if filter.archetype_filter_fetch(location.index) { Ok(fetch.archetype_fetch(location.index)) } else { - Err(QueryEntityError::QueryDoesNotMatch) + Err(QueryEntityError::QueryDoesNotMatch(entity)) } } + /// Gets the read-only query results for the given [`World`] and array of [`Entity`], where the last change and + /// the current change tick are given. + /// + /// # Safety + /// + /// This must be called on the same `World` that the `Query` was generated from: + /// use `QueryState::validate_world` to verify this. + pub(crate) unsafe fn get_multiple_read_only_manual<'s, 'w, const N: usize>( + &'s self, + world: &'w World, + entities: [Entity; N], + last_change_tick: u32, + change_tick: u32, + ) -> Result<[>::Item; N], QueryEntityError> { + // SAFE: fetch is read-only + // and world must be validated + let array_of_results = entities.map(|entity| { + self.get_unchecked_manual::( + world, + entity, + last_change_tick, + change_tick, + ) + }); + + // TODO: Replace with TryMap once https://github.com/rust-lang/rust/issues/79711 is stabilized + // If any of the get calls failed, bubble up the error + for result in &array_of_results { + match result { + Ok(_) => (), + Err(error) => return Err(*error), + } + } + + // Since we have verified that all entities are present, we can safely unwrap + Ok(array_of_results.map(|result| result.unwrap())) + } + + /// Gets the query results for the given [`World`] and array of [`Entity`], where the last change and + /// the current change tick are given. + /// + /// # Safety + /// + /// This does not check for unique access to subsets of the entity-component data. + /// To be safe, make sure mutable queries have unique access to the components they query. + /// + /// This must be called on the same `World` that the `Query` was generated from: + /// use `QueryState::validate_world` to verify this. + pub(crate) unsafe fn get_multiple_unchecked_manual<'s, 'w, const N: usize>( + &'s self, + world: &'w World, + entities: [Entity; N], + last_change_tick: u32, + change_tick: u32, + ) -> Result<[>::Item; N], QueryEntityError> { + // Verify that all entities are unique + for i in 0..N { + for j in 0..i { + if entities[i] == entities[j] { + return Err(QueryEntityError::AliasedMutability(entities[i])); + } + } + } + + let array_of_results = entities.map(|entity| { + self.get_unchecked_manual::(world, entity, last_change_tick, change_tick) + }); + + // If any of the get calls failed, bubble up the error + for result in &array_of_results { + match result { + Ok(_) => (), + Err(error) => return Err(*error), + } + } + + // Since we have verified that all entities are present, we can safely unwrap + Ok(array_of_results.map(|result| result.unwrap())) + } + /// Returns an [`Iterator`] over the query results for the given [`World`]. /// /// This can only be called for read-only queries, see [`Self::iter_mut`] for write-queries. @@ -733,10 +925,117 @@ where } /// An error that occurs when retrieving a specific [`Entity`]'s query result. -#[derive(Error, Debug)] +// TODO: return the type_name as part of this error +#[derive(Error, Debug, PartialEq, Clone, Copy)] pub enum QueryEntityError { #[error("The given entity does not have the requested component.")] - QueryDoesNotMatch, + QueryDoesNotMatch(Entity), #[error("The requested entity does not exist.")] - NoSuchEntity, + NoSuchEntity(Entity), + #[error("The entity was requested mutably more than once.")] + AliasedMutability(Entity), +} + +#[cfg(test)] +mod tests { + use crate::{prelude::*, query::QueryEntityError}; + + #[test] + fn get_multiple_unchecked_manual_uniqueness() { + let mut world = World::new(); + + let entities: Vec = (0..10).map(|_| world.spawn().id()).collect(); + + let query_state = world.query::(); + + // These don't matter for the test + let last_change_tick = world.last_change_tick(); + let change_tick = world.read_change_tick(); + + // It's best to test get_multiple_unchecked_manual directly, + // as it is shared and unsafe + // We don't care about aliased mutabilty for the read-only equivalent + assert!(unsafe { + query_state + .get_multiple_unchecked_manual::<10>( + &world, + entities.clone().try_into().unwrap(), + last_change_tick, + change_tick, + ) + .is_ok() + }); + + assert_eq!( + unsafe { + query_state + .get_multiple_unchecked_manual( + &world, + [entities[0], entities[0]], + last_change_tick, + change_tick, + ) + .unwrap_err() + }, + QueryEntityError::AliasedMutability(entities[0]) + ); + + assert_eq!( + unsafe { + query_state + .get_multiple_unchecked_manual( + &world, + [entities[0], entities[1], entities[0]], + last_change_tick, + change_tick, + ) + .unwrap_err() + }, + QueryEntityError::AliasedMutability(entities[0]) + ); + + assert_eq!( + unsafe { + query_state + .get_multiple_unchecked_manual( + &world, + [entities[9], entities[9]], + last_change_tick, + change_tick, + ) + .unwrap_err() + }, + QueryEntityError::AliasedMutability(entities[9]) + ); + } + + #[test] + #[should_panic] + fn right_world_get() { + let mut world_1 = World::new(); + let world_2 = World::new(); + + let mut query_state = world_1.query::(); + let _panics = query_state.get(&world_2, Entity::from_raw(0)); + } + + #[test] + #[should_panic] + fn right_world_get_multiple() { + let mut world_1 = World::new(); + let world_2 = World::new(); + + let mut query_state = world_1.query::(); + let _panics = query_state.get_multiple(&world_2, []); + } + + #[test] + #[should_panic] + fn right_world_get_multiple_mut() { + let mut world_1 = World::new(); + let mut world_2 = World::new(); + + let mut query_state = world_1.query::(); + let _panics = query_state.get_multiple_mut(&mut world_2, []); + } } diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index 349121b06bd87..7fcaf72a0ac57 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -618,6 +618,73 @@ where } } + /// Returns the read-only query results for the given array of [`Entity`]. + /// + /// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is + /// returned instead. + /// + /// Note that the unlike [`Query::get_multiple_mut`], the entities passed in do not need to be unique. + /// + /// See [`Query::multiple`] for the infallible equivalent. + #[inline] + pub fn get_multiple( + &self, + entities: [Entity; N], + ) -> Result<[>::Item; N], QueryEntityError> { + // SAFE: it is the scheduler's responsibility to ensure that `Query` is never handed out on the wrong `World`. + unsafe { + self.state.get_multiple_read_only_manual( + self.world, + entities, + self.last_change_tick, + self.change_tick, + ) + } + } + + /// Returns the read-only query items for the provided array of [`Entity`] + /// + /// See [`Query::get_multiple`] for the [`Result`]-returning equivalent. + /// + /// # Examples + /// ```rust, no_run + /// use bevy_ecs::prelude::*; + /// + /// #[derive(Component)] + /// struct Targets([Entity; 3]); + /// + /// #[derive(Component)] + /// struct Position{ + /// x: i8, + /// y: i8 + /// }; + /// + /// impl Position { + /// fn distance(&self, other: &Position) -> i8 { + /// // Manhattan distance is way easier to compute! + /// (self.x - other.x).abs() + (self.y - other.y).abs() + /// } + /// } + /// + /// fn check_all_targets_in_range(targeting_query: Query<(Entity, &Targets, &Position)>, targets_query: Query<&Position>){ + /// for (targeting_entity, targets, origin) in targeting_query.iter(){ + /// // We can use "destructuring" to unpack the results nicely + /// let [target_1, target_2, target_3] = targets_query.multiple(targets.0); + /// + /// assert!(target_1.distance(origin) <= 5); + /// assert!(target_2.distance(origin) <= 5); + /// assert!(target_3.distance(origin) <= 5); + /// } + /// } + /// ``` + #[inline] + pub fn multiple( + &self, + entities: [Entity; N], + ) -> [>::Item; N] { + self.get_multiple(entities).unwrap() + } + /// Returns the query result for the given [`Entity`]. /// /// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is @@ -659,6 +726,77 @@ where } } + /// Returns the query results for the given array of [`Entity`]. + /// + /// In case of a nonexisting entity, duplicate entities or mismatched component, a [`QueryEntityError`] is + /// returned instead. + /// + /// See [`Query::multiple_mut`] for the infallible equivalent. + #[inline] + pub fn get_multiple_mut( + &mut self, + entities: [Entity; N], + ) -> Result<[>::Item; N], QueryEntityError> { + // SAFE: scheduler ensures safe Query world access + unsafe { + self.state.get_multiple_unchecked_manual( + self.world, + entities, + self.last_change_tick, + self.change_tick, + ) + } + } + + /// Returns the query items for the provided array of [`Entity`] + /// + /// See [`Query::get_multiple_mut`] for the [`Result`]-returning equivalent. + /// + /// # Examples + /// + /// ```rust, no_run + /// use bevy_ecs::prelude::*; + /// + /// #[derive(Component)] + /// struct Spring{ + /// connected_entities: [Entity; 2], + /// strength: f32, + /// } + /// + /// #[derive(Component)] + /// struct Position { + /// x: f32, + /// y: f32, + /// } + /// + /// #[derive(Component)] + /// struct Force { + /// x: f32, + /// y: f32, + /// } + /// + /// fn spring_forces(spring_query: Query<&Spring>, mut mass_query: Query<(&Position, &mut Force)>){ + /// for spring in spring_query.iter(){ + /// // We can use "destructuring" to unpack our query items nicely + /// let [(position_1, mut force_1), (position_2, mut force_2)] = mass_query.multiple_mut(spring.connected_entities); + /// + /// force_1.x += spring.strength * (position_1.x - position_2.x); + /// force_1.y += spring.strength * (position_1.y - position_2.y); + /// + /// // Silence borrow-checker: I have split your mutable borrow! + /// force_2.x += spring.strength * (position_2.x - position_1.x); + /// force_2.y += spring.strength * (position_2.y - position_1.y); + /// } + /// } + /// ``` + #[inline] + pub fn multiple_mut( + &mut self, + entities: [Entity; N], + ) -> [>::Item; N] { + self.get_multiple_mut(entities).unwrap() + } + /// Returns the query result for the given [`Entity`]. /// /// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_lifetime_safety.rs b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_lifetime_safety.rs new file mode 100644 index 0000000000000..d50bf4da82d24 --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_lifetime_safety.rs @@ -0,0 +1,13 @@ +use bevy_ecs::prelude::*; + +#[derive(Component)] +struct A(usize); + +fn system(mut query: Query<&mut A>, e: Res) { + let a1 = query.get_multiple([*e, *e]).unwrap(); + let a2 = query.get_mut(*e).unwrap(); + // this should fail to compile + println!("{} {}", a1[0].0, a2.0); +} + +fn main() {} diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_lifetime_safety.stderr b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_lifetime_safety.stderr new file mode 100644 index 0000000000000..229c88181040f --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_lifetime_safety.stderr @@ -0,0 +1,10 @@ +error[E0502]: cannot borrow `query` as mutable because it is also borrowed as immutable + --> tests/ui/system_query_get_multiple_lifetime_safety.rs:8:14 + | +7 | let a1 = query.get_multiple([*e, *e]).unwrap(); + | ---------------------------- immutable borrow occurs here +8 | let a2 = query.get_mut(*e).unwrap(); + | ^^^^^^^^^^^^^^^^^ mutable borrow occurs here +9 | // this should fail to compile +10 | println!("{} {}", a1[0].0, a2.0); + | ----- immutable borrow later used here diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_mut_lifetime_safety.rs b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_mut_lifetime_safety.rs new file mode 100644 index 0000000000000..4a1219f4e8df6 --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_mut_lifetime_safety.rs @@ -0,0 +1,13 @@ +use bevy_ecs::prelude::*; + +#[derive(Component)] +struct A(usize); + +fn system(mut query: Query<&mut A>, e: Res) { + let a1 = query.get_multiple_mut([*e, *e]).unwrap(); + let a2 = query.get_mut(*e).unwrap(); + // this should fail to compile + println!("{} {}", a1[0].0, a2.0); +} + +fn main() {} diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_mut_lifetime_safety.stderr b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_mut_lifetime_safety.stderr new file mode 100644 index 0000000000000..4ef0380eeeae4 --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/system_query_get_multiple_mut_lifetime_safety.stderr @@ -0,0 +1,10 @@ +error[E0499]: cannot borrow `query` as mutable more than once at a time + --> tests/ui/system_query_get_multiple_mut_lifetime_safety.rs:8:14 + | +7 | let a1 = query.get_multiple_mut([*e, *e]).unwrap(); + | -------------------------------- first mutable borrow occurs here +8 | let a2 = query.get_mut(*e).unwrap(); + | ^^^^^^^^^^^^^^^^^ second mutable borrow occurs here +9 | // this should fail to compile +10 | println!("{} {}", a1[0].0, a2.0); + | ----- first borrow later used here