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

1 files changed, 145 insertions, 0 deletions

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