Repository: Ralith/hecs
Pull request #222 (2021-12-23 09:42:03)

Add query map providing random access to query results
This is basically a generalisation of the `World::column(_mut)` API for arbitrary queries involving multiple components and/or query combinators.

adamreichold: Performance is as expected, i.e.
```
> cargo bench --bench bench --all-features -- access
   Compiling hecs v0.7.1 (/home/adam/Projekte/hecs)
    Finished bench [optimized + debuginfo] target(s) in 35.03s
     Running unittests (target/release/deps/bench-4fb23b93b457eec3)

running 2 tests
test access_column                  ... bench:           7 ns/iter (+/- 0)
test access_query_map               ... bench:           1 ns/iter (+/- 0)

test result: ok. 0 passed; 0 failed; 0 ignored; 2 measured
```

I think `Column(Mut)::get` could be faster if it were to store the base pointers instead of the state?
adamreichold: > I think Column(Mut)::get could be faster if it were to store the base pointers instead of the state?

The last commit brings them on an equal footing:
```
> cargo bench --bench bench --all-features -- access
   Compiling hecs v0.7.1 (/home/adam/Projekte/hecs)
    Finished bench [optimized + debuginfo] target(s) in 34.82s
     Running unittests (target/release/deps/bench-4fb23b93b457eec3)

running 2 tests
test access_column                  ... bench:           1 ns/iter (+/- 0)
test access_query_map               ... bench:           1 ns/iter (+/- 0)

test result: ok. 0 passed; 0 failed; 0 ignored; 2 measured
```

If you'd like I could split that one out into a separate PR as it is independent of the design questions w.r.t. `QueryMap`?
Ralith: If you make such a PR I'll happily merge it, but if we're going to replace `Column(Mut)` with this anyway I'm not sure it's worth the effort.
Ralith: LGTM aside from the outstanding questions about the view type; thanks!

Patch:
diff --git a/benches/bench.rs b/benches/bench.rs
index 997f402..d30a98a 100644
--- a/benches/bench.rs
+++ b/benches/bench.rs
@@ -218,6 +218,21 @@ fn access_column(b: &mut Bencher) {
         let _comp = bencher::black_box(column.get(entc).unwrap());
     });
 }
