perf(ecs): use component index map when creating component sequences

1 files changed, 38 insertions, 15 deletions

diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index aa424e5..c4b4cf0 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -1,4 +1,4 @@
-use std::iter::{Filter, Flatten, Map};
+use std::iter::{repeat_n, Filter, Flatten, Map, RepeatN, Zip};
 use std::marker::PhantomData;
 
 use crate::component::storage::{
@@ -51,9 +51,14 @@ where
             entities: self
                 .component_storage
                 .find_entities(Comps::metadata())
-                .map(Archetype::entities as ComponentIterMapFn)
+                .map(
+                    (|archetype| {
+                        repeat_n(archetype, archetype.entity_cnt())
+                            .zip(archetype.entities())
+                    }) as ComponentIterMapFn,
+                )
                 .flatten()
-                .filter(|entity| OptionsT::entity_filter(entity.components())),
+                .filter(|(_, entity)| OptionsT::entity_filter(entity.components())),
             comps_pd: PhantomData,
         }
     }
@@ -71,9 +76,14 @@ where
             entities: self
                 .component_storage
                 .find_entities(Comps::metadata())
-                .map(Archetype::entities as ComponentIterMapFn)
+                .map(
+                    (|archetype| {
+                        repeat_n(archetype, archetype.entity_cnt())
+                            .zip(archetype.entities())
+                    }) as ComponentIterMapFn,
+                )
                 .flatten()
-                .filter(|entity| OptionsT::entity_filter(entity.components())),
+                .filter(|(_, entity)| OptionsT::entity_filter(entity.components())),
             comps_pd: PhantomData,
         }
     }
@@ -103,9 +113,14 @@ where
                         .chain(extra_components)
                         .collect::<Vec<_>>(),
                 )
-                .map(Archetype::entities as ComponentIterMapFn)
+                .map(
+                    (|archetype| {
+                        repeat_n(archetype, archetype.entity_cnt())
+                            .zip(archetype.entities())
+                    }) as ComponentIterMapFn,
+                )
                 .flatten()
-                .filter(|entity| OptionsT::entity_filter(entity.components())),
+                .filter(|(_, entity)| OptionsT::entity_filter(entity.components())),
             comps_pd: PhantomData,
         }
     }
@@ -175,9 +190,11 @@ where
     }
 }
 
-type ComponentIterMapFn = for<'a> fn(&'a Archetype) -> EntityIter<'a>;
+type ComponentIterMapFn =
+    for<'a> fn(&'a Archetype) -> Zip<RepeatN<&'a Archetype>, EntityIter<'a>>;
 
-type ComponentIterFilterFn = for<'a, 'b> fn(&'a &'b ArchetypeEntity) -> bool;
+type ComponentIterFilterFn =
+    for<'a, 'b> fn(&'a (&'b Archetype, &'b ArchetypeEntity)) -> bool;
 
 type QueryEntityIter<'world> = Filter<
     Flatten<Map<ArchetypeRefIter<'world>, ComponentIterMapFn>>,
@@ -186,7 +203,7 @@ type QueryEntityIter<'world> = Filter<
 
 pub struct ComponentIterMut<'world, Comps, EntityIter>
 where
-    EntityIter: Iterator<Item = &'world ArchetypeEntity>,
+    EntityIter: Iterator<Item = (&'world Archetype, &'world ArchetypeEntity)>,
 {
     world: &'world World,
     entities: EntityIter,
@@ -196,14 +213,17 @@ where
 impl<'world, Comps, EntityIter> Iterator for ComponentIterMut<'world, Comps, EntityIter>
 where
     Comps: ComponentSequence + 'world,
-    EntityIter: Iterator<Item = &'world ArchetypeEntity>,
+    EntityIter: Iterator<Item = (&'world Archetype, &'world ArchetypeEntity)>,
 {
     type Item = Comps::MutRefs<'world>;
 
     fn next(&mut self) -> Option<Self::Item>
     {
+        let (archetype, entity) = self.entities.next()?;
+
         Some(Comps::from_components_mut(
-            self.entities.next()?.components().iter(),
+            entity.components(),
+            |component_uid| archetype.get_index_for_component(component_uid),
             self.world,
             lock_component_rw,
         ))
@@ -227,7 +247,7 @@ fn lock_component_rw(
 
 pub struct ComponentIter<'world, Comps, EntityIter>
 where
-    EntityIter: Iterator<Item = &'world ArchetypeEntity>,
+    EntityIter: Iterator<Item = (&'world Archetype, &'world ArchetypeEntity)>,
 {
     world: &'world World,
     entities: EntityIter,
@@ -237,14 +257,17 @@ where
 impl<'world, Comps, EntityIter> Iterator for ComponentIter<'world, Comps, EntityIter>
 where
     Comps: ComponentSequence + 'world,
-    EntityIter: Iterator<Item = &'world ArchetypeEntity>,
+    EntityIter: Iterator<Item = (&'world Archetype, &'world ArchetypeEntity)>,
 {
     type Item = Comps::Refs<'world>;
 
     fn next(&mut self) -> Option<Self::Item>
     {
+        let (archetype, entity) = self.entities.next()?;
+
         Some(Comps::from_components(
-            self.entities.next()?.components().iter(),
+            entity.components(),
+            |component_uid| archetype.get_index_for_component(component_uid),
             self.world,
             lock_component_ro,
         ))