Repository: Ralith/hecs (rust)

Issue:
get_mut2/query_one2 (or similar functionality?)
There are a few cases where it would be very useful to be able to reference multiple distinct entities w/ potentially the same components, at the same time. The only way this is currently possible is through the `get_unchecked` methods, or with a query which is later filtered out or collected into an auxiliary structure. Unless I've missed something?

Pull request description:
RFC: Add View::get_mut_n as a safe alterantive to View::get_mut_unchecked.
Fixes #217 

Patch:
diff --git a/src/lib.rs b/src/lib.rs
index 624dd25..c452bad 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -80,8 +80,9 @@ pub use entities::{Entity, NoSuchEntity};
 pub use entity_builder::{BuiltEntity, BuiltEntityClone, EntityBuilder, EntityBuilderClone};
 pub use entity_ref::{EntityRef, Ref, RefMut};
 pub use query::{
-    Access, Batch, BatchedIter, Or, PreparedQuery, PreparedQueryBorrow, PreparedQueryIter, Query,
-    QueryBorrow, QueryItem, QueryIter, QueryMut, QueryShared, Satisfies, View, With, Without,
+    Access, Batch, BatchedIter, Or, PreparedQuery, PreparedQueryBorrow, PreparedQueryIter,
+    PreparedView, Query, QueryBorrow, QueryItem, QueryIter, QueryMut, QueryShared, Satisfies, View,
+    With, Without,
 };
 pub use query_one::QueryOne;
 pub use world::{
diff --git a/src/query.rs b/src/query.rs
index fafa602..1604b02 100644
--- a/src/query.rs
+++ b/src/query.rs
@@ -1324,6 +1324,45 @@ impl<'q, Q: Query> View<'q, Q> {
             .as_ref()
             .map(|fetch| fetch.get(meta.location.index as usize))
     }
+
+    /// Like `get_mut`, but allows checked simultaneous access to multiple entities
+    ///
+    /// For N > 3, the check for distinct entities will clone the array and take O(N log N) time.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use hecs::World;
+    /// let mut world = World::new();
+    ///
+    /// let a = world.spawn((1, 1.0));
+    /// let b = world.spawn((2, 4.0));
+    /// let c = world.spawn((3, 9.0));
+    ///
+    /// let mut query = world.query_mut::<&mut i32>();
+    /// let mut view = query.view();
+    /// let [a,b,c] = view.get_mut_n([a, b, c]);
+    ///
+    /// assert_eq!(*a.unwrap(), 1);
+    /// assert_eq!(*b.unwrap(), 2);
+    /// assert_eq!(*c.unwrap(), 3);
+    /// ```
+    pub fn get_mut_n<const N: usize>(
+        &mut self,
+        entities: [Entity; N],
+    ) -> [Option<QueryItem<'q, Q>>; 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
+    }
 }
 
 /// Provides random access to the results of a prepared query
@@ -1394,8 +1433,46 @@ impl<'q, Q: Query> PreparedView<'q, Q> {
             .as_ref()
             .map(|fetch| fetch.get(meta.location.index as usize))
     }
+
+    /// Like `get_mut`, but allows checked simultaneous access to multiple entities
+    ///
+    /// See [`View::get_mut_n`] for details.
+    pub fn get_mut_n<const N: usize>(
+        &mut self,
+        entities: [Entity; N],
+    ) -> [Option<QueryItem<'q, Q>>; 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
+    }
 }
 
+fn assert_distinct<const N: usize>(entities: &[Entity; N]) {
+    match N {
+        1 => (),
+        2 => assert_ne!(entities[0], entities[1]),
+        3 => {
+            assert_ne!(entities[0], entities[1]);
+            assert_ne!(entities[1], entities[2]);
+            assert_ne!(entities[2], entities[0]);
+        }
+        _ => {
+            let mut entities = *entities;
+            entities.sort_unstable();
+            for index in 0..N - 1 {
+                assert_ne!(entities[index], entities[index + 1]);
+            }
+        }
+    }
+}
 #[cfg(test)]
 mod tests {
     use super::*;


Test patch:
diff --git a/tests/tests.rs b/tests/tests.rs
index b4c2194..d10ef7e 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -255,6 +255,21 @@ fn random_access_via_view_mut() {
     assert!(view.get_mut(f).is_none());
 }
 
+#[test]
+#[should_panic]
+fn simultaneous_access_must_be_non_overlapping() {
+    let mut world = World::new();
+    let a = world.spawn((1,));
+    let b = world.spawn((2,));
+    let c = world.spawn((3,));
+    let d = world.spawn((4,));
+
+    let mut query = world.query_mut::<&mut i32>();
+    let mut view = query.view();
+
+    view.get_mut_n([a, d, c, b, a]);
+}
+
 #[test]
 fn build_entity() {
     let mut world = World::new();