refactor(ecs): create archetype lookup entries on-the-go

1 files changed, 168 insertions, 246 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index 2bd3826..6e06ded 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -1,4 +1,6 @@
 use std::any::type_name;
+use std::borrow::Borrow;
+use std::cell::RefCell;
 use std::collections::{HashMap, HashSet};
 use std::slice::Iter as SliceIter;
 
@@ -12,35 +14,83 @@ use crate::component::{
 use crate::entity::Uid as EntityUid;
 use crate::lock::Lock;
 use crate::type_name::TypeName;
-use crate::util::Sortable;
+use crate::util::{RefCellRefMap, Sortable};
 use crate::EntityComponent;
 
 #[derive(Debug, Default)]
 pub struct Storage
 {
     archetypes: Vec<Archetype>,
-    archetype_lookup: HashMap<ArchetypeId, Vec<usize>>,
-    pending_archetype_lookup_entries: Vec<Vec<ComponentMetadata>>,
+    archetype_lookup: RefCell<HashMap<ArchetypeId, ArchetypeLookupEntry>>,
 }
 
 impl Storage
 {
-    pub fn find_entities(
+    pub fn find_entities<CompMetadataList>(
         &self,
-        mut components_metadata: impl IntoIterator<Item = ComponentMetadata>
-            + Sortable<Item = ComponentMetadata>,
+        mut components_metadata: CompMetadataList,
     ) -> ArchetypeRefIter<'_>
+    where
+        CompMetadataList: Sortable<Item = ComponentMetadata>,
+        CompMetadataList: Borrow<[ComponentMetadata]>,
     {
         components_metadata.sort_by_key_b(|component_metadata| component_metadata.id);
 
-        self.archetype_lookup
-            .get(&ArchetypeId::from_components_metadata(components_metadata))
-            .map_or_else(ArchetypeRefIter::new_empty, |archetypes_indices| {
-                ArchetypeRefIter {
-                    inner: archetypes_indices.iter(),
-                    archetypes: &self.archetypes,
+        let archetype_id = ArchetypeId::from_components_metadata(
+            components_metadata.borrow().into_iter().cloned(),
+        );
+
+        // This looks stupid but the borrow checker complains otherwise
+        if self.archetype_lookup.borrow().contains_key(&archetype_id) {
+            let archetype_lookup = self.archetype_lookup.borrow();
+
+            return ArchetypeRefIter {
+                inner: RefCellRefMap::new(archetype_lookup, |archetype_lookup| {
+                    archetype_lookup
+                        .get(&archetype_id)
+                        .unwrap()
+                        .archetype_indices
+                        .iter()
+                }),
+                archetypes: &self.archetypes,
+            };
+        }
+
+        let comp_ids_set = create_non_opt_component_id_set(components_metadata.borrow());
+
+        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,
+            },
+        );
+
+        let archetype_lookup = self.archetype_lookup.borrow();
+
+        ArchetypeRefIter {
+            inner: RefCellRefMap::new(archetype_lookup, |archetype_lookup| {
+                archetype_lookup
+                    .get(&archetype_id)
+                    .unwrap()
+                    .archetype_indices
+                    .iter()
+            }),
+            archetypes: &self.archetypes,
+        }
     }
 
     #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
@@ -67,48 +117,25 @@ impl Storage
                 .map(|component| ComponentMetadata::of(&**component)),
         );
 
-        let archetype_indices =
-            self.archetype_lookup
-                .entry(archetype_id)
-                .or_insert_with(|| {
-                    self.archetypes.push(Archetype::new(
-                        components.iter().map(|component| component.id()),
-                    ));
-
-                    vec![self.archetypes.len() - 1]
-                });
+        let comp_ids_set = create_non_opt_component_id_set(
+            components
+                .iter()
+                .map(|component| ComponentMetadata::of(&**component)),
+        );
 
-        if archetype_indices.is_empty() {
-            self.archetypes.push(Archetype::new(
-                components.iter().map(|component| component.id()),
-            ));
+        let archetype_index =
+            self.get_or_create_archetype(archetype_id, &comp_ids_set, &components);
 
-            archetype_indices.push(self.archetypes.len() - 1);
-        }
+        self.populate_matching_archetype_lookup_entries(&comp_ids_set, archetype_index);
 
         let archetype = self
             .archetypes
-            .get_mut(
-                archetype_indices
-                    .first()
-                    .copied()
-                    .expect("No archetype index found"),
-            )
+            .get_mut(archetype_index)
             .expect("Archetype is gone");
 
         let entity_uid = EntityUid::new_unique();
 
-        archetype.entities.push(ArchetypeEntity {
-            uid: entity_uid,
-            components: components
-                .into_iter()
-                .map(|component| EntityComponent {
-                    id: component.id(),
-                    name: component.type_name(),
-                    component: Lock::new(component),
-                })
-                .collect(),
-        });
+        archetype.push_entity(entity_uid, components);
 
         archetype
             .entity_lookup
@@ -117,54 +144,54 @@ impl Storage
         (archetype_id, entity_uid)
     }
 
-    pub fn add_archetype_lookup_entry(
+    fn populate_matching_archetype_lookup_entries(
         &mut self,
-        mut components_metadata: impl IntoIterator<Item = ComponentMetadata>
-            + Sortable<Item = ComponentMetadata>,
+        comp_ids_set: &HashSet<ComponentId>,
+        archetype_index: usize,
     )
     {
-        components_metadata.sort_by_key_b(|component_metadata| component_metadata.id);
+        let mut archetype_lookup = self.archetype_lookup.borrow_mut();
 
-        self.pending_archetype_lookup_entries
-            .push(components_metadata.into_iter().collect());
+        for (_, lookup_entry) in archetype_lookup.iter_mut() {
+            if &lookup_entry.component_ids == comp_ids_set {
+                continue;
+            }
+
+            if lookup_entry.component_ids.is_subset(&comp_ids_set) {
+                lookup_entry.archetype_indices.push(archetype_index);
+            }
+        }
     }
 
-    pub fn make_archetype_lookup_entries(&mut self)
+    fn get_or_create_archetype(
+        &mut self,
+        archetype_id: ArchetypeId,
+        comp_ids_set: &HashSet<ComponentId>,
+        components: &[Box<dyn Component>],
+    ) -> usize
     {
-        for pending_entry in self.pending_archetype_lookup_entries.drain(..) {
-            let pending_entry_ids_set = pending_entry
-                .iter()
-                .filter_map(|component_metadata| {
-                    if component_metadata.is_optional == ComponentIsOptional::Yes {
-                        return None;
-                    }
-
-                    Some(component_metadata.id)
-                })
-                .collect::<HashSet<_>>();
-
-            let matching_archetype_indices = self
-                .archetypes
-                .iter()
-                .enumerate()
-                .filter_map(|(index, archetype)| {
-                    if archetype.component_ids_is_superset(&pending_entry_ids_set) {
-                        return Some(index);
-                    }
-
-                    None
-                });
+        let mut archetype_lookup = self.archetype_lookup.borrow_mut();
 
-            let archetype_id =
-                ArchetypeId::from_components_metadata(pending_entry.into_iter());
-
-            if self.archetype_lookup.contains_key(&archetype_id) {
-                continue;
+        let lookup_entry = archetype_lookup.entry(archetype_id).or_insert_with(|| {
+            ArchetypeLookupEntry {
+                component_ids: comp_ids_set.clone(),
+                archetype_indices: Vec::with_capacity(1),
             }
+        });
 
-            self.archetype_lookup
-                .insert(archetype_id, matching_archetype_indices.collect());
+        if lookup_entry.archetype_indices.is_empty() {
+            self.archetypes.push(Archetype::new(
+                components.iter().map(|component| component.id()),
+            ));
+
+            lookup_entry
+                .archetype_indices
+                .push(self.archetypes.len() - 1);
         }
+
+        // SAFETY: Above, we push a archetype index if archetype_indices is empty so this
+        // cannot fail
+        unsafe { *lookup_entry.archetype_indices.first().unwrap_unchecked() }
     }
 }
 
@@ -177,6 +204,13 @@ impl TypeName for Storage
 }
 
 #[derive(Debug)]
+struct ArchetypeLookupEntry
+{
+    component_ids: HashSet<ComponentId>,
+    archetype_indices: Vec<usize>,
+}
+
+#[derive(Debug)]
 pub struct Archetype
 {
     component_ids: HashMap<ComponentId, usize>,
@@ -229,6 +263,25 @@ impl Archetype
     {
         self.component_ids.get(component_id).copied()
     }
+
+    fn push_entity(
+        &mut self,
+        entity_uid: EntityUid,
+        components: impl IntoIterator<Item = Box<dyn Component>>,
+    )
+    {
+        self.entities.push(ArchetypeEntity {
+            uid: entity_uid,
+            components: components
+                .into_iter()
+                .map(|component| EntityComponent {
+                    id: component.id(),
+                    name: component.type_name(),
+                    component: Lock::new(component),
+                })
+                .collect(),
+        });
+    }
 }
 
 #[derive(Debug)]
@@ -259,18 +312,14 @@ impl ArchetypeEntity
 #[derive(Debug)]
 pub struct ArchetypeRefIter<'component_storage>
 {
-    inner: SliceIter<'component_storage, usize>,
+    inner: RefCellRefMap<
+        'component_storage,
+        HashMap<ArchetypeId, ArchetypeLookupEntry>,
+        SliceIter<'component_storage, usize>,
+    >,
     archetypes: &'component_storage [Archetype],
 }
 
-impl<'component_storage> ArchetypeRefIter<'component_storage>
-{
-    fn new_empty() -> Self
-    {
-        Self { inner: [].iter(), archetypes: &[] }
-    }
-}
-
 impl<'component_storage> Iterator for ArchetypeRefIter<'component_storage>
 {
     type Item = &'component_storage Archetype;
@@ -303,24 +352,40 @@ impl<'archetype> Iterator for EntityIter<'archetype>
     }
 }
 
+fn create_non_opt_component_id_set<'a, Item>(
+    component_metadata_iter: impl IntoIterator<Item = Item>,
+) -> HashSet<ComponentId>
+where
+    Item: Borrow<ComponentMetadata>,
+{
+    component_metadata_iter
+        .into_iter()
+        .filter_map(|item| {
+            let component_metadata = item.borrow();
+
+            if component_metadata.is_optional == ComponentIsOptional::Yes {
+                return None;
+            }
+
+            Some(component_metadata.id)
+        })
+        .collect::<HashSet<_>>()
+}
+
 #[cfg(test)]
 mod tests
 {
-    use std::collections::HashMap;
 
     use ecs_macros::Component;
 
-    use super::{Archetype, Storage};
+    use super::Storage;
     use crate::archetype::Id as ArchetypeId;
-    use crate::component::storage::ArchetypeEntity;
     use crate::component::{
         Id as ComponentId,
         IsOptional as ComponentIsOptional,
         Metadata as ComponentMetadata,
     };
-    use crate::entity::Uid as EntityUid;
-    use crate::lock::Lock;
-    use crate::{self as ecs, EntityComponent};
+    use crate::{self as ecs};
 
     #[derive(Debug, Component)]
     struct HealthPotion
@@ -378,7 +443,7 @@ mod tests
 
         assert_eq!(entity_components.components.len(), 2);
 
-        assert_eq!(component_storage.archetype_lookup.len(), 1);
+        assert_eq!(component_storage.archetype_lookup.borrow().len(), 1);
 
         let mut components_metadata = [
             ComponentMetadata {
@@ -393,160 +458,17 @@ mod tests
 
         components_metadata.sort_by_key(|comp_metadata| comp_metadata.id);
 
-        let lookup = component_storage
-            .archetype_lookup
+        let archetype_lookup = component_storage.archetype_lookup.borrow();
+
+        let lookup_entry = archetype_lookup
             .get(&ArchetypeId::from_components_metadata(components_metadata))
             .expect("Expected entry in archetype lookup map");
 
-        let first_archetype_index = lookup
+        let first_archetype_index = lookup_entry
+            .archetype_indices
             .first()
             .expect("Expected archetype lookup to contain a archetype reference");
 
         assert_eq!(*first_archetype_index, 0);
     }
-
-    #[test]
-    fn lookup_works()
-    {
-        let mut component_storage = Storage::default();
-
-        let entity_uid_a = EntityUid::new_unique();
-        let entity_uid_b = EntityUid::new_unique();
-        let entity_uid_c = EntityUid::new_unique();
-
-        component_storage.archetypes.push(Archetype {
-            component_ids: HashMap::from([
-                (ComponentId::of::<IronBoots>(), 0),
-                (ComponentId::of::<HealthPotion>(), 1),
-                (ComponentId::of::<Hookshot>(), 2),
-            ]),
-            entity_lookup: HashMap::from([
-                (entity_uid_a, 0),
-                (entity_uid_b, 1),
-                (entity_uid_c, 2),
-            ]),
-            entities: vec![
-                ArchetypeEntity {
-                    uid: entity_uid_a,
-                    components: vec![EntityComponent {
-                        id: ComponentId::of::<IronBoots>(),
-                        name: "Iron boots",
-                        component: Lock::new(Box::new(IronBoots)),
-                    }],
-                },
-                ArchetypeEntity {
-                    uid: entity_uid_b,
-                    components: vec![EntityComponent {
-                        id: ComponentId::of::<HealthPotion>(),
-                        name: "Health potion",
-                        component: Lock::new(Box::new(HealthPotion {
-                            _hp_restoration: 20,
-                        })),
-                    }],
-                },
-                ArchetypeEntity {
-                    uid: entity_uid_c,
-                    components: vec![EntityComponent {
-                        id: ComponentId::of::<Hookshot>(),
-                        name: "Hookshot",
-                        component: Lock::new(Box::new(Hookshot { _range: 67 })),
-                    }],
-                },
-            ],
-        });
-
-        let entity_uid_d = EntityUid::new_unique();
-        let entity_uid_e = EntityUid::new_unique();
-        let entity_uid_f = EntityUid::new_unique();
-        let entity_uid_g = EntityUid::new_unique();
-
-        component_storage.archetypes.push(Archetype {
-            component_ids: HashMap::from([
-                (ComponentId::of::<DekuNut>(), 0),
-                (ComponentId::of::<IronBoots>(), 1),
-                (ComponentId::of::<Bow>(), 2),
-                (ComponentId::of::<Hookshot>(), 3),
-            ]),
-            entity_lookup: HashMap::from([
-                (entity_uid_d, 0),
-                (entity_uid_e, 1),
-                (entity_uid_f, 2),
-                (entity_uid_g, 3),
-            ]),
-            entities: vec![
-                ArchetypeEntity {
-                    uid: entity_uid_d,
-                    components: vec![EntityComponent {
-                        id: ComponentId::of::<DekuNut>(),
-                        name: "Deku nut",
-                        component: Lock::new(Box::new(DekuNut { _throwing_damage: 5 })),
-                    }],
-                },
-                ArchetypeEntity {
-                    uid: entity_uid_e,
-                    components: vec![EntityComponent {
-                        id: ComponentId::of::<IronBoots>(),
-                        name: "Iron boots",
-                        component: Lock::new(Box::new(IronBoots)),
-                    }],
-                },
-                ArchetypeEntity {
-                    uid: entity_uid_f,
-                    components: vec![EntityComponent {
-                        id: ComponentId::of::<Bow>(),
-                        name: "Bow",
-                        component: Lock::new(Box::new(Bow { _damage: 20 })),
-                    }],
-                },
-                ArchetypeEntity {
-                    uid: entity_uid_g,
-                    components: vec![EntityComponent {
-                        id: ComponentId::of::<Hookshot>(),
-                        name: "Hookshot",
-                        component: Lock::new(Box::new(Hookshot { _range: 67 })),
-                    }],
-                },
-            ],
-        });
-
-        component_storage.add_archetype_lookup_entry([
-            ComponentMetadata {
-                id: ComponentId::of::<IronBoots>(),
-                is_optional: ComponentIsOptional::No,
-            },
-            ComponentMetadata {
-                id: ComponentId::of::<Hookshot>(),
-                is_optional: ComponentIsOptional::No,
-            },
-        ]);
-
-        assert_eq!(component_storage.pending_archetype_lookup_entries.len(), 1);
-
-        component_storage.make_archetype_lookup_entries();
-
-        assert_eq!(component_storage.archetype_lookup.len(), 1);
-
-        let mut comps_metadata = [
-            ComponentMetadata {
-                id: ComponentId::of::<IronBoots>(),
-                is_optional: ComponentIsOptional::No,
-            },
-            ComponentMetadata {
-                id: ComponentId::of::<Hookshot>(),
-                is_optional: ComponentIsOptional::No,
-            },
-        ];
-
-        comps_metadata.sort_by_key(|comp_metadata| comp_metadata.id);
-
-        let archetypes = component_storage
-            .archetype_lookup
-            .get(&ArchetypeId::from_components_metadata(comps_metadata))
-            .expect(concat!(
-                "Expected a archetype for IronBoots & Hookshot to be found in the ",
-                "archetype lookup map"
-            ));
-
-        assert_eq!(archetypes.len(), 2);
-    }
 }