Add get_multiple and get_multiple_mut APIs for Query and QueryState (#4298)

# Objective

- The inability to have multiple active mutable borrows into a query is a common source of borrow-checker pain for users.
- This is a pointless restriction if and only if we can guarantee that the entities they are accessing are unique.
- This could already by bypassed with get_unchecked, but that is an extremely unsafe API.
- Closes https://github.com/bevyengine/bevy/issues/2042.

## Solution

- Add `get_multiple`, `get_multiple_mut` and their unchecked equivalents (`multiple` and `multiple_mut`) to `Query` and `QueryState`.
- Improve the `QueryEntityError` type to provide more useful error information.

## Changelog

- Added `get_multiple`, `get_multiple_mut` and their unchecked equivalents (`multiple` and `multiple_mut`) to Query and QueryState.

## Migration Guide

- The `QueryEntityError` enum now has a `AliasedMutability variant, and returns the offending entity.

## Context

This is a fresh attempt at #3333; rebasing was behaving very badly and it was important to rebase on top of the recent query soundness fixes. Many thanks to all the reviewers in that thread, especially @BoxyUwU for the help with lifetimes.

## To-do

- [x] Add compile fail tests
- [x] Successfully deduplicate code
- [x] Decide what to do about failing doc tests
- [x] Get some reviews for lifetime soundness

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<Entity> = (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<[<Q::ReadOnlyFetch as Fetch<'w, 's>>::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<Entity> = (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<[<Q::Fetch as Fetch<'w, 's>>::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<[<Q::ReadOnlyFetch as Fetch<'w, 's>>::Item; N], QueryEntityError> {
+        // SAFE: fetch is read-only
+        // and world must be validated
+        let array_of_results = entities.map(|entity| {
+            self.get_unchecked_manual::<Q::ReadOnlyFetch>(
+                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<[<Q::Fetch as Fetch<'w, 's>>::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::<Q::Fetch>(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<Entity> = (0..10).map(|_| world.spawn().id()).collect();
+
+        let query_state = world.query::<Entity>();
+
+        // 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::<Entity>();
+        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::<Entity>();
+        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::<Entity>();
+        let _panics = query_state.get_multiple_mut(&mut world_2, []);
+    }
 }

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<const N: usize>(
+        &self,
+        entities: [Entity; N],
+    ) -> Result<[<Q::ReadOnlyFetch as Fetch<'_, 's>>::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<const N: usize>(
+        &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 +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<const N: usize>(
+        &mut self,
+        entities: [Entity; N],
+    ) -> Result<[<Q::Fetch as Fetch<'_, 's>>::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<const N: usize>(
+        &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

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<Entity>) {
+    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() {}

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

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<Entity>) {
+    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() {}

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