Query::get_multiple and get_multiple_mut

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<[<Q::ReadOnlyFetch as Fetch<'w, 's>>::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],
+    ) -> [<Q::ReadOnlyFetch as Fetch<'w, 's>>::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::<A>(entity_1).unwrap(), A(11));
+    /// assert_eq!(*world.get::<A>(entity_2).unwrap(), A(2));
+    /// assert_eq!(*world.get::<A>(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<[<Q::Fetch as Fetch<'w, 's>>::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],
+    ) -> [<Q::Fetch as Fetch<'w, 's>>::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),
 }

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<<Q::ReadOnlyFetch as Fetch<'w, 's>>::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<const N: usize>(
+        &self,
+        entities: [Entity; N],
+    ) -> Result<[<Q::ReadOnlyFetch as Fetch<'_, 's>>::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<const N: usize>(
+        &self,
+        entities: [Entity; N],
+    ) -> [<Q::ReadOnlyFetch as Fetch<'_, 's>>::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<const N: usize>(
+        &mut self,
+        entities: [Entity; N],
+    ) -> Result<[<Q::Fetch as Fetch<'_, 's>>::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<const N: usize>(
+        &mut self,
+        entities: [Entity; N],
+    ) -> [<Q::Fetch as Fetch<'_, 's>>::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<<Q::Fetch as Fetch<'w, 's>>::Item, QueryEntityError> {
         // SEMI-SAFE: system runs without conflicts with other systems.

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

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