+
+fn access_view(b: &mut Bencher) {
+    let mut world = World::new();
+    let _enta = world.spawn((Position(0.0), Velocity(0.0)));
+    let _entb = world.spawn((true, 12));
+    let entc = world.spawn((Position(3.0),));
+    let _entd = world.spawn((13, true, 4.0));
+    let mut query = PreparedQuery::<&Position>::new();
+    let mut query = query.query(&world);
+    let view = query.view();
+    b.iter(|| {
+        let _comp = bencher::black_box(view.get(entc).unwrap());
+    });
+}
+
 fn spawn_buffered(b: &mut Bencher) {
     let mut world = World::new();
     let mut buffer = CommandBuffer::new();
@@ -244,6 +259,7 @@ benchmark_group!(
     build,
     build_cloneable,
     access_column,
+    access_view,
     spawn_buffered,
 );
 benchmark_main!(benches);
diff --git a/src/lib.rs b/src/lib.rs
index bbc39d1..624dd25 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -81,7 +81,7 @@ pub use entity_builder::{BuiltEntity, BuiltEntityClone, EntityBuilder, EntityBui
 pub use entity_ref::{EntityRef, Ref, RefMut};
 pub use query::{
     Access, Batch, BatchedIter, Or, PreparedQuery, PreparedQueryBorrow, PreparedQueryIter, Query,
-    QueryBorrow, QueryItem, QueryIter, QueryMut, Satisfies, With, Without,
+    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 5236663..b82658e 100644
--- a/src/query.rs
+++ b/src/query.rs
@@ -11,7 +11,7 @@ use core::mem;
 use core::ptr::NonNull;
 use core::slice::Iter as SliceIter;
 
-use crate::alloc::boxed::Box;
+use crate::alloc::{boxed::Box, vec::Vec};
 use crate::archetype::Archetype;
 use crate::entities::EntityMeta;
 use crate::{Component, Entity, World};
@@ -22,6 +22,10 @@ pub trait Query {
     type Fetch: for<'a> Fetch<'a>;
 }
 
+/// Marker trait indicating whether a given [`Query`] will not produce unique references
+#[allow(clippy::missing_safety_doc)]
+pub unsafe trait QueryShared {}
+
 /// Type of values yielded by a query
 ///
 /// Once rust offers generic associated types, this will be moved into [`Query`].
@@ -80,6 +84,8 @@ impl<'a, T: Component> Query for &'a T {
     type Fetch = FetchRead<T>;
 }
 
+unsafe impl<'a, T> QueryShared for &'a T {}
+
 #[doc(hidden)]
 pub struct FetchRead<T>(NonNull<T>);
 
@@ -173,6 +179,8 @@ impl<T: Query> Query for Option<T> {
     type Fetch = TryFetch<T::Fetch>;
 }
 
+unsafe impl<T: QueryShared> QueryShared for Option<T> {}
+
 #[doc(hidden)]
 pub struct TryFetch<T>(Option<T>);
 
@@ -311,6 +319,8 @@ impl<L: Query, R: Query> Query for Or<L, R> {
     type Fetch = FetchOr<L::Fetch, R::Fetch>;
 }
 
+unsafe impl<L: QueryShared, R: QueryShared> QueryShared for Or<L, R> {}
+
 #[doc(hidden)]
 pub struct FetchOr<L, R>(Or<L, R>);
 
@@ -376,6 +386,8 @@ impl<T: Component, Q: Query> Query for Without<T, Q> {
     type Fetch = FetchWithout<T, Q::Fetch>;
 }
 
+unsafe impl<T, Q: QueryShared> QueryShared for Without<T, Q> {}
+
 #[doc(hidden)]
 pub struct FetchWithout<T, F>(F, PhantomData<fn(T)>);
 
@@ -446,6 +458,8 @@ impl<T: Component, Q: Query> Query for With<T, Q> {
     type Fetch = FetchWith<T, Q::Fetch>;
 }
 
+unsafe impl<T, Q: QueryShared> QueryShared for With<T, Q> {}
+
 #[doc(hidden)]
 pub struct FetchWith<T, F>(F, PhantomData<fn(T)>);
 
@@ -517,6 +531,8 @@ impl<Q: Query> Query for Satisfies<Q> {
     type Fetch = FetchSatisfies<Q::Fetch>;
 }
 
+unsafe impl<Q> QueryShared for Satisfies<Q> {}
+
 #[doc(hidden)]
 pub struct FetchSatisfies<F>(bool, PhantomData<F>);
 
@@ -576,6 +592,12 @@ impl<'w, Q: Query> QueryBorrow<'w, Q> {
         unsafe { QueryIter::new(self.meta, self.archetypes.iter()) }
     }
 
+    /// Provide random access to the query results
+    pub fn view(&mut self) -> View<'_, Q> {
+        self.borrow();
+        unsafe { View::new(self.meta, self.archetypes) }
+    }
+
     /// Like `iter`, but returns child iterators of at most `batch_size` elements
     ///
     /// Useful for distributing work over a threadpool.
@@ -771,6 +793,11 @@ impl<'q, Q: Query> QueryMut<'q, Q> {
         }
     }
 
+    /// Provide random access to the query results
+    pub fn view(&mut self) -> View<'_, Q> {
+        unsafe { View::new(self.iter.meta, self.iter.archetypes.as_slice()) }
+    }
+
     /// Transform the query into one that requires a certain component without borrowing it
     ///
     /// See `QueryBorrow::with`
@@ -1005,6 +1032,8 @@ macro_rules! tuple_impl {
         impl<$($name: Query),*> Query for ($($name,)*) {
             type Fetch = ($($name::Fetch,)*);
         }
+
+        unsafe impl<$($name: QueryShared),*> QueryShared for ($($name,)*) {}
     };
 }
 
