refactor(ecs): add struct for querying using component metadata

5 files changed, 207 insertions, 117 deletions

diff --git a/ecs/src/component.rs b/ecs/src/component.rs
index 4d37cb8..35e5430 100644
--- a/ecs/src/component.rs
+++ b/ecs/src/component.rs
@@ -180,7 +180,9 @@ pub trait RefSequence
 {
     type Handles<'component>;
 
-    fn metadata() -> impl Array<Metadata>;
+    type Metadata: Array<Metadata>;
+
+    fn metadata() -> Self::Metadata;
 
     fn from_components<'component>(
         components: &'component [EntityComponent],
@@ -330,7 +332,9 @@ macro_rules! inner {
             {
                 type Handles<'component> = (#(CompRef~I::Handle<'component>,)*);
 
-                fn metadata() -> impl Array<Metadata>
+                type Metadata = [Metadata; $c + 1];
+
+                fn metadata() -> Self::Metadata
                 {
                     [#(
                         Metadata {
@@ -397,8 +401,9 @@ impl Sequence for ()
 impl RefSequence for ()
 {
     type Handles<'component> = ();
+    type Metadata = [Metadata; 0];
 
-    fn metadata() -> impl Array<Metadata>
+    fn metadata() -> Self::Metadata
     {
         []
     }
diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index 4b7a6b4..5ce587f 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -14,7 +14,6 @@ use crate::component::{
 };
 use crate::type_name::TypeName;
 use crate::uid::Uid;
-use crate::util::Sortable;
 use crate::EntityComponent;
 
 #[derive(Debug, Default)]
@@ -27,22 +26,21 @@ pub struct Storage
 
 impl Storage
 {
-    pub fn iter_archetypes_with_comps<CompMetadataList>(
+    pub fn iter_archetypes_with_comps(
         &self,
-        mut comp_metadata_list: CompMetadataList,
+        comp_metadata: impl AsRef<[ComponentMetadata]>,
     ) -> ArchetypeRefIter<'_>
-    where
-        CompMetadataList: Sortable<Item = ComponentMetadata>,
-        CompMetadataList: AsRef<[ComponentMetadata]>,
     {
-        comp_metadata_list.sort_by_key_b(|metadata| metadata.id);
+        debug_assert!(comp_metadata
+            .as_ref()
+            .is_sorted_by_key(|metadata| metadata.id));
 
-        let archetype_id = ArchetypeId::from_components_metadata(&comp_metadata_list);
+        let archetype_id = ArchetypeId::from_components_metadata(&comp_metadata);
 
         if !self.archetype_lookup.borrow().contains_key(&archetype_id) {
             self.archetype_lookup.borrow_mut().insert(
                 archetype_id,
-                self.create_populated_archetype_lookup_entry(comp_metadata_list.as_ref()),
+                self.create_populated_archetype_lookup_entry(comp_metadata.as_ref()),
             );
         }
 
diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index d217525..1b9a31b 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -23,6 +23,7 @@ use crate::event::component::Kind as ComponentEventKind;
 use crate::extension::{Collector as ExtensionCollector, Extension};
 use crate::lock::{Lock, WriteGuard};
 use crate::phase::{Phase, START as START_PHASE};
+use crate::query::flexible::Query as FlexibleQuery;
 use crate::query::options::{Not, Options as QueryOptions, With};
 use crate::relationship::{ChildOf, DependsOn, Relationship};
 use crate::sole::Sole;
@@ -30,6 +31,7 @@ use crate::stats::Stats;
 use crate::system::{System, SystemComponent};
 use crate::type_name::TypeName;
 use crate::uid::{Kind as UidKind, Uid};
+use crate::util::Sortable;
 
 pub mod actions;
 pub mod component;
@@ -182,6 +184,16 @@ impl World
         Query::new(self)
     }
 
+    pub fn flexible_query<CompMetadata>(
+        &self,
+        comp_metadata: CompMetadata,
+    ) -> FlexibleQuery<CompMetadata>
+    where
+        CompMetadata: Sortable<Item = ComponentMetadata> + AsRef<[ComponentMetadata]>,
+    {
+        FlexibleQuery::new(self, comp_metadata)
+    }
+
     /// Performs a single tick.
     ///
     /// # Panics
@@ -440,14 +452,18 @@ impl World
 
     fn emit_event_by_id(&self, event_id: Uid)
     {
-        for (system,) in self
-            .query::<(&SystemComponent,), ()>()
-            .iter_with_extra_comps([ComponentMetadata {
+        let query = self.flexible_query([
+            ComponentMetadata::of::<SystemComponent>(),
+            ComponentMetadata {
                 id: event_id,
                 is_optional: ComponentIsOptional::No,
-            }])
+            },
+        ]);
+
+        for (system,) in query
+            .iter::<()>()
+            .into_component_iter::<(&SystemComponent,)>(self)
         {
-            // SAFETY: The world lives long enough
             unsafe {
                 system.system.run(self);
             }
diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index 53d4b74..2f8b285 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -1,23 +1,13 @@
-use std::iter::{repeat_n, Filter, Flatten, Map, RepeatN, Zip};
 use std::marker::PhantomData;
 
-use crate::component::storage::{
-    Archetype,
-    ArchetypeEntity,
-    ArchetypeRefIter,
-    EntityIter,
-    Storage as ComponentStorage,
-};
-use crate::component::{
-    Metadata as ComponentMetadata,
-    RefSequence as ComponentRefSequence,
-};
-use crate::lock::ReadGuard;
+use crate::component::RefSequence as ComponentRefSequence;
+use crate::query::flexible::{Iter as FlexibleQueryIter, Query as FlexibleQuery};
 use crate::query::options::Options;
 use crate::system::{Param as SystemParam, System};
 use crate::uid::Uid;
 use crate::World;
 
+pub mod flexible;
 pub mod options;
 
 #[derive(Debug)]
@@ -26,7 +16,7 @@ where
     Comps: ComponentRefSequence,
 {
     world: &'world World,
-    component_storage: ReadGuard<'world, ComponentStorage>,
+    inner: FlexibleQuery<'world, Comps::Metadata>,
     _pd: PhantomData<(Comps, OptionsT)>,
 }
 
@@ -37,9 +27,7 @@ where
 {
     /// Iterates over the entities matching this query.
     #[must_use]
-    pub fn iter<'query>(
-        &'query self,
-    ) -> ComponentIter<'query, 'world, Comps, QueryEntityIter<'query>>
+    pub fn iter<'query>(&'query self) -> ComponentIter<'query, 'world, Comps>
     {
         #[cfg(feature = "debug")]
         tracing::debug!("Searching for {}", std::any::type_name::<Comps>());
@@ -47,54 +35,7 @@ where
         #[allow(clippy::map_flatten)]
         ComponentIter {
             world: self.world,
-            entities: self
-                .component_storage
-                .iter_archetypes_with_comps(Comps::metadata())
-                .map(
-                    (|archetype| {
-                        repeat_n(archetype, archetype.entity_cnt())
-                            .zip(archetype.entities())
-                    }) as ComponentIterMapFn,
-                )
-                .flatten()
-                .filter(|(_, entity)| OptionsT::entity_filter(entity.components())),
-            comps_pd: PhantomData,
-        }
-    }
-
-    /// Iterates over the entities matching this query and has the provided extra
-    /// component.
-    #[must_use]
-    pub fn iter_with_extra_comps<'query>(
-        &'query self,
-        extra_components: impl IntoIterator<Item = ComponentMetadata>,
-    ) -> ComponentIter<'query, 'world, Comps, QueryEntityIter<'query>>
-    {
-        #[cfg(feature = "debug")]
-        tracing::debug!(
-            "Searching for {} + extra components",
-            std::any::type_name::<Comps>()
-        );
-
-        #[allow(clippy::map_flatten)]
-        ComponentIter {
-            world: self.world,
-            entities: self
-                .component_storage
-                .iter_archetypes_with_comps(
-                    Comps::metadata()
-                        .into_iter()
-                        .chain(extra_components)
-                        .collect::<Vec<_>>(),
-                )
-                .map(
-                    (|archetype| {
-                        repeat_n(archetype, archetype.entity_cnt())
-                            .zip(archetype.entities())
-                    }) as ComponentIterMapFn,
-                )
-                .flatten()
-                .filter(|(_, entity)| OptionsT::entity_filter(entity.components())),
+            iter: self.inner.iter::<OptionsT>(),
             comps_pd: PhantomData,
         }
     }
@@ -103,25 +44,14 @@ where
     #[must_use]
     pub fn get_entity_uid(&self, entity_index: usize) -> Option<Uid>
     {
-        Some(
-            self.component_storage
-                .iter_archetypes_with_comps(Comps::metadata())
-                .flat_map(|archetype| archetype.entities())
-                .filter(|entity| OptionsT::entity_filter(entity.components()))
-                .nth(entity_index)?
-                .uid(),
-        )
+        Some(self.inner.iter::<OptionsT>().nth(entity_index)?.uid())
     }
 
     pub(crate) fn new(world: &'world World) -> Self
     {
         Self {
             world,
-            component_storage: world
-                .data
-                .component_storage
-                .read_nonblock()
-                .expect("Failed to acquire read-only component storage lock"),
+            inner: world.flexible_query(Comps::metadata()),
             _pd: PhantomData,
         }
     }
@@ -133,7 +63,7 @@ where
     Comps: ComponentRefSequence + 'world,
     OptionsT: Options,
 {
-    type IntoIter = ComponentIter<'query, 'world, Comps, QueryEntityIter<'query>>;
+    type IntoIter = ComponentIter<'query, 'world, Comps>;
     type Item = Comps::Handles<'query>;
 
     fn into_iter(self) -> Self::IntoIter
@@ -165,42 +95,38 @@ where
     }
 }
 
-type ComponentIterMapFn =
-    for<'a> fn(&'a Archetype) -> Zip<RepeatN<&'a Archetype>, EntityIter<'a>>;
-
-type ComponentIterFilterFn =
-    for<'a, 'b> fn(&'a (&'b Archetype, &'b ArchetypeEntity)) -> bool;
-
-type QueryEntityIter<'query> = Filter<
-    Flatten<Map<ArchetypeRefIter<'query>, ComponentIterMapFn>>,
-    ComponentIterFilterFn,
->;
-
-pub struct ComponentIter<'query, 'world, Comps, EntityIter>
-where
-    EntityIter: Iterator<Item = (&'query Archetype, &'query ArchetypeEntity)>,
+pub struct ComponentIter<'query, 'world, Comps>
 {
     world: &'world World,
-    entities: EntityIter,
+    iter: FlexibleQueryIter<'query>,
     comps_pd: PhantomData<Comps>,
 }
 
-impl<'query, 'world, Comps, EntityIter> Iterator
-    for ComponentIter<'query, 'world, Comps, EntityIter>
+impl<'query, 'world, Comps> ComponentIter<'query, 'world, Comps>
+where
+    Comps: ComponentRefSequence + 'world,
+    'world: 'query,
+{
+    pub(crate) fn new(world: &'world World, iter: FlexibleQueryIter<'query>) -> Self
+    {
+        Self { world, iter, comps_pd: PhantomData }
+    }
+}
+
+impl<'query, 'world, Comps> Iterator for ComponentIter<'query, 'world, Comps>
 where
     Comps: ComponentRefSequence + 'world,
-    EntityIter: Iterator<Item = (&'query Archetype, &'query ArchetypeEntity)>,
     'world: 'query,
 {
     type Item = Comps::Handles<'query>;
 
     fn next(&mut self) -> Option<Self::Item>
     {
-        let (archetype, entity) = self.entities.next()?;
+        let entity_handle = self.iter.next()?;
 
         Some(Comps::from_components(
-            entity.components(),
-            |component_uid| archetype.get_index_for_component(component_uid),
+            entity_handle.components(),
+            |component_uid| entity_handle.get_component_index(component_uid),
             self.world,
         ))
     }
diff --git a/ecs/src/query/flexible.rs b/ecs/src/query/flexible.rs
new file mode 100644
index 0000000..d081d31
--- /dev/null
+++ b/ecs/src/query/flexible.rs
@@ -0,0 +1,145 @@
+//! Low-level querying.
+use std::iter::{repeat_n, Filter, Flatten, Map, RepeatN, Zip};
+
+use crate::component::storage::{
+    Archetype,
+    ArchetypeEntity,
+    ArchetypeRefIter,
+    EntityIter,
+    Storage as ComponentStorage,
+};
+use crate::component::{
+    Metadata as ComponentMetadata,
+    RefSequence as ComponentRefSequence,
+};
+use crate::lock::ReadGuard;
+use crate::query::options::Options;
+use crate::query::ComponentIter;
+use crate::uid::Uid;
+use crate::util::Sortable;
+use crate::{EntityComponent, World};
+
+/// Low-level entity query structure.
+#[derive(Debug)]
+pub struct Query<'world, CompMetadata>
+where
+    CompMetadata: Sortable<Item = ComponentMetadata> + AsRef<[ComponentMetadata]>,
+{
+    component_storage: ReadGuard<'world, ComponentStorage>,
+    comp_metadata: CompMetadata,
+}
+
+impl<'world, CompMetadata> Query<'world, CompMetadata>
+where
+    CompMetadata: Sortable<Item = ComponentMetadata> + AsRef<[ComponentMetadata]>,
+{
+    /// Iterates over the entities matching this query.
+    #[must_use]
+    pub fn iter<'query, OptionsT: Options>(&'query self) -> Iter<'query>
+    {
+        Iter {
+            iter: self
+                .component_storage
+                .iter_archetypes_with_comps(&self.comp_metadata)
+                .map(
+                    (|archetype| {
+                        repeat_n(archetype, archetype.entity_cnt())
+                            .zip(archetype.entities())
+                    }) as ComponentIterMapFn,
+                )
+                .flatten()
+                .filter(|(_, entity)| OptionsT::entity_filter(entity.components())),
+        }
+    }
+
+    pub(crate) fn new(world: &'world World, mut comp_metadata: CompMetadata) -> Self
+    {
+        comp_metadata.sort_by_key_b(|metadata| metadata.id);
+
+        Self {
+            component_storage: world
+                .data
+                .component_storage
+                .read_nonblock()
+                .expect("Failed to acquire read-only component storage lock"),
+            comp_metadata,
+        }
+    }
+}
+
+pub struct Iter<'query>
+{
+    iter: QueryEntityIter<'query>,
+}
+
+impl<'query> Iter<'query>
+{
+    /// Converts this iterator into a [`ComponentIter`].
+    ///
+    /// Note: The matching entities of this iterator should have all of the non-[`Option`]
+    /// components in `Comps`, otherwise iterating the [`ComponentIter`] will cause a
+    /// panic.
+    #[must_use]
+    #[inline]
+    pub fn into_component_iter<'world, Comps>(
+        self,
+        world: &'world World,
+    ) -> ComponentIter<'query, 'world, Comps>
+    where
+        Comps: ComponentRefSequence + 'world,
+        'world: 'query,
+    {
+        ComponentIter::new(world, self)
+    }
+}
+
+impl<'query> Iterator for Iter<'query>
+{
+    type Item = EntityHandle<'query>;
+
+    fn next(&mut self) -> Option<Self::Item>
+    {
+        let (archetype, entity) = self.iter.next()?;
+
+        Some(EntityHandle { archetype, entity })
+    }
+}
+
+pub struct EntityHandle<'query>
+{
+    archetype: &'query Archetype,
+    entity: &'query ArchetypeEntity,
+}
+
+impl<'query> EntityHandle<'query>
+{
+    /// Returns the [`Uid`] of this entity.
+    #[inline]
+    pub fn uid(&self) -> Uid
+    {
+        self.entity.uid()
+    }
+
+    #[inline]
+    pub fn components(&self) -> &'query [EntityComponent]
+    {
+        self.entity.components()
+    }
+
+    #[inline]
+    pub fn get_component_index(&self, component_uid: Uid) -> Option<usize>
+    {
+        self.archetype.get_index_for_component(component_uid)
+    }
+}
+
+type ComponentIterMapFn =
+    for<'a> fn(&'a Archetype) -> Zip<RepeatN<&'a Archetype>, EntityIter<'a>>;
+
+type ComponentIterFilterFn =
+    for<'a, 'b> fn(&'a (&'b Archetype, &'b ArchetypeEntity)) -> bool;
+
+type QueryEntityIter<'query> = Filter<
+    Flatten<Map<ArchetypeRefIter<'query>, ComponentIterMapFn>>,
+    ComponentIterFilterFn,
+>;