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

3 files changed, 75 insertions, 60 deletions

diff --git a/ecs/src/component.rs b/ecs/src/component.rs
index 5da510a..85c556f 100644
--- a/ecs/src/component.rs
+++ b/ecs/src/component.rs
@@ -179,19 +179,17 @@ pub trait Sequence
     fn removed_event_ids() -> Vec<Uid>;
 
     fn from_components_mut<'component>(
-        components: impl Iterator<Item = &'component EntityComponent>,
+        components: &'component [EntityComponent],
+        component_index_lookup: impl Fn(Uid) -> Option<usize>,
         world: &'component World,
-        lock_component: fn(
-            entity_component: &EntityComponent,
-        ) -> WriteGuard<'_, Box<dyn Component>>,
+        lock_component: fn(&EntityComponent) -> WriteGuard<'_, Box<dyn Component>>,
     ) -> Self::MutRefs<'component>;
 
     fn from_components<'component>(
-        components: impl Iterator<Item = &'component EntityComponent>,
+        components: &'component [EntityComponent],
+        component_index_lookup: impl Fn(Uid) -> Option<usize>,
         world: &'component World,
-        lock_component: fn(
-            entity_component: &EntityComponent,
-        ) -> ReadGuard<'_, Box<dyn Component>>,
+        lock_component: fn(&EntityComponent) -> ReadGuard<'_, Box<dyn Component>>,
     ) -> Self::Refs<'component>;
 }
 
@@ -304,56 +302,45 @@ macro_rules! inner {
                     ]
                 }
 
+                #[inline]
                 fn from_components_mut<'component>(
-                    components: impl Iterator<Item = &'component EntityComponent>,
+                    components: &'component [EntityComponent],
+                    component_index_lookup: impl Fn(Uid) -> Option<usize>,
                     world: &'component World,
-                    lock_component: fn(
-                        entity_component: &EntityComponent,
-                    ) -> WriteGuard<'_, Box<dyn Component>>,
+                    lock_component:
+                        fn(&EntityComponent) -> WriteGuard<'_, Box<dyn Component>>,
                 ) -> Self::MutRefs<'component>
                 {
-                    #(
-                        let mut comp_~I: Option<WriteGuard<Box<dyn Component>>> = None;
-                    )*
-
-                    for comp in components {
-                        #(
-                            if comp.id == Comp~I::Component::id() {
-                                comp_~I = Some(lock_component(comp));
-                                continue;
-                            }
-                        )*
-                    }
-
                     (#(
-                        Comp~I::RefMut::from_optional_mut_component(comp_~I, world),
+                        Comp~I::RefMut::from_optional_mut_component(
+                            component_index_lookup(Comp~I::Component::id())
+                                .and_then(|component_index| {
+                                    components.get(component_index)
+                                        .map(lock_component)
+                                }),
+                            world
+                        ),
                     )*)
                 }
 
+                #[inline]
                 fn from_components<'component>(
-                    components: impl Iterator<Item = &'component EntityComponent>,
+                    components: &'component [EntityComponent],
+                    component_index_lookup: impl Fn(Uid) -> Option<usize>,
                     world: &'component World,
-                    lock_component: fn(
-                        entity_component: &EntityComponent,
-                    ) -> ReadGuard<'_, Box<dyn Component>>,
+                    lock_component:
+                        fn(&EntityComponent) -> ReadGuard<'_, Box<dyn Component>>,
                 ) -> Self::Refs<'component>
                 {
-
-                    #(
-                        let mut comp_~I: Option<ReadGuard<Box<dyn Component>>> = None;
-                    )*
-
-                    for comp in components {
-                        #(
-                            if comp.id == Comp~I::Component::id() {
-                                comp_~I = Some(lock_component(comp));
-                                continue;
-                            }
-                        )*
-                    }
-
                     (#(
-                        Comp~I::Ref::from_optional_component(comp_~I, world),
+                        Comp~I::Ref::from_optional_component(
+                            component_index_lookup(Comp~I::Component::id())
+                                .and_then(|component_index| {
+                                    components.get(component_index)
+                                        .map(lock_component)
+                                }),
+                            world
+                        ),
                     )*)
                 }
             }
diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index dcf0181..b46d053 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -414,6 +414,11 @@ impl Archetype
         EntityIter { iter: self.entities.iter() }
     }
 
+    pub fn entity_cnt(&self) -> usize
+    {
+        self.entities.len()
+    }
+
     pub fn get_index_for_component(&self, component_id: Uid) -> Option<usize>
     {
         self.component_ids.get(&component_id).copied()
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,
         ))