perf(ecs): make archetype index by ID lookup faster

1 files changed, 40 insertions, 33 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index 5098c5f..53b1168 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -53,8 +53,7 @@ impl Storage
     {
         let archetype_id = self.entity_archetype_lookup.get(&entity_uid)?;
 
-        let archetype_index =
-            self.find_archetype_index_with_entity(*archetype_id, entity_uid)?;
+        let archetype_index = self.get_archetype_index_by_id(*archetype_id)?;
 
         self.archetypes.get(archetype_index)
     }
@@ -65,9 +64,7 @@ impl Storage
             return;
         };
 
-        let Some(archetype_index) =
-            self.find_archetype_index_with_entity(*archetype_id, entity_uid)
-        else {
+        let Some(archetype_index) = self.get_archetype_index_by_id(*archetype_id) else {
             return;
         };
 
@@ -148,8 +145,7 @@ impl Storage
     {
         let archetype_id = self.entity_archetype_lookup.get(&entity_uid)?;
 
-        let archetype_index =
-            self.find_archetype_index_with_entity(*archetype_id, entity_uid)?;
+        let archetype_index = self.get_archetype_index_by_id(*archetype_id)?;
 
         let archetype = self.archetypes.get_mut(archetype_index)?;
 
@@ -200,8 +196,7 @@ impl Storage
     {
         let archetype_id = self.entity_archetype_lookup.get(&entity_uid)?;
 
-        let archetype_index =
-            self.find_archetype_index_with_entity(*archetype_id, entity_uid)?;
+        let archetype_index = self.get_archetype_index_by_id(*archetype_id)?;
 
         let archetype = self.archetypes.get_mut(archetype_index)?;
 
@@ -280,29 +275,24 @@ impl Storage
         unsafe { *lookup_entry.archetype_indices.first().unwrap_unchecked() }
     }
 
-    fn find_archetype_index_with_entity(
-        &self,
-        archetype_id: ArchetypeId,
-        entity_uid: Uid,
-    ) -> Option<usize>
+    fn get_archetype_index_by_id(&self, archetype_id: ArchetypeId) -> Option<usize>
     {
-        let archetype_lookup = self.archetype_lookup.borrow_mut();
+        let archetype_lookup = self.archetype_lookup.borrow();
 
         let archetype_lookup_entry = archetype_lookup.get(&archetype_id)?;
 
-        // TODO: Also have a hashmap for precise archetype ID -> archetype index lookup.
-        // This way is slow
-        archetype_lookup_entry
+        let index = *archetype_lookup_entry
             .archetype_indices
-            .iter()
-            .find(|archetype_index| {
-                let Some(archetype) = self.archetypes.get(**archetype_index) else {
-                    return false;
-                };
+            .first()
+            .expect("No archetype indices in archetype lookup entry");
 
-                archetype.has_entity(entity_uid)
-            })
-            .copied()
+        debug_assert!(
+            self.archetypes.get(index).is_some_and(|archetype| archetype
+                .component_ids_is(&archetype_lookup_entry.component_ids)),
+            "Archetype components is not exact match"
+        );
+
+        Some(index)
     }
 
     fn iter_archetypes_by_lookup(&self, archetype_id: ArchetypeId)
@@ -334,24 +324,35 @@ impl Storage
     {
         let comp_ids_set = create_non_opt_component_id_set(comp_metadata_list.as_ref());
 
+        let mut exact_matching_archetype_index = None;
+
         let matching_archetype_indices = self
             .archetypes
             .iter()
             .enumerate()
             .filter_map(|(index, archetype)| {
+                if archetype.component_ids_is(&comp_ids_set) {
+                    exact_matching_archetype_index = Some(index);
+
+                    return None;
+                }
+
                 if archetype.component_ids_is_superset(&comp_ids_set) {
                     return Some(index);
                 }
 
                 None
             })
-            .collect();
+            .collect::<Vec<_>>();
 
         self.archetype_lookup.borrow_mut().insert(
             archetype_id,
             ArchetypeLookupEntry {
                 component_ids: comp_ids_set,
-                archetype_indices: matching_archetype_indices,
+                archetype_indices: exact_matching_archetype_index
+                    .into_iter()
+                    .chain(matching_archetype_indices)
+                    .collect(),
             },
         );
     }
@@ -414,6 +415,17 @@ impl Archetype
         }
     }
 
+    pub fn component_ids_is(&self, other_component_ids: &HashSet<Uid>) -> bool
+    {
+        if other_component_ids.len() == self.component_ids.len() {
+            other_component_ids
+                .iter()
+                .all(|v| self.component_ids.contains_key(v))
+        } else {
+            false
+        }
+    }
+
     pub fn get_entity(&self, entity_uid: Uid) -> Option<&ArchetypeEntity>
     {
         let entity_index = *self.entity_lookup.get(&entity_uid)?;
@@ -462,11 +474,6 @@ impl Archetype
 
         Some(entity)
     }
-
-    fn has_entity(&self, entity_uid: Uid) -> bool
-    {
-        self.entity_lookup.contains_key(&entity_uid)
-    }
 }
 
 #[derive(Debug)]