1 files changed, 26 insertions, 100 deletions

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,
         ))
     }