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

1 files changed, 37 insertions, 4 deletions

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`].
 ///