[Merged by Bors] - Add get_multiple and get_multiple_mut APIs for Query and QueryState

crates/bevy_ecs/src/query/state.rs

@@ -153,6 +153,74 @@ 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 mut entities: [Entity; 3] = [Entity::from_raw(0); 3];
+    ///
+    /// world.spawn().insert(A(42));
+    ///
+    /// for i in 0..3 {
+    ///     entities[i] = world.spawn().insert(A(i)).id();
+    /// }
+    ///
+    /// 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<[<Q::ReadOnlyFetch as Fetch<'w, 's>>::Item; N], QueryEntityError> {
+        self.update_archetypes(world);
+
+        let array_of_results = entities.map(|entity| {
+            // SAFETY: query is read only
+            unsafe {
+                self.get_unchecked_manual::<Q::ReadOnlyFetch>(
+                    world,
+                    entity,
+                    world.last_change_tick(),
+                    world.read_change_tick(),
+                )
+            }
+        });
+
+        // If any of the entities were not present, return an error
+        for result in &array_of_results {
+            if let Err(QueryEntityError::NoSuchEntity(entity)) = result {
+                return Err(QueryEntityError::NoSuchEntity(*entity));
+            }
+        }
+
+        // Since we have verified that all entities are present, we can safely unwrap
+        Ok(array_of_results.map(|result| result.unwrap()))
+    }
+
     /// Gets the query result for the given [`World`] and [`Entity`].
     #[inline]
     pub fn get_mut<'w, 's>(
@@ -172,6 +240,85 @@ 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 mut entities: [Entity; 3] = [Entity::from_raw(0); 3];
+    ///
+    /// world.spawn().insert(A(42));
+    ///
+    /// for i in 0..3 {
+    ///     entities[i] = world.spawn().insert(A(i)).id();
+    /// }
+    ///
+    /// 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(365);
+    ///
+    /// assert_eq!(query_state.get_multiple_mut(&mut world, [wrong_entity]), Err(QueryEntityError::NoSuchEntity(wrong_entity)));
+    /// assert_eq!(query_state.get_multiple_mut(&mut world, [entities[0], entities[0]]), 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<[<Q::Fetch as Fetch<'w, 's>>::Item; N], QueryEntityError> {
+        self.update_archetypes(world);
+
+        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| {
+            // SAFETY: entity list is checked for uniqueness, and we require exclusive access to the World
+            unsafe {
+                self.get_unchecked_manual::<Q::Fetch>(
+                    world,
+                    entity,
+                    world.last_change_tick(),
+                    world.read_change_tick(),
+                )
+            }
+        });
+
+        // If any of the entities were not present, return an error
+        for result in &array_of_results {
+            if let Err(QueryEntityError::NoSuchEntity(entity)) = result {
+                return Err(QueryEntityError::NoSuchEntity(*entity));
+            }
+        }
+
+        // Since we have verified that all entities are present, we can safely unwrap
+        Ok(array_of_results.map(|result| result.unwrap()))
+    }
+
     #[inline]
     pub fn get_manual<'w, 's>(
         &'s self,
@@ -228,7 +375,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 +880,13 @@ 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,
     #[error("The requested entity does not exist.")]
-    NoSuchEntity,
+    NoSuchEntity(Entity),
+    #[error("The entity was requested mutably more than once.")]
+    AliasedMutability(Entity),
 }

crates/bevy_ecs/src/system/query.rs

@@ -618,6 +618,85 @@ 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<const N: usize>(
+        &'s self,
+        entities: [Entity; N],
+    ) -> Result<[<Q::ReadOnlyFetch as Fetch<'_, 's>>::Item; N], QueryEntityError> {
+        let array_of_results = entities.map(|entity| {
+            // SAFETY: query is read only
+            unsafe {
+                self.state.get_unchecked_manual::<Q::ReadOnlyFetch>(
+                    self.world,
+                    entity,
+                    self.last_change_tick,
+                    self.change_tick,
+                )
+            }
+        });
+
+        // If any of the entities were not present, return an error
+        for result in &array_of_results {
+            if let Err(QueryEntityError::NoSuchEntity(entity)) = result {
+                return Err(QueryEntityError::NoSuchEntity(*entity));
+            }
+        }
+
+        // Since we have verified that all entities are present, we can safely unwrap
+        Ok(array_of_results.map(|result| result.unwrap()))
+    }
+
+    /// 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<const N: usize>(
+        &'s self,
+        entities: [Entity; N],
+    ) -> [<Q::ReadOnlyFetch as Fetch<'_, 's>>::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 +738,99 @@ 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<const N: usize>(
+        &'s mut self,
+        entities: [Entity; N],
+    ) -> Result<[<Q::Fetch as Fetch>::Item; N], QueryEntityError> {
+        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| {
+            // SAFETY: Entities are checked for uniqueness above,
+            // the scheduler ensure that we do not have conflicting world access,
+            // and we require &mut self to avoid any other simultaneous operations on this Query
+            unsafe {
+                self.state.get_unchecked_manual::<Q::Fetch>(
+                    self.world,
+                    entity,
+                    self.last_change_tick,
+                    self.change_tick,
+                )
+            }
+        });
+
+        // If any of the entities were not present, return an error
+        for result in &array_of_results {
+            if let Err(QueryEntityError::NoSuchEntity(entity)) = result {
+                return Err(QueryEntityError::NoSuchEntity(*entity));
+            }
+        }
+
+        // Since we have verified that all entities are present, we can safely unwrap
+        Ok(array_of_results.map(|result| result.unwrap()))
+    }
+
+    /// 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<const N: usize>(
+        &'s mut self,
+        entities: [Entity; N],
+    ) -> [<Q::Fetch as Fetch<'_, 's>>::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