fix(ecs): update archetype lookups when entity is spawned with Actions

4 files changed, 83 insertions, 16 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index 237dd73..86f28fc 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -131,9 +131,13 @@ impl ComponentStorage
 
     pub fn make_archetype_lookup_entries(&mut self)
     {
-        for pending_entry in &self.pending_archetype_lookup_entries {
+        // TODO: Re-doing the whole archetype lookup table is dumb and slow. Only re-do
+        // the archetype lookup entries that need to be changed
+        self.archetype_lookup.clear();
+
+        for pending_entry in self.pending_archetype_lookup_entries.drain(..) {
             let components_set: HashSet<_> = pending_entry
-                .into_iter()
+                .iter()
                 .map(|component_id| *component_id)
                 .collect();
 
@@ -151,9 +155,7 @@ impl ComponentStorage
 
             let archetype_indices = self
                 .archetype_lookup
-                .entry(ArchetypeComponentsHash::new(
-                    pending_entry.into_iter().copied().clone(),
-                ))
+                .entry(ArchetypeComponentsHash::new(pending_entry.into_iter()))
                 .or_default();
 
             archetype_indices.extend(matching_archetype_indices);
diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index e427150..42cdef3 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -142,12 +142,29 @@ impl World
             .drain(..)
         {
             match action {
-                Action::Spawn(components) => self
-                    .data
-                    .component_storage
-                    .write_nonblock()
-                    .expect("Failed to acquire read-write component storage lock")
-                    .push_entity(components),
+                Action::Spawn(components) => {
+                    {
+                        let mut component_storage_lock =
+                            self.data.component_storage.write_nonblock().expect(
+                                "Failed to acquire read-write component storage lock",
+                            );
+
+                        component_storage_lock.push_entity(components);
+                    }
+
+                    for system in &self.systems {
+                        unsafe {
+                            system.prepare(&self.data);
+                        }
+                    }
+
+                    let mut component_storage_lock =
+                        self.data.component_storage.write_nonblock().expect(
+                            "Failed to acquire read-write component storage lock",
+                        );
+
+                    component_storage_lock.make_archetype_lookup_entries();
+                }
                 Action::Stop => {
                     self.stop.store(true, Ordering::Relaxed);
                 }
diff --git a/ecs/src/system.rs b/ecs/src/system.rs
index 868afa7..f8885fc 100644
--- a/ecs/src/system.rs
+++ b/ecs/src/system.rs
@@ -85,7 +85,7 @@ macro_rules! impl_system {
                 {
                     TypeErased {
                         data: Box::new(self),
-                        func: Box::new(|data, world_data| {
+                        run: Box::new(|data, world_data| {
                             // SAFETY: The caller of TypeErased::run ensures the lifetime
                             // is correct
                             let data = unsafe { &*addr_of!(*data) };
@@ -102,6 +102,23 @@ macro_rules! impl_system {
 
                             me.run(world_data);
                         }),
+                        prepare: Box::new(|data, world_data| {
+                            // SAFETY: The caller of TypeErased::run ensures the lifetime
+                            // is correct
+                            let data = unsafe { &*addr_of!(*data) };
+
+                            let me = data
+                                .downcast_ref::<Func>()
+                                .expect("Function downcast failed");
+
+                            // SAFETY: The caller of TypeErased::run ensures the lifetime
+                            // is correct
+                            let world_data = unsafe {
+                                &*(world_data as *const WorldData)
+                            };
+
+                            me.prepare(world_data);
+                        }),
                     }
                 }
 
@@ -137,7 +154,8 @@ pub trait Into<Impl>
 pub struct TypeErased
 {
     data: Box<dyn Any + RefUnwindSafe + UnwindSafe>,
-    func: Box<TypeErasedFunc>,
+    run: Box<TypeErasedRunFn>,
+    prepare: Box<TypeErasedPrepareFn>,
 }
 
 impl TypeErased
@@ -151,7 +169,19 @@ impl TypeErased
         // You have to dereference for downcasting to work for some reason
         let data = &*self.data;
 
-        (self.func)(data, world_data);
+        (self.run)(data, world_data);
+    }
+
+    /// Prepares the system.
+    ///
+    /// # Safety
+    /// `world_data` must live at least as long as the [`World`] the system belongs to.
+    pub unsafe fn prepare(&self, world_data: &WorldData)
+    {
+        // You have to dereference for downcasting to work for some reason
+        let data = &*self.data;
+
+        (self.prepare)(data, world_data);
     }
 }
 
@@ -164,7 +194,10 @@ impl Debug for TypeErased
 }
 
 /// Function in [`TypeErased`] used to run the system.
-type TypeErasedFunc = dyn Fn(&dyn Any, &WorldData) + RefUnwindSafe + UnwindSafe;
+type TypeErasedRunFn = dyn Fn(&dyn Any, &WorldData) + RefUnwindSafe + UnwindSafe;
+
+/// Function in [`TypeErased`] used to prepare the system.
+type TypeErasedPrepareFn = dyn Fn(&dyn Any, &WorldData) + RefUnwindSafe + UnwindSafe;
 
 /// A parameter to a [`System`].
 ///
diff --git a/ecs/src/system/stateful.rs b/ecs/src/system/stateful.rs
index ae6a5b5..49fd3bf 100644
--- a/ecs/src/system/stateful.rs
+++ b/ecs/src/system/stateful.rs
@@ -111,7 +111,7 @@ macro_rules! impl_system {
                 {
                     TypeErased {
                         data: Box::new(self),
-                        func: Box::new(|data, world_data| {
+                        run: Box::new(|data, world_data| {
                             // SAFETY: The caller of TypeErased::run ensures the lifetime
                             // is correct
                             let data = unsafe { &*(data as *const dyn Any) };
@@ -126,6 +126,21 @@ macro_rules! impl_system {
 
                             me.run(world_data);
                         }),
+                        prepare: Box::new(|data, world_data| {
+                            // SAFETY: The caller of TypeErased::run ensures the lifetime
+                            // is correct
+                            let data = unsafe { &*(data as *const dyn Any) };
+
+                            let me = data.downcast_ref::<Self>().unwrap();
+
+                            // SAFETY: The caller of TypeErased::run ensures the lifetime
+                            // is correct
+                            let world_data = unsafe {
+                                &*(world_data as *const WorldData)
+                            };
+
+                            me.prepare(world_data);
+                        }),
                     }
                 }