@@ -1015,6 +1044,7 @@ smaller_tuples_too!(tuple_impl, O, N, M, L, K, J, I, H, G, F, E, D, C, B, A);
 pub struct PreparedQuery<Q: Query> {
     memo: (u64, u64),
     state: Box<[(usize, <Q::Fetch as Fetch<'static>>::State)]>,
+    fetch: Box<[Option<Q::Fetch>]>,
 }
 
 impl<Q: Query> Default for PreparedQuery<Q> {
@@ -1027,9 +1057,10 @@ impl<Q: Query> PreparedQuery<Q> {
     /// Create a prepared query which is not yet attached to any world
     pub fn new() -> Self {
         Self {
-            // This memo will not match any world as the first ID will 1.
+            // This memo will not match any world as the first ID will be 1.
             memo: (0, 0),
             state: Default::default(),
+            fetch: Default::default(),
         }
     }
 
@@ -1043,7 +1074,9 @@ impl<Q: Query> PreparedQuery<Q> {
             .filter_map(|(idx, x)| Q::Fetch::prepare(x).map(|state| (idx, state)))
             .collect();
 
-        Self { memo, state }
+        let fetch = world.archetypes().map(|_| None).collect();
+
+        Self { memo, state, fetch }
     }
 
     /// Query `world`, using dynamic borrow checking
@@ -1058,7 +1091,7 @@ impl<Q: Query> PreparedQuery<Q> {
         let meta = world.entities_meta();
         let archetypes = world.archetypes_inner();
 
-        PreparedQueryBorrow::new(meta, archetypes, &*self.state)
+        PreparedQueryBorrow::new(meta, archetypes, &*self.state, &mut *self.fetch)
     }
 
     /// Query a uniquely borrowed world
@@ -1079,6 +1112,23 @@ impl<Q: Query> PreparedQuery<Q> {
 
         unsafe { PreparedQueryIter::new(meta, archetypes, state.iter()) }
     }
+
+    /// Provide random access to query results for a uniquely borrow world
+    pub fn view_mut<'q>(&'q mut self, world: &'q mut World) -> PreparedView<'q, Q> {
+        assert_borrow::<Q>();
+
+        if self.memo != world.memo() {
+            *self = Self::prepare(world);
+        }
+
+        let meta = world.entities_meta();
+        let archetypes = world.archetypes_inner();
+
+        let state: &'q [(usize, <Q::Fetch as Fetch<'q>>::State)] =
+            unsafe { mem::transmute(&*self.state) };
+
+        unsafe { PreparedView::new(meta, archetypes, state.iter(), &mut *self.fetch) }
+    }
 }
 
 /// Combined borrow of a [`PreparedQuery`] and a [`World`]
@@ -1086,6 +1136,7 @@ pub struct PreparedQueryBorrow<'q, Q: Query> {
     meta: &'q [EntityMeta],
     archetypes: &'q [Archetype],
     state: &'q [(usize, <Q::Fetch as Fetch<'static>>::State)],
+    fetch: &'q mut [Option<Q::Fetch>],
 }
 
 impl<'q, Q: Query> PreparedQueryBorrow<'q, Q> {
@@ -1093,6 +1144,7 @@ impl<'q, Q: Query> PreparedQueryBorrow<'q, Q> {
         meta: &'q [EntityMeta],
         archetypes: &'q [Archetype],
         state: &'q [(usize, <Q::Fetch as Fetch<'static>>::State)],
+        fetch: &'q mut [Option<Q::Fetch>],
     ) -> Self {
         for (idx, state) in state {
             Q::Fetch::borrow(&archetypes[*idx], *state);
@@ -1102,6 +1154,7 @@ impl<'q, Q: Query> PreparedQueryBorrow<'q, Q> {
             meta,
             archetypes,
             state,
+            fetch,
         }
     }
 
@@ -1113,6 +1166,14 @@ impl<'q, Q: Query> PreparedQueryBorrow<'q, Q> {
 
         unsafe { PreparedQueryIter::new(self.meta, self.archetypes, state.iter()) }
     }
+
+    /// Provides random access to the results of the prepared query
+    pub fn view<'i>(&'i mut self) -> PreparedView<'i, Q> {
+        let state: &'i [(usize, <Q::Fetch as Fetch<'i>>::State)] =
+            unsafe { mem::transmute(self.state) };
+
+        unsafe { PreparedView::new(self.meta, self.archetypes, state.iter(), self.fetch) }
+    }
 }
 
 impl<Q: Query> Drop for PreparedQueryBorrow<'_, Q> {
@@ -1200,6 +1261,123 @@ impl<Q: Query> ExactSizeIterator for PreparedQueryIter<'_, Q> {
     }
 }
 
