fix(ecs): always populate archetype lookup entries

1 files changed, 25 insertions, 33 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index e939ba7..4b7a6b4 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -40,9 +40,9 @@ impl Storage
         let archetype_id = ArchetypeId::from_components_metadata(&comp_metadata_list);
 
         if !self.archetype_lookup.borrow().contains_key(&archetype_id) {
-            self.create_populated_archetype_lookup_entry(
+            self.archetype_lookup.borrow_mut().insert(
                 archetype_id,
-                comp_metadata_list,
+                self.create_populated_archetype_lookup_entry(comp_metadata_list.as_ref()),
             );
         }
 
@@ -115,8 +115,7 @@ impl Storage
                 .map(|component| ComponentMetadata::get(&**component)),
         );
 
-        let archetype_index =
-            self.get_or_create_archetype(archetype_id, &comp_ids_set, &components);
+        let archetype_index = self.get_or_create_archetype(archetype_id, &components);
 
         self.populate_matching_archetype_lookup_entries(&comp_ids_set, archetype_index);
 
@@ -243,29 +242,25 @@ impl Storage
     fn get_or_create_archetype(
         &mut self,
         archetype_id: ArchetypeId,
-        comp_ids_set: &HashSet<Uid>,
         components: &[Box<dyn Component>],
     ) -> usize
     {
         let mut archetype_lookup = self.archetype_lookup.borrow_mut();
 
-        let lookup_entry = archetype_lookup.entry(archetype_id).or_insert_with(|| {
-            ArchetypeLookupEntry {
-                component_ids: comp_ids_set.clone(),
-                archetype_indices: Vec::with_capacity(1),
-            }
-        });
-
-        if lookup_entry.archetype_indices.is_empty() {
+        if !archetype_lookup.contains_key(&archetype_id) {
             self.archetypes.push(Archetype::new(
                 components.iter().map(|component| component.self_id()),
             ));
-
-            lookup_entry
-                .archetype_indices
-                .push(self.archetypes.len() - 1);
         }
 
+        let lookup_entry = archetype_lookup.entry(archetype_id).or_insert_with(|| {
+            self.create_populated_archetype_lookup_entry(
+                components
+                    .iter()
+                    .map(|component| ComponentMetadata::get(&**component)),
+            )
+        });
+
         // SAFETY: Above, we push a archetype index if archetype_indices is empty so this
         // cannot fail
         unsafe { *lookup_entry.archetype_indices.first().unwrap_unchecked() }
@@ -311,14 +306,14 @@ impl Storage
         }
     }
 
-    fn create_populated_archetype_lookup_entry<CompMetadataList>(
+    fn create_populated_archetype_lookup_entry<CompMetadataIter>(
         &self,
-        archetype_id: ArchetypeId,
-        comp_metadata_list: CompMetadataList,
-    ) where
-        CompMetadataList: AsRef<[ComponentMetadata]>,
+        comp_metadata_iter: CompMetadataIter,
+    ) -> ArchetypeLookupEntry
+    where
+        CompMetadataIter: IntoIterator<Item: Borrow<ComponentMetadata>>,
     {
-        let comp_ids_set = create_non_opt_component_id_set(comp_metadata_list.as_ref());
+        let comp_ids_set = create_non_opt_component_id_set(comp_metadata_iter);
 
         let mut exact_matching_archetype_index = None;
 
@@ -341,16 +336,13 @@ impl Storage
             })
             .collect::<Vec<_>>();
 
-        self.archetype_lookup.borrow_mut().insert(
-            archetype_id,
-            ArchetypeLookupEntry {
-                component_ids: comp_ids_set,
-                archetype_indices: exact_matching_archetype_index
-                    .into_iter()
-                    .chain(matching_archetype_indices)
-                    .collect(),
-            },
-        );
+        ArchetypeLookupEntry {
+            component_ids: comp_ids_set,
+            archetype_indices: exact_matching_archetype_index
+                .into_iter()
+                .chain(matching_archetype_indices)
+                .collect(),
+        }
     }
 }