refactor(ecs): reduce amount of code in Storage::find_entities

1 files changed, 39 insertions, 29 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index b46d053..e1523df 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -29,44 +29,23 @@ impl Storage
 {
     pub fn find_entities<CompMetadataList>(
         &self,
-        mut components_metadata: CompMetadataList,
+        mut comp_metadata_list: CompMetadataList,
     ) -> ArchetypeRefIter<'_>
     where
         CompMetadataList: Sortable<Item = ComponentMetadata>,
         CompMetadataList: AsRef<[ComponentMetadata]>,
     {
-        components_metadata.sort_by_key_b(|component_metadata| component_metadata.id);
+        comp_metadata_list.sort_by_key_b(|metadata| metadata.id);
 
-        let archetype_id = ArchetypeId::from_components_metadata(&components_metadata);
+        let archetype_id = ArchetypeId::from_components_metadata(&comp_metadata_list);
 
-        // This looks stupid but the borrow checker complains otherwise
-        if self.archetype_lookup.borrow().contains_key(&archetype_id) {
-            return self.iter_archetypes_by_lookup(archetype_id);
+        if !self.archetype_lookup.borrow().contains_key(&archetype_id) {
+            self.create_populated_archetype_lookup_entry(
+                archetype_id,
+                comp_metadata_list,
+            );
         }
 
-        let comp_ids_set = create_non_opt_component_id_set(components_metadata.as_ref());
-
-        let matching_archetype_indices = self
-            .archetypes
-            .iter()
-            .enumerate()
-            .filter_map(|(index, archetype)| {
-                if archetype.component_ids_is_superset(&comp_ids_set) {
-                    return Some(index);
-                }
-
-                None
-            })
-            .collect();
-
-        self.archetype_lookup.borrow_mut().insert(
-            archetype_id,
-            ArchetypeLookupEntry {
-                component_ids: comp_ids_set,
-                archetype_indices: matching_archetype_indices,
-            },
-        );
-
         self.iter_archetypes_by_lookup(archetype_id)
     }
 
@@ -343,6 +322,37 @@ impl Storage
             archetypes: &self.archetypes,
         }
     }
+
+    fn create_populated_archetype_lookup_entry<CompMetadataList>(
+        &self,
+        archetype_id: ArchetypeId,
+        comp_metadata_list: CompMetadataList,
+    ) where
+        CompMetadataList: AsRef<[ComponentMetadata]>,
+    {
+        let comp_ids_set = create_non_opt_component_id_set(comp_metadata_list.as_ref());
+
+        let matching_archetype_indices = self
+            .archetypes
+            .iter()
+            .enumerate()
+            .filter_map(|(index, archetype)| {
+                if archetype.component_ids_is_superset(&comp_ids_set) {
+                    return Some(index);
+                }
+
+                None
+            })
+            .collect();
+
+        self.archetype_lookup.borrow_mut().insert(
+            archetype_id,
+            ArchetypeLookupEntry {
+                component_ids: comp_ids_set,
+                archetype_indices: matching_archetype_indices,
+            },
+        );
+    }
 }
 
 /// Component storage error