perf(ecs): store components using archetypes

6 files changed, 393 insertions, 56 deletions

diff --git a/ecs/src/component.rs b/ecs/src/component.rs
index d3b00ef..5c0b9ce 100644
--- a/ecs/src/component.rs
+++ b/ecs/src/component.rs
@@ -28,6 +28,11 @@ pub trait Component: SystemInput + Any + TypeName
 
     #[doc(hidden)]
     fn as_any(&self) -> &dyn Any;
+
+    fn is_optional(&self) -> bool
+    {
+        false
+    }
 }
 
 impl dyn Component
@@ -80,6 +85,11 @@ where
     {
         self
     }
+
+    fn is_optional(&self) -> bool
+    {
+        true
+    }
 }
 
 impl<ComponentT> TypeName for Option<ComponentT>
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);
+    }
 }
diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index 1983f66..0dd1c02 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -83,6 +83,8 @@ impl World
     ) where
         EventT: Event,
     {
+        system.prepare(&self.data);
+
         self.systems.push(system.into_type_erased());
 
         self.data
@@ -159,6 +161,12 @@ impl World
     /// Will panic if a internal lock cannot be acquired.
     pub fn event_loop<EventSeq: EventSequence>(&self)
     {
+        self.data
+            .component_storage
+            .write_nonblock()
+            .expect("Failed to acquire read-write component storage lock")
+            .make_archetype_lookup_entries();
+
         let event_seq = EventSeq::ids();
 
         loop {
diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index 88ca180..de6c832 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -1,8 +1,11 @@
-use std::any::{Any, TypeId};
+use std::any::{type_name, Any, TypeId};
 use std::collections::HashSet;
+use std::iter::{Flatten, Map};
 use std::marker::PhantomData;
+use std::slice::Iter as SliceIter;
 use std::sync::{Arc, Weak};
 
+use crate::component::storage::Archetype;
 use crate::component::{
     IsOptional as ComponentIsOptional,
     Sequence as ComponentSequence,
@@ -13,7 +16,7 @@ use crate::system::{
     Param as SystemParam,
     System,
 };
-use crate::{ComponentStorage, WorldData};
+use crate::{ComponentStorage, EntityComponent, WorldData};
 
 #[derive(Debug)]
 pub struct Query<'world, Comps>
@@ -34,10 +37,17 @@ where
     where
         'this: 'world,
     {
+        #[cfg(feature = "debug")]
+        tracing::debug!("Searching for {}", type_name::<Comps>());
+
         ComponentIter {
-            component_storage: &self.component_storage,
-            current_entity_index: 0,
-            component_type_ids: Comps::type_ids(),
+            entities: self
+                .component_storage
+                .find_entities(&Comps::type_ids())
+                .unwrap_or_else(|| panic!("Could not find {:?}", type_name::<Comps>()))
+                .iter()
+                .map((|archetype| archetype.components.as_slice()) as ComponentIterMapFn)
+                .flatten(),
             comps_pd: PhantomData,
         }
     }
@@ -118,6 +128,33 @@ where
             component_type_ids: Comps::type_ids(),
         })
     }
+
+    fn handle_pre_run(world_data: &WorldData)
+    {
+        let mut component_storage_lock = world_data
+            .component_storage
+            .write_nonblock()
+            .expect("Failed to acquire read-write component storage lock");
+
+        #[cfg(feature = "debug")]
+        tracing::debug!(
+            "Adding archetypes lookup entry for components: ({})",
+            type_name::<Comps>()
+        );
+
+        component_storage_lock.add_archetype_lookup_entry(
+            &Comps::type_ids()
+                .into_iter()
+                .filter_map(|(component_id, is_optional)| {
+                    if is_optional == ComponentIsOptional::Yes {
+                        return None;
+                    }
+
+                    Some(component_id)
+                })
+                .collect::<Vec<_>>(),
+        );
+    }
 }
 
 /// A entity query containing a weak reference to the world.
@@ -182,11 +219,11 @@ where
     }
 }
 
+type ComponentIterMapFn = for<'a> fn(&'a &'a Archetype) -> &'a [Vec<EntityComponent>];
+
 pub struct ComponentIter<'world, Comps>
 {
-    component_storage: &'world ComponentStorage,
-    current_entity_index: usize,
-    component_type_ids: Vec<(TypeId, ComponentIsOptional)>,
+    entities: Flatten<Map<SliceIter<'world, &'world Archetype>, ComponentIterMapFn>>,
     comps_pd: PhantomData<Comps>,
 }
 
@@ -198,17 +235,7 @@ where
 
     fn next(&mut self) -> Option<Self::Item>
     {
-        let (matching_entity_index, matching_entity_components) =
-            self.component_storage.find_entity_with_components(
-                self.current_entity_index,
-                &self.component_type_ids,
-            )?;
-
-        self.current_entity_index = matching_entity_index + 1;
-
-        Some(Comps::from_components(
-            matching_entity_components.iter().map(|component| component),
-        ))
+        Some(Comps::from_components(self.entities.next()?.iter()))
     }
 }
 
diff --git a/ecs/src/system.rs b/ecs/src/system.rs
index 47872a6..868afa7 100644
--- a/ecs/src/system.rs
+++ b/ecs/src/system.rs
@@ -25,6 +25,8 @@ pub trait System<'world, Impl>: 'static
     #[must_use]
     fn initialize(self, input: Self::Input) -> Self;
 
+    fn prepare(&self, world_data: &WorldData);
+
     fn run<'this>(&'this self, world_data: &'world WorldData)
     where
         'this: 'world;
@@ -57,6 +59,13 @@ macro_rules! impl_system {
                     self
                 }
 
+                fn prepare(&self, world_data: &WorldData)
+                {
+                    #(
+                        TParam~I::handle_pre_run(world_data);
+                    )*
+                }
+
                 fn run<'this>(&'this self, world_data: &'world WorldData)
                 where
                     'this: 'world
@@ -179,6 +188,8 @@ pub unsafe trait Param<'world>
     fn is_compatible<Other: Param<'world>>() -> bool;
 
     fn get_comparable() -> Box<dyn Any>;
+
+    fn handle_pre_run(_world_data: &WorldData) {}
 }
 
 pub struct NoInitParamFlag {}
diff --git a/ecs/src/system/stateful.rs b/ecs/src/system/stateful.rs
index 6c0b554..b765ff0 100644
--- a/ecs/src/system/stateful.rs
+++ b/ecs/src/system/stateful.rs
@@ -85,6 +85,13 @@ macro_rules! impl_system {
                     self
                 }
 
+                fn prepare(&self, world_data: &WorldData)
+                {
+                    #(
+                        TParam~I::handle_pre_run(world_data);
+                    )*
+                }
+
                 fn run<'this>(&'this self, world_data: &'world WorldData)
                 where
                     'this: 'world