perf(ecs): store components using archetypes

1 files changed, 311 insertions, 37 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index cdff09e..cc9e911 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -1,5 +1,7 @@
 use std::any::{type_name, TypeId};
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
+use std::hash::{DefaultHasher, Hash, Hasher};
+use std::ptr::NonNull;
 
 use crate::component::{Component, IsOptional as ComponentIsOptional};
 use crate::lock::Lock;
@@ -9,52 +11,82 @@ use crate::EntityComponent;
 #[derive(Debug, Default)]
 pub struct ComponentStorage
 {
-    entities: Vec<Entity>,
+    archetypes: Vec<Archetype>,
+    archetype_lookup: HashMap<ArchetypeComponentsHash, Vec<NonNull<Archetype>>>,
+    pending_archetype_lookup_entries: Vec<Vec<TypeId>>,
 }
 
 impl ComponentStorage
 {
-    pub fn find_entity_with_components(
+    pub fn find_entities(
         &self,
-        start_index: usize,
-        component_type_ids: &[(TypeId, ComponentIsOptional)],
-    ) -> Option<(usize, &[EntityComponent])>
+        component_ids: &[(TypeId, ComponentIsOptional)],
+    ) -> Option<&[&Archetype]>
     {
-        // TODO: This is a really dumb and slow way to do this. Refactor the world
-        // to store components in archetypes
-        self.entities
+        let ids = component_ids
             .iter()
-            .enumerate()
-            .skip(start_index)
-            .find(move |(_index, entity)| {
-                let entity_components = entity
-                    .components
-                    .iter()
-                    .map(|component| component.id)
-                    .collect::<HashSet<_>>();
-
-                if component_type_ids
-                    .iter()
-                    .filter(|(_, is_optional)| *is_optional == ComponentIsOptional::No)
-                    .all(|(component_type_id, _)| {
-                        entity_components.contains(component_type_id)
-                    })
-                {
-                    return true;
+            .filter_map(|(component_id, is_optional)| {
+                if *is_optional == ComponentIsOptional::Yes {
+                    return None;
                 }
 
-                false
+                Some(*component_id)
+            });
+
+        self.archetype_lookup
+            .get(&ArchetypeComponentsHash::new(ids))
+            .map(|archetypes|
+                // SAFETY: All NonNull<Archetype>s are references to items of the
+                // archetypes field and the items won't be dropped until the whole
+                // struct is dropped
+                unsafe {
+                nonnull_slice_to_ref_slice(archetypes.as_slice())
             })
-            .map(|(index, entity)| (index, &*entity.components))
     }
 
-    pub fn push_entity(
-        &mut self,
-        components: impl IntoIterator<Item = Box<dyn Component>>,
-    )
+    #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
+    pub fn push_entity(&mut self, components: Vec<Box<dyn Component>>)
     {
-        self.entities.push(Entity {
-            components: components
+        #[cfg(feature = "debug")]
+        tracing::debug!(
+            "Pushing entity with components: ({})",
+            components
+                .iter()
+                .map(|component| component.type_name())
+                .collect::<Vec<_>>()
+                .join(", ")
+        );
+
+        let archetypes = self
+            .archetype_lookup
+            .entry(ArchetypeComponentsHash::new(
+                components
+                    .iter()
+                    .filter(|component| !component.is_optional())
+                    .map(|component| (*component).type_id()),
+            ))
+            .or_insert_with(|| {
+                self.archetypes.push(Archetype::default());
+
+                vec![NonNull::from(self.archetypes.last().unwrap())]
+            });
+
+        // SAFETY: All NonNull<Archetype>s are references to items of the
+        // archetypes field and the items won't be dropped until the whole
+        // struct is dropped
+        let archetype = unsafe {
+            archetypes
+                .first_mut()
+                .expect("Archetype has disappeared")
+                .as_mut()
+        };
+
+        archetype
+            .component_ids
+            .extend(components.iter().map(|component| (*component).type_id()));
+
+        archetype.components.push(
+            components
                 .into_iter()
                 .map(|component| EntityComponent {
                     id: (*component).type_id(),
@@ -62,7 +94,40 @@ impl ComponentStorage
                     component: Lock::new(component),
                 })
                 .collect(),
-        });
+        );
+    }
+
+    pub fn add_archetype_lookup_entry(&mut self, component_ids: &[TypeId])
+    {
+        self.pending_archetype_lookup_entries
+            .push(component_ids.to_vec());
+    }
+
+    pub fn make_archetype_lookup_entries(&mut self)
+    {
+        for pending_entry in &self.pending_archetype_lookup_entries {
+            let components_set: HashSet<_> = pending_entry
+                .into_iter()
+                .map(|component_id| *component_id)
+                .collect();
+
+            let matching_archetypes = self.archetypes.iter().filter_map(|archetype| {
+                if archetype.component_ids.is_superset(&components_set) {
+                    return Some(NonNull::from(archetype));
+                }
+
+                None
+            });
+
+            let lookup_archetypes = self
+                .archetype_lookup
+                .entry(ArchetypeComponentsHash::new(
+                    pending_entry.into_iter().copied().clone(),
+                ))
+                .or_default();
+
+            lookup_archetypes.extend(matching_archetypes);
+        }
     }
 }
 
@@ -75,7 +140,216 @@ impl TypeName for ComponentStorage
 }
 
 #[derive(Debug, Default)]
-struct Entity
+pub struct Archetype
+{
+    component_ids: HashSet<TypeId>,
+    pub components: Vec<Vec<EntityComponent>>,
+}
+
+#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
+struct ArchetypeComponentsHash
+{
+    hash: u64,
+}
+
+impl ArchetypeComponentsHash
+{
+    fn new(component_ids: impl IntoIterator<Item = TypeId>) -> Self
+    {
+        let mut hasher = DefaultHasher::new();
+
+        for component_id in component_ids {
+            component_id.hash(&mut hasher);
+        }
+
+        let hash = hasher.finish();
+
+        Self { hash }
+    }
+}
+
+/// Casts a `&[NonNull<Item>]` to a `&[&Item]`.
+///
+/// # Safety
+/// All items in the slice must be initialized, properly aligned and follow Rust's
+/// aliasing rules.
+const unsafe fn nonnull_slice_to_ref_slice<Item>(slice: &[NonNull<Item>]) -> &[&Item]
 {
-    components: Vec<EntityComponent>,
+    unsafe { &*(std::ptr::from_ref(slice) as *const [&Item]) }
+}
+
+#[cfg(test)]
+mod tests
+{
+    use std::any::TypeId;
+    use std::collections::HashSet;
+    use std::ptr::addr_of;
+
+    use ecs_macros::Component;
+
+    use super::{Archetype, ArchetypeComponentsHash, ComponentStorage};
+    use crate::lock::Lock;
+    use crate::{self as ecs, EntityComponent};
+
+    #[derive(Debug, Component)]
+    struct HealthPotion
+    {
+        _hp_restoration: u32,
+    }
+
+    #[derive(Debug, Component)]
+    struct Hookshot
+    {
+        _range: u32,
+    }
+
+    #[derive(Debug, Component)]
+    struct DekuNut
+    {
+        _throwing_damage: u32,
+    }
+
+    #[derive(Debug, Component)]
+    struct Bow
+    {
+        _damage: u32,
+    }
+
+    #[derive(Debug, Component)]
+    struct IronBoots;
+
+    #[test]
+    fn push_entity_works()
+    {
+        let mut component_storage = ComponentStorage::default();
+
+        component_storage.push_entity(vec![
+            Box::new(HealthPotion { _hp_restoration: 12 }),
+            Box::new(Hookshot { _range: 50 }),
+        ]);
+
+        assert_eq!(component_storage.archetypes.len(), 1);
+
+        let archetype = component_storage
+            .archetypes
+            .first()
+            .expect("Expected a archetype in archetypes Vec");
+
+        assert_eq!(archetype.component_ids.len(), 2);
+
+        // One entity
+        assert_eq!(archetype.components.len(), 1);
+
+        let entity_components = archetype
+            .components
+            .first()
+            .expect("Expected a entity in archetype");
+
+        assert_eq!(entity_components.len(), 2);
+
+        assert_eq!(component_storage.archetype_lookup.len(), 1);
+
+        let lookup = component_storage
+            .archetype_lookup
+            .get(&ArchetypeComponentsHash::new([
+                TypeId::of::<HealthPotion>(),
+                TypeId::of::<Hookshot>(),
+            ]))
+            .expect("Expected entry in archetype lookup map");
+
+        let archetype_from_lookup = lookup
+            .first()
+            .expect("Expected archetype lookup to contain a archetype reference");
+
+        assert_eq!(
+            archetype_from_lookup.as_ptr() as usize,
+            addr_of!(*archetype) as usize
+        );
+    }
+
+    #[test]
+    fn lookup_works()
+    {
+        let mut component_storage = ComponentStorage::default();
+
+        component_storage.archetypes.push(Archetype {
+            component_ids: HashSet::from([
+                TypeId::of::<IronBoots>(),
+                TypeId::of::<HealthPotion>(),
+                TypeId::of::<Hookshot>(),
+            ]),
+            components: vec![
+                vec![EntityComponent {
+                    id: TypeId::of::<IronBoots>(),
+                    name: "Iron boots",
+                    component: Lock::new(Box::new(IronBoots)),
+                }],
+                vec![EntityComponent {
+                    id: TypeId::of::<HealthPotion>(),
+                    name: "Health potion",
+                    component: Lock::new(Box::new(HealthPotion { _hp_restoration: 20 })),
+                }],
+                vec![EntityComponent {
+                    id: TypeId::of::<Hookshot>(),
+                    name: "Hookshot",
+                    component: Lock::new(Box::new(Hookshot { _range: 67 })),
+                }],
+            ],
+        });
+
+        component_storage.archetypes.push(Archetype {
+            component_ids: HashSet::from([
+                TypeId::of::<DekuNut>(),
+                TypeId::of::<IronBoots>(),
+                TypeId::of::<Bow>(),
+                TypeId::of::<Hookshot>(),
+            ]),
+            components: vec![
+                vec![EntityComponent {
+                    id: TypeId::of::<DekuNut>(),
+                    name: "Deku nut",
+                    component: Lock::new(Box::new(DekuNut { _throwing_damage: 5 })),
+                }],
+                vec![EntityComponent {
+                    id: TypeId::of::<IronBoots>(),
+                    name: "Iron boots",
+                    component: Lock::new(Box::new(IronBoots)),
+                }],
+                vec![EntityComponent {
+                    id: TypeId::of::<Bow>(),
+                    name: "Bow",
+                    component: Lock::new(Box::new(Bow { _damage: 20 })),
+                }],
+                vec![EntityComponent {
+                    id: TypeId::of::<Hookshot>(),
+                    name: "Hookshot",
+                    component: Lock::new(Box::new(Hookshot { _range: 67 })),
+                }],
+            ],
+        });
+
+        component_storage.add_archetype_lookup_entry(&[
+            TypeId::of::<IronBoots>(),
+            TypeId::of::<Hookshot>(),
+        ]);
+
+        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 archetypes = component_storage
+            .archetype_lookup
+            .get(&ArchetypeComponentsHash::new([
+                TypeId::of::<IronBoots>(),
+                TypeId::of::<Hookshot>(),
+            ]))
+            .expect(concat!(
+                "Expected a archetype for IronBoots & Hookshot to be found in the ",
+                "archetype lookup map"
+            ));
+
+        assert_eq!(archetypes.len(), 2);
+    }
 }