refactor(ecs): make entity finding a ComponentStorage method

1 files changed, 41 insertions, 30 deletions

diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index a2fc36e..9d79b91 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -6,7 +6,6 @@ use std::fmt::Debug;
 use std::marker::PhantomData;
 use std::mem::ManuallyDrop;
 use std::ops::RangeBounds;
-use std::slice::Iter as SliceIter;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Weak};
 use std::vec::Drain;
@@ -262,7 +261,8 @@ where
         'this: 'world,
     {
         QueryComponentIter {
-            entity_iter: self.component_storage.entities.iter(),
+            component_storage: &self.component_storage,
+            current_entity_index: 0,
             component_type_ids: Comps::type_ids(),
             comps_pd: PhantomData,
         }
@@ -434,7 +434,8 @@ impl QueryComponentIds
 
 pub struct QueryComponentIter<'world, Comps>
 {
-    entity_iter: SliceIter<'world, Entity>,
+    component_storage: &'world ComponentStorage,
+    current_entity_index: usize,
     component_type_ids: Vec<TypeId>,
     comps_pd: PhantomData<Comps>,
 }
@@ -447,8 +448,13 @@ where
 
     fn next(&mut self) -> Option<Self::Item>
     {
-        let matching_entity =
-            find_entity_with_components(&mut self.entity_iter, &self.component_type_ids)?;
+        let (matching_entity_index, matching_entity) =
+            self.component_storage.find_entity_with_components(
+                self.current_entity_index,
+                &self.component_type_ids,
+            )?;
+
+        self.current_entity_index = matching_entity_index + 1;
 
         Some(Comps::from_components(
             matching_entity
@@ -459,37 +465,42 @@ where
     }
 }
 
-fn find_entity_with_components<'world>(
-    entity_iter: &mut impl Iterator<Item = &'world Entity>,
-    component_type_ids: &[TypeId],
-) -> Option<&'world Entity>
-{
-    // TODO: This is a really dumb and slow way to do this. Refactor the world
-    // to store components in archetypes
-    entity_iter.find(|entity| {
-        let entity_components = entity
-            .components
-            .iter()
-            .map(|component| component.id)
-            .collect::<HashSet<_>>();
-
-        if component_type_ids
-            .iter()
-            .all(|component_type_id| entity_components.contains(component_type_id))
-        {
-            return true;
-        }
-
-        false
-    })
-}
-
 #[derive(Debug, Default)]
 pub struct ComponentStorage
 {
     entities: Vec<Entity>,
 }
 
+impl ComponentStorage
+{
+    fn find_entity_with_components(
+        &self,
+        start_index: usize,
+        component_type_ids: &[TypeId],
+    ) -> Option<(usize, &Entity)>
+    {
+        // TODO: This is a really dumb and slow way to do this. Refactor the world
+        // to store components in archetypes
+        self.entities.iter().enumerate().skip(start_index).find(
+            move |(_index, entity)| {
+                let entity_components = entity
+                    .components
+                    .iter()
+                    .map(|component| component.id)
+                    .collect::<HashSet<_>>();
+
+                if component_type_ids.iter().all(|component_type_id| {
+                    entity_components.contains(component_type_id)
+                }) {
+                    return true;
+                }
+
+                false
+            },
+        )
+    }
+}
+
 impl TypeName for ComponentStorage
 {
     fn type_name(&self) -> &'static str