ecs: Working implementation of static queries

It took some wrestling with lifetimes, but static queries now work. They have roughly the same
interface as `hecs`.

Author: caelunshun

quill/ecs/src/ecs.rs

@@ -1,4 +1,4 @@
-use std::{any::type_name, iter};
+use std::{any::type_name, borrow::Cow, iter, marker::PhantomData};
 
 use ahash::AHashMap;
 use itertools::Either;
@@ -8,6 +8,7 @@ use crate::{
     component::{Component, ComponentMeta},
     entity::{Entities, EntityId},
     entity_builder::EntityBuilder,
+    query::{QueryDriverIter, QueryTuple},
     storage::SparseSetStorage,
     QueryDriver,
 };
@@ -140,36 +141,75 @@ impl Ecs {
         Ok(())
     }
 
-    /*
     /// Queries for all entities that have the given set of components.
     ///
     /// Returns an iterator over tuples of `(entity, components)`.
-    pub fn query<'a, Q: QueryTuple>(
-        &'a self,
-    ) -> impl Iterator<Item = (EntityId, Q::Output<'a>)> + 'a
-    where
-        Q::Output<'a>: 'a,
-    {
-        let sparse_sets = match Q::sparse_sets(&self.components) {
-            Some(s) => s,
-            None => return Either::Left(iter::empty()),
-        };
-        let sparse_set_refs: Vec<_> = sparse_sets.iter().map(|s| s.to_ref()).collect();
+    pub fn query<'w, 'q, Q: QueryTuple<'w>>(&'w self) -> Query<'w, 'q, Q> {
+        let sparse_sets = Q::sparse_sets(&self.components).unwrap_or_else(|| todo!());
+        let sparse_set_refs: Vec<_> = sparse_sets.iter().map(|set| set.to_ref()).collect();
         let dense_indices = Q::dense_indices();
 
-        let driver = QueryDriver::new(&sparse_set_refs, &dense_indices);
+        let driver = QueryDriver::new(Cow::Owned(sparse_set_refs), Cow::Owned(dense_indices));
 
-        Either::Right(driver.iter().map(move |item| {
-            let components = unsafe { Q::make_output(&sparse_sets, item.dense_indices) };
-            let entity = self.entities.get(item.sparse_index);
-            (entity, components)
-        }))
+        Query {
+            driver,
+            sparse_sets,
+            entities: &self.entities,
+            _marker: PhantomData,
+        }
     }
-    */
 
     fn check_entity(&self, entity: EntityId) -> Result<(), EntityDead> {
         self.entities
             .check_generation(entity)
             .map_err(|_| EntityDead)
     }
 }
+
+/// An iterator over a statically-typed query.
+///
+/// Call [`iter`] to iterate over the items.
+pub struct Query<'w, 'q, Q> {
+    driver: QueryDriver<'w, 'q>,
+    sparse_sets: Vec<&'w SparseSetStorage>,
+    entities: &'w Entities,
+    _marker: PhantomData<Q>,
+}
+
+impl<'w, 'q, Q> Query<'w, 'q, Q>
+where
+    Q: QueryTuple<'w>,
+{
+    pub fn iter(&'q mut self) -> QueryIter<'w, 'q, Q> {
+        QueryIter {
+            driver: self.driver.iter(),
+            sparse_sets: &self.sparse_sets,
+            entities: self.entities,
+            _marker: self._marker,
+        }
+    }
+}
+
+pub struct QueryIter<'w, 'q, Q> {
+    driver: QueryDriverIter<'w, 'q>,
+    sparse_sets: &'q [&'w SparseSetStorage],
+    entities: &'w Entities,
+    _marker: PhantomData<Q>,
+}
+
+impl<'w, 'q, Q> Iterator for QueryIter<'w, 'q, Q>
+where
+    Q: QueryTuple<'w>,
+{
+    type Item = (EntityId, Q::Output);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let item = self.driver.next()?;
+
+        let components = unsafe { Q::make_output(self.sparse_sets, item.dense_indices) };
+
+        let entity = self.entities.get(item.sparse_index);
+
+        Some((entity, components))
+    }
+}

