1 files changed, 48 insertions, 39 deletions

diff --git a/engine-ecs/src/query/flexible.rs b/engine-ecs/src/query/flexible.rs
index fd1e30e..0198fb4 100644
--- a/engine-ecs/src/query/flexible.rs
+++ b/engine-ecs/src/query/flexible.rs
@@ -1,41 +1,34 @@
 //! Low-level querying.
-use std::iter::{repeat_n, FlatMap, RepeatN, Zip};
-
 use crate::component::storage::archetype::{Archetype, EntityIter};
-use crate::component::storage::{ArchetypeRefIter, ArchetypeSearchTerms};
+use crate::component::storage::{ArchetypeRefIter, TraversalResult};
 use crate::entity::Handle as EntityHandle;
 use crate::query::Terms;
+use crate::util::array_vec::ArrayVec;
 use crate::World;
 
 /// Low-level entity query structure.
 #[derive(Debug)]
-pub struct Query<'world, const MAX_TERM_CNT: usize>
+pub struct Query<'world, const TERM_CAP: usize>
 {
     world: &'world World,
-    terms: Terms<MAX_TERM_CNT>,
+    terms: Terms<TERM_CAP>,
 }
 
-impl<'world, const MAX_TERM_CNT: usize> Query<'world, MAX_TERM_CNT>
+impl<'world, const TERM_CAP: usize> Query<'world, TERM_CAP>
 {
     /// Iterates over the entities matching this query.
     #[must_use]
-    pub fn iter(&self) -> Iter<'_, 'world>
+    pub fn iter<'this>(&'this self) -> Iter<'this, TERM_CAP>
+    where
+        'world: 'this,
     {
         Iter {
-            iter: self
+            archetype_iter: self
                 .world
                 .data
                 .component_storage
-                .search_archetypes(ArchetypeSearchTerms {
-                    present: &self.terms.present,
-                    absent: &self.terms.absent
-                })
-                .flat_map(
-                    (|archetype| {
-                        repeat_n(archetype, archetype.entity_cnt())
-                            .zip(archetype.entities())
-                    }) as ComponentIterMapFn,
-                ),
+                .search_archetypes(&self.terms),
+            current: None,
             world: self.world,
         }
     }
@@ -46,16 +39,19 @@ impl<'world, const MAX_TERM_CNT: usize> Query<'world, MAX_TERM_CNT>
         self.world
     }
 
-    pub(crate) fn new(world: &'world World, terms: Terms<MAX_TERM_CNT>) -> Self
+    pub(crate) fn new(world: &'world World, terms: Terms<TERM_CAP>) -> Self
     {
         Self { world, terms }
     }
 }
 
-impl<'query, 'world, const MAX_TERM_CNT: usize> IntoIterator for &'query Query<'world, MAX_TERM_CNT>
+impl<'query, 'world, const TERM_CAP: usize> IntoIterator
+    for &'query Query<'world, TERM_CAP>
+where
+    'world: 'query,
 {
-    type IntoIter = Iter<'query, 'world>;
-    type Item = EntityHandle<'world>;
+    type IntoIter = Iter<'query, TERM_CAP>;
+    type Item = EntityHandle<'query>;
 
     fn into_iter(self) -> Self::IntoIter
     {
@@ -63,30 +59,43 @@ impl<'query, 'world, const MAX_TERM_CNT: usize> IntoIterator for &'query Query<'
     }
 }
 
-pub struct Iter<'query, 'world>
+pub struct Iter<'query, const TERM_CAP: usize>
 {
-    iter: QueryEntityIter<'query, 'world>,
-    world: &'world World,
+    archetype_iter: ArchetypeRefIter<'query, 'query, TERM_CAP>,
+    current: Option<(
+        &'query Archetype,
+        ArrayVec<TraversalResult, TERM_CAP>,
+        EntityIter<'query>,
+    )>,
+    world: &'query World,
 }
 
-impl<'query, 'world> Iterator for Iter<'query, 'world>
+impl<'query, const TERM_CAP: usize> Iter<'query, TERM_CAP>
 {
-    type Item = EntityHandle<'world>;
-
-    fn next(&mut self) -> Option<Self::Item>
+    pub fn traversal_results(&self) -> Option<&[TraversalResult]>
     {
-        let (archetype, entity) = self.iter.next()?;
-
-        Some(EntityHandle::new(archetype, entity, self.world))
+        self.current
+            .as_ref()
+            .map(|(_, traversal_results, _)| traversal_results.as_ref())
     }
 }
 
-type ComponentIterMapFnOutput<'a> = Zip<RepeatN<&'a Archetype>, EntityIter<'a>>;
+impl<'query, const TERM_CAP: usize> Iterator for Iter<'query, TERM_CAP>
+{
+    type Item = EntityHandle<'query>;
 
-type ComponentIterMapFn = for<'a> fn(&'a Archetype) -> ComponentIterMapFnOutput<'a>;
+    fn next(&mut self) -> Option<Self::Item>
+    {
+        if let Some((archetype, _, entity_iter)) = self.current.as_mut() {
+            if let Some(entity) = entity_iter.next() {
+                return Some(EntityHandle::new(archetype, entity, self.world));
+            }
+        }
+
+        let (archetype, traversal_results) = self.archetype_iter.next()?;
 
-type QueryEntityIter<'query, 'world> = FlatMap<
-    ArchetypeRefIter<'world, 'query>,
-    ComponentIterMapFnOutput<'world>,
-    ComponentIterMapFn,
->;
+        self.current = Some((archetype, traversal_results, archetype.entities()));
+
+        self.next()
+    }
+}