From 3f075f91c532e369435a06b7670f9fb5ab5f35cb Mon Sep 17 00:00:00 2001
From: Jonathan Spira <jjspira@gmail.com>
Date: Sat, 23 Dec 2023 12:08:02 -0800
Subject: [PATCH] query-n-mut

---
 src/query.rs | 57 +++++++++++++++++++++++++---------------------------
 src/world.rs | 21 ++++++++++++++++++-
 2 files changed, 47 insertions(+), 31 deletions(-)

diff --git a/src/query.rs b/src/query.rs
index d46fca1c..6b852840 100644
--- a/src/query.rs
+++ b/src/query.rs
@@ -902,24 +902,6 @@ impl<'q, Q: Query> IntoIterator for QueryMut<'q, Q> {
     }
 }
 
-pub(crate) fn assert_borrow<Q: Query>() {
-    // This looks like an ugly O(n^2) loop, but everything's constant after inlining, so in
-    // practice LLVM optimizes it out entirely.
-    let mut i = 0;
-    Q::Fetch::for_each_borrow(|a, unique| {
-        if unique {
-            let mut j = 0;
-            Q::Fetch::for_each_borrow(|b, _| {
-                if i != j {
-                    core::assert!(a != b, "query violates a unique borrow");
-                }
-                j += 1;
-            })
-        }
-        i += 1;
-    });
-}
-
 struct ChunkIter<Q: Query> {
     entities: NonNull<u32>,
     fetch: Q::Fetch,
@@ -1414,9 +1396,13 @@ impl<'q, Q: Query> View<'q, Q> {
             .map(|fetch| Q::get(fetch, meta.location.index as usize))
     }
 
-    /// Like `get_mut`, but allows checked simultaneous access to multiple entities
+    /// Like `get_mut`, but allows checked simultaneous access to multiple entities.
+    /// Note: passing the same entity twice will result in a panic.
+    ///
+    /// # Panics
     ///
-    /// For N > 3, the check for distinct entities will clone the array and take O(N log N) time.
+    /// Passing two or more of the same entities will result in a panic. Additionally,
+    /// for `N > 3`, the check for distinct entities will copy the array and take O(N log N) time.
     ///
     /// # Examples
     ///
@@ -1439,15 +1425,7 @@ impl<'q, Q: Query> View<'q, Q> {
     pub fn get_mut_n<const N: usize>(&mut self, entities: [Entity; N]) -> [Option<Q::Item<'_>>; N] {
         assert_distinct(&entities);
 
-        let mut items = [(); N].map(|()| None);
-
-        for (item, entity) in items.iter_mut().zip(entities) {
-            unsafe {
-                *item = self.get_unchecked(entity);
-            }
-        }
-
-        items
+        core::array::from_fn(|idx| unsafe { self.get_unchecked(entities[idx]) })
     }
 
     /// Iterate over all entities satisfying `Q`
@@ -1637,7 +1615,7 @@ impl<'a, 'q, Q: Query> IntoIterator for &'a mut PreparedView<'q, Q> {
     }
 }
 
-fn assert_distinct<const N: usize>(entities: &[Entity; N]) {
+pub(crate) fn assert_distinct<const N: usize>(entities: &[Entity; N]) {
     match N {
         1 => (),
         2 => assert_ne!(entities[0], entities[1]),
@@ -1655,6 +1633,25 @@ fn assert_distinct<const N: usize>(entities: &[Entity; N]) {
         }
     }
 }
+
+pub(crate) fn assert_borrow<Q: Query>() {
+    // This looks like an ugly O(n^2) loop, but everything's constant after inlining, so in
+    // practice LLVM optimizes it out entirely.
+    let mut i = 0;
+    Q::Fetch::for_each_borrow(|a, unique| {
+        if unique {
+            let mut j = 0;
+            Q::Fetch::for_each_borrow(|b, _| {
+                if i != j {
+                    core::assert!(a != b, "query violates a unique borrow");
+                }
+                j += 1;
+            })
+        }
+        i += 1;
+    });
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/world.rs b/src/world.rs
index e29ce95d..5710a6f2 100644
--- a/src/world.rs
+++ b/src/world.rs
@@ -22,7 +22,7 @@ use hashbrown::hash_map::{Entry, HashMap};
 use crate::alloc::boxed::Box;
 use crate::archetype::{Archetype, TypeIdMap, TypeInfo};
 use crate::entities::{Entities, EntityMeta, Location, ReserveEntitiesIterator};
-use crate::query::assert_borrow;
+use crate::query::{assert_borrow, assert_distinct};
 use crate::{
     Bundle, ColumnBatch, ComponentRef, DynamicBundle, Entity, EntityRef, Fetch, MissingComponent,
     NoSuchEntity, Query, QueryBorrow, QueryMut, QueryOne, TakenEntity,
@@ -472,6 +472,25 @@ impl World {
         unsafe { Ok(Q::get(&fetch, loc.index as usize)) }
     }
 
+    /// Query multiple entities in a uniquely borrow world
+    pub fn query_n_mut<Q: Query, const N: usize>(
+        &mut self,
+        entities: [Entity; N],
+    ) -> [Result<Q::Item<'_>, QueryOneError>; N] {
+        assert_borrow::<Q>();
+        assert_distinct(&entities);
+
+        core::array::from_fn(|i| {
+            let entity = entities[i];
+
+            let loc = self.entities.get(entity)?;
+            let archetype = &self.archetypes.archetypes[loc.archetype as usize];
+            let state = Q::Fetch::prepare(archetype).ok_or(QueryOneError::Unsatisfied)?;
+            let fetch = Q::Fetch::execute(archetype, state);
+            unsafe { Ok(Q::get(&fetch, loc.index as usize)) }
+        })
+    }
+
     /// Short-hand for [`entity`](Self::entity) followed by [`EntityRef::get`]
     pub fn get<'a, T: ComponentRef<'a>>(
         &'a self,