Repository: Ralith/hecs
Pull request #362 (2023-12-23 22:35:18)

Introduce World::query_many_mut
CC @sanbox-irl

Bikeshed: It's a bit unfortunate that `View::get_mut_n` is named in a different pattern, but `query_one_mut_n` would be really silly. `query_n_mut` is also a possibility, but is inconsistent with the proposed std [`[T]::get_many_mut`](https://doc.rust-lang.org/std/primitive.slice.html#method.get_many_mut), and if we're going to be internally inconsistent, we might as well take the opportunity to mirror std.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9cd443b..470b589 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,9 @@
 # Unreleased
 
-# Changed
+### Added
+- `World::get_many_mut` to conveniently query a few entities concurrently
+
+### Changed
 - Renamed `{View, PreparedView}::get_mut_n` to `get_many_mut` for consistency with the proposed std
   slice method
 
diff --git a/src/query.rs b/src/query.rs
index 46492b7..1b470b8 100644
--- a/src/query.rs
+++ b/src/query.rs
@@ -902,6 +902,7 @@ impl<'q, Q: Query> IntoIterator for QueryMut<'q, Q> {
     }
 }
 
+/// Check that Q doesn't alias a `&mut T` on its own. Currently over-conservative for `Or` queries.
 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.
@@ -1655,7 +1656,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]),
diff --git a/src/world.rs b/src/world.rs
index e29ce95..8118ab6 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,
@@ -454,7 +454,7 @@ impl World {
         })
     }
 
-    /// Query a single entity in a uniquely borrow world
+    /// Query a single entity in a uniquely borrowed world
     ///
     /// Like [`query_one`](Self::query_one), but faster because dynamic borrow checks can be
     /// skipped. Note that, unlike [`query_one`](Self::query_one), on success this returns the
@@ -472,6 +472,27 @@ impl World {
         unsafe { Ok(Q::get(&fetch, loc.index as usize)) }
     }
 
+    /// Query a fixed number of distinct entities in a uniquely borrowed world
+    ///
+    /// Like [`query_one_mut`](Self::query_one_mut), but for multiple entities, which would
+    /// otherwise be forbidden by the unique borrow. Panics if the same entity occurs more than
+    /// once.
+    pub fn query_many_mut<Q: Query, const N: usize>(
+        &mut self,
+        entities: [Entity; N],
+    ) -> [Result<Q::Item<'_>, QueryOneError>; N] {
+        assert_borrow::<Q>();
+        assert_distinct(&entities);
+
+        entities.map(|entity| {
+            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,


Test patch:
diff --git a/tests/tests.rs b/tests/tests.rs
index 06ad969..bece118 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -969,3 +969,19 @@ fn component_ref_map() {
         assert_eq!(*id, 31);
     }
 }
+
+#[test]
+fn query_many() {
+    let mut world = World::new();
+    let a = world.spawn((42, true));
+    let b = world.spawn((17,));
+    assert_eq!(world.query_many_mut::<&i32, 2>([a, b]), [Ok(&42), Ok(&17)]);
+}
+
+#[test]
+#[should_panic]
+fn query_many_duplicate() {
+    let mut world = World::new();
+    let e = world.spawn(());
+    _ = world.query_many_mut::<(), 2>([e, e]);
+}