+/// Provides random access to the results of a query
+pub struct View<'q, Q: Query> {
+    meta: &'q [EntityMeta],
+    fetch: Vec<Option<Q::Fetch>>,
+}
+
+impl<'q, Q: Query> View<'q, Q> {
+    /// # Safety
+    ///
+    /// `'q` must be sufficient to guarantee that `Q` cannot violate borrow safety, either with
+    /// dynamic borrow checks or by representing exclusive access to the `World`.
+    unsafe fn new(meta: &'q [EntityMeta], archetypes: &'q [Archetype]) -> Self {
+        let fetch = archetypes
+            .iter()
+            .map(|archetype| {
+                Q::Fetch::prepare(archetype).map(|state| Q::Fetch::execute(archetype, state))
+            })
+            .collect();
+
+        Self { meta, fetch }
+    }
+
+    /// Retrieve the query results corresponding to `entity`
+    ///
+    /// Will yield `None` if the entity does not exist or does not match the query.
+    ///
+    /// Does not require exclusive access to the map, but is defined only for queries yielding only shared references.
+    pub fn get<'m>(&'m self, entity: Entity) -> Option<QueryItem<'q, Q>>
+    where
+        Q: QueryShared,
+    {
+        let meta = self.meta.get(entity.id as usize)?;
+        if meta.generation != entity.generation {
+            return None;
+        }
+
+        self.fetch[meta.location.archetype as usize]
+            .as_ref()
+            .map(|fetch| unsafe { fetch.get(meta.location.index as usize) })
+    }
+
+    /// Retrieve the query results corresponding to `entity`
+    ///
+    /// Will yield `None` if the entity does not exist or does not match the query.
+    pub fn get_mut<'m>(&'m mut self, entity: Entity) -> Option<QueryItem<'q, Q>> {
+        let meta = self.meta.get(entity.id as usize)?;
+        if meta.generation != entity.generation {
+            return None;
+        }
+
+        self.fetch[meta.location.archetype as usize]
+            .as_ref()
+            .map(|fetch| unsafe { fetch.get(meta.location.index as usize) })
+    }
+}
+
+/// Provides random access to the results of a prepared query
+pub struct PreparedView<'q, Q: Query> {
+    meta: &'q [EntityMeta],
+    fetch: &'q mut [Option<Q::Fetch>],
+}
+
+impl<'q, Q: Query> PreparedView<'q, Q> {
+    /// # Safety
+    ///
+    /// `'q` must be sufficient to guarantee that `Q` cannot violate borrow safety, either with
+    /// dynamic borrow checks or by representing exclusive access to the `World`.
+    unsafe fn new(
+        meta: &'q [EntityMeta],
+        archetypes: &'q [Archetype],
+        state: SliceIter<'q, (usize, <Q::Fetch as Fetch<'q>>::State)>,
+        fetch: &'q mut [Option<Q::Fetch>],
+    ) -> Self {
+        fetch.iter_mut().for_each(|fetch| *fetch = None);
+
+        for (idx, state) in state {
+            let archetype = &archetypes[*idx];
+            fetch[*idx] = Some(Q::Fetch::execute(archetype, *state));
+        }
+
+        Self { meta, fetch }
+    }
+
+    /// Retrieve the query results corresponding to `entity`
+    ///
+    /// Will yield `None` if the entity does not exist or does not match the query.
+    ///
+    /// Does not require exclusive access to the map, but is defined only for queries yielding only shared references.
+    pub fn get<'m>(&'m self, entity: Entity) -> Option<QueryItem<'q, Q>>
+    where
+        Q: QueryShared,
+    {
+        let meta = self.meta.get(entity.id as usize)?;
+        if meta.generation != entity.generation {
+            return None;
+        }
+
+        self.fetch[meta.location.archetype as usize]
+            .as_ref()
+            .map(|fetch| unsafe { fetch.get(meta.location.index as usize) })
+    }
+
+    /// Retrieve the query results corresponding to `entity`
+    ///
+    /// Will yield `None` if the entity does not exist or does not match the query.
+    pub fn get_mut<'m>(&'m mut self, entity: Entity) -> Option<QueryItem<'q, Q>> {
+        let meta = self.meta.get(entity.id as usize)?;
+        if meta.generation != entity.generation {
+            return None;
+        }
+
+        self.fetch[meta.location.archetype as usize]
+            .as_ref()
+            .map(|fetch| unsafe { fetch.get(meta.location.index as usize) })
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;


Test patch:
diff --git a/tests/tests.rs b/tests/tests.rs
index 394b4e0..e9ba980 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -221,6 +221,39 @@ fn invalidate_prepared_query() {
     assert!(ents.contains(&(g, 789, "ghi")));
 }
 
+#[test]
+fn random_access_via_view() {
+    let mut world = World::new();
+    let e = world.spawn(("abc", 123));
+    let f = world.spawn(("def",));
+
+    let mut query = PreparedQuery::<(&i32, &&str)>::default();
+    let mut query = query.query(&world);
+    let mut view = query.view();
+
+    let (i, s) = view.get(e).unwrap();
+    assert_eq!(*i, 123);
+    assert_eq!(*s, "abc");
+
+    assert!(view.get_mut(f).is_none());
+}
+
+#[test]
+fn random_access_via_view_mut() {
+    let mut world = World::new();
+    let e = world.spawn(("abc", 123));
+    let f = world.spawn(("def",));
+
+    let mut query = PreparedQuery::<(&i32, &&str)>::default();
+    let mut view = query.view_mut(&mut world);
+
+    let (i, s) = view.get(e).unwrap();
+    assert_eq!(*i, 123);
+    assert_eq!(*s, "abc");
+
+    assert!(view.get_mut(f).is_none());
+}
+
 #[test]
 fn build_entity() {
     let mut world = World::new();