quill/ecs/src/query.rs

@@ -1,22 +1,22 @@
 //! Dynamic query infrastructure.
 
-use std::{any::TypeId, cell::Cell, ops::Deref};
+use std::{any::TypeId, borrow::Cow, cell::Cell, ops::Deref};
 
 use crate::{storage::sparse_set, Component, Components, Ecs, SparseSetRef, SparseSetStorage};
 
 /// Drives a query by yielding the entities
 /// whose components satisfy the query parameters.
 pub struct QueryDriver<'w, 'q> {
     /// A sparse set for each component in the query.
-    sparse_sets: &'q [SparseSetRef<'w>],
+    sparse_sets: Cow<'q, [SparseSetRef<'w>]>,
 
     /// The "lead" sparse set, chosen as the set with
     /// the smallest number of components.
-    lead: SparseSetRef<'q>,
+    lead: SparseSetRef<'w>,
 
     /// Used as the yielded value for the iterator.
     /// (We can't own this because of the lack of GATs.)
-    dense_indices: &'q [Cell<u32>],
+    dense_indices: Cow<'q, [Cell<u32>]>,
 }
 
 impl<'w, 'q> QueryDriver<'w, 'q> {
@@ -27,15 +27,19 @@ impl<'w, 'q> QueryDriver<'w, 'q> {
     ///
     /// # Panics
     /// Panics if `sparse_sets.len() != dense_indices.len()`.
-    pub fn new(sparse_sets: &'q [SparseSetRef<'w>], dense_indices: &'q [Cell<u32>]) -> Self {
+    pub fn new(
+        sparse_sets: Cow<'q, [SparseSetRef<'w>]>,
+        dense_indices: Cow<'q, [Cell<u32>]>,
+    ) -> Self {
         let lead = sparse_sets
             .iter()
             .min_by_key(|set| set.len())
-            .unwrap_or(SparseSetRef::empty());
+            .copied()
+            .unwrap_or(*SparseSetRef::empty());
 
         Self {
             sparse_sets,
-            lead: *lead,
+            lead,
             dense_indices,
         }
     }
@@ -60,10 +64,8 @@ pub struct QueryDriverIter<'w, 'q> {
     lead_iter: sparse_set::Iter<'q>,
 }
 
-impl<'w, 'q> Iterator for QueryDriverIter<'w, 'q> {
-    type Item = QueryItem<'q>;
-
-    fn next(&mut self) -> Option<Self::Item> {
+impl<'w, 'q> QueryDriverIter<'w, 'q> {
+    pub fn next(&mut self) -> Option<QueryItem> {
         loop {
             let (sparse_index, lead_dense_index) = self.lead_iter.next()?;
 
@@ -98,36 +100,35 @@ pub struct QueryItem<'q> {
 
 // -- Static queries
 
-/*
 /// A typed query element.
-pub trait QueryParameter {
-    type Output;
+pub trait QueryParameter<'a> {
+    type Output: 'a;
     type Component: Component;
 
     unsafe fn get_unchecked_by_dense_index(
-        storage: &SparseSetStorage,
+        storage: &'a SparseSetStorage,
         dense_index: u32,
     ) -> Self::Output;
 }
 
-impl<'a, T> QueryParameter for &'a T
+impl<'a, T> QueryParameter<'a> for &'a T
 where
     T: Component,
 {
-    type Output<'b> = &'b T;
+    type Output = &'a T;
     type Component = T;
 
     unsafe fn get_unchecked_by_dense_index(
-        storage: &SparseSetStorage,
+        storage: &'a SparseSetStorage,
         dense_index: u32,
-    ) -> Self::Output<'_> {
+    ) -> Self::Output {
         storage.get_unchecked_by_dense_index(dense_index)
     }
 }
 
 /// A tuple of query parameters.
-pub trait QueryTuple {
-    type Output<'s>;
+pub trait QueryTuple<'a> {
+    type Output: 'a;
 
     // avoiding allocations here is blocked on const generics and/or GATs
     fn sparse_sets(components: &Components) -> Option<Vec<&SparseSetStorage>>;
@@ -141,16 +142,16 @@ pub trait QueryTuple {
     /// `dense_indices` and `sparse_sets` must have a length equal
     /// to the length of the vectors returned by the corresponding methods
     /// of this trait.
-    unsafe fn make_output<'s>(
-        sparse_sets: &'s [&'s SparseSetStorage],
+    unsafe fn make_output(
+        sparse_sets: &[&'a SparseSetStorage],
         dense_indices: &[Cell<u32>],
-    ) -> Self::Output<'s>;
+    ) -> Self::Output;
 }
 
 macro_rules! query_tuple_impl {
     ($count:literal, $(($ty:ident, $index:literal)),* $(,)?) => {
-        impl <$($ty: QueryParameter),*> QueryTuple for ($($ty),*) {
-            type Output<'s> = ($($ty::Output<'s>),*);
+        impl <'a, $($ty: QueryParameter<'a>),*> QueryTuple<'a> for ($($ty),*) {
+            type Output = ($($ty::Output),*);
 
             fn sparse_sets(components: &Components) -> Option<Vec<&SparseSetStorage>> {
                 Some(vec![
@@ -164,10 +165,10 @@ macro_rules! query_tuple_impl {
                 vec![Cell::new(0); $count]
             }
 
-            unsafe fn make_output<'s>(
-                sparse_sets: &'s [&'s SparseSetStorage],
+            unsafe fn make_output(
+                sparse_sets: &[&'a SparseSetStorage],
                 dense_indices: &[Cell<u32>],
-            ) -> Self::Output<'s> {
+            ) -> Self::Output {
                 (
                     $(
                         $ty::get_unchecked_by_dense_index(sparse_sets.get_unchecked($index), dense_indices.get_unchecked($index).get())
@@ -179,4 +180,4 @@ macro_rules! query_tuple_impl {
 }
 
 query_tuple_impl!(1, (T1, 0));
-*/
+query_tuple_impl!(2, (T1, 0), (T2, 1));

quill/ecs/tests/ecs.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use quill_ecs::*;
 
 #[test]
@@ -83,3 +85,48 @@ fn remove_nonexisting() {
     let entity = ecs.spawn_bundle((10i32,));
     assert!(ecs.remove::<usize>(entity).is_err());
 }
+
+#[test]
+fn query_basic() {
+    let mut ecs = Ecs::new();
+
+    let entity1 = ecs.spawn_bundle((10i32, "name1"));
+    let entity2 = ecs.spawn_bundle((15i32, "name2", 50.0f32));
+
+    let mut query = ecs.query::<(&i32, &&'static str)>();
+    let mut iter = query.iter();
+
+    assert_eq!(iter.next(), Some((entity1, (&10, &"name1"))));
+    assert_eq!(iter.next(), Some((entity2, (&15, &"name2"))));
+    assert_eq!(iter.next(), None);
+}
+
+#[test]
+fn query_big_ecs_after_despawn() {
+    let mut ecs = Ecs::new();
+
+    let mut entities = Vec::new();
+    for i in 0..10_000usize {
+        let mut builder = EntityBuilder::new();
+        if i % 3 == 0 {
+            builder.add(format!("entity #{}", i));
+        }
+        builder.add(i);
+        let entity = builder.spawn_into(&mut ecs);
+        if i % 3 == 0 {
+            entities.push(entity);
+        }
+    }
+
+    let last = entities.len() - 1;
+    ecs.despawn(entities.remove(last)).unwrap();
+
+    let queried: HashMap<EntityId, (&String, &usize)> =
+        ecs.query::<(&String, &usize)>().iter().collect();
+
+    for (i, entity) in entities.iter().copied().enumerate() {
+        assert_eq!(queried[&entity], (&format!("entity #{}", i * 3), &(i * 3)));
+    }
+
+    assert_eq!(queried.len(), entities.len());
+}