Introduce World::query_many_mut

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
 

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]),

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,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,