Deduplicate logic

Author: alice-i-cecile

crates/bevy_ecs/src/query/state.rs

@@ -198,27 +198,15 @@ where
     ) -> 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));
-            }
+        // SAFE: query is read-only
+        unsafe {
+            self.get_multiple_unchecked_manual::<Q::ReadOnlyFetch>(
+                world,
+                entities,
+                world.last_change_tick(),
+                world.change_tick(),
+            )
         }
-
-        // 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`].
@@ -288,35 +276,17 @@ where
     ) -> 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]));
-                }
-            }
-        }
+        verify_entities_unique(entities)?;
 
-        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));
-            }
+        // SAFE: method requires exclusive world access, and entities are checked for uniqueness
+        unsafe {
+            self.get_multiple_unchecked_manual::<Q::Fetch>(
+                world,
+                entities,
+                world.last_change_tick(),
+                world.change_tick(),
+            )
         }
-
-        // Since we have verified that all entities are present, we can safely unwrap
-        Ok(array_of_results.map(|result| result.unwrap()))
     }
 
     #[inline]
@@ -396,6 +366,47 @@ where
         }
     }
 
+    /// 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 mutable query correctness. To be safe, make sure mutable queries
+    /// have unique access to the components they query.
+    ///
+    /// If you are calling this method with a mutable query, you must verify that `entities` is free of duplicates:
+    /// use [`verify_entities_unique`] to do so efficiently for small `N`.
+    pub(crate) unsafe fn get_multiple_unchecked_manual<
+        's,
+        'w,
+        QF: Fetch<'w, 's, State = Q::State>,
+        const N: usize,
+    >(
+        &'s self,
+        world: &'w World,
+        entities: [Entity; N],
+        last_change_tick: u32,
+        change_tick: u32,
+    ) -> Result<[<QF as Fetch>::Item; N], QueryEntityError> {
+        let array_of_results = entities.map(|entity| {
+            self.get_unchecked_manual::<QF>(
+                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()))
+    }
+
     /// 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.
@@ -890,3 +901,17 @@ pub enum QueryEntityError {
     #[error("The entity was requested mutably more than once.")]
     AliasedMutability(Entity),
 }
+
+#[inline]
+pub(crate) fn verify_entities_unique<const N: usize>(
+    entities: [Entity; N],
+) -> Result<(), QueryEntityError> {
+    for i in 0..N {
+        for j in 0..i {
+            if entities[i] == entities[j] {
+                return Err(QueryEntityError::AliasedMutability(entities[i]));
+            }
+        }
+    }
+    Ok(())
+}

crates/bevy_ecs/src/system/query.rs

@@ -2,8 +2,8 @@ use crate::{
     component::Component,
     entity::Entity,
     query::{
-        Fetch, FilterFetch, NopFetch, QueryCombinationIter, QueryEntityError, QueryIter,
-        QueryState, ReadOnlyFetch, WorldQuery,
+        verify_entities_unique, Fetch, FilterFetch, NopFetch, QueryCombinationIter,
+        QueryEntityError, QueryIter, QueryState, ReadOnlyFetch, WorldQuery,
     },
     world::{Mut, World},
 };
@@ -631,27 +631,15 @@ where
         &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));
-            }
+        // SAFE: Query is read-only
+        unsafe {
+            self.state.get_multiple_unchecked_manual::<Q::Fetch>(
+                self.world,
+                entities,
+                self.last_change_tick,
+                self.change_tick,
+            )
         }
-
-        // 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`]
@@ -749,37 +737,18 @@ where
         &mut self,
         entities: [Entity; N],
     ) -> Result<[<Q::Fetch as Fetch<'_, 's>>::Item; N], QueryEntityError> {
-        for i in 0..N {
-            for j in 0..i {
-                if entities[i] == entities[j] {
-                    return Err(QueryEntityError::AliasedMutability(entities[i]));
-                }
-            }
-        }
+        verify_entities_unique(entities)?;
 
-        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>(
+        // SAFE: scheduler ensures safe Query world access, and entities are checked for uniqueness
+        unsafe {
+            self.state
+                .get_multiple_unchecked_manual::<<Q::Fetch>::Item>(
                     self.world,
-                    entity,
+                    entities,
                     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`]