1 files changed, 34 insertions, 25 deletions

diff --git a/ecs/src/system/stateful.rs b/ecs/src/system/stateful.rs
index 3f71000..7ce87fa 100644
--- a/ecs/src/system/stateful.rs
+++ b/ecs/src/system/stateful.rs
@@ -1,12 +1,12 @@
-use std::any::{type_name, TypeId};
+use std::any::{type_name, Any, TypeId};
+use std::cell::{RefCell, RefMut};
 use std::collections::HashMap;
-use std::ptr::addr_of_mut;
 
 use seq_macro::seq;
 
 use crate::component::Component;
 use crate::system::util::check_params_are_compatible;
-use crate::system::{Into as IntoSystem, Param, System, TypeErased};
+use crate::system::{ComponentRefMut, Into as IntoSystem, Param, System, TypeErased};
 use crate::tuple::{
     Filter as TupleFilter,
     FilterExclude as TupleFilterExclude,
@@ -20,13 +20,13 @@ use crate::WorldData;
 pub struct Stateful<Func>
 {
     func: Func,
-    local_components: HashMap<TypeId, Box<dyn Component>>,
+    local_components: HashMap<TypeId, RefCell<Box<dyn Component>>>,
 }
 
 macro_rules! impl_system {
     ($c: tt) => {
         seq!(I in 0..$c {
-            impl<'world, Func, #(TParam~I,)*> System<fn(&'world (), #(TParam~I,)*)>
+            impl<'world, Func, #(TParam~I,)*> System<'world, fn(&'world (), #(TParam~I,)*)>
                 for Stateful<Func>
             where
                 Func: Fn(#(TParam~I,)*) + Copy + 'static,
@@ -76,7 +76,9 @@ macro_rules! impl_system {
                     self
                 }
 
-                fn run(&mut self, world_data: &mut WorldData)
+                fn run<'this>(&'this self, world_data: &'world WorldData)
+                where
+                    'this: 'world
                 {
                     #(
                         check_params_are_compatible!(I, TParam~I, $c);
@@ -85,17 +87,7 @@ macro_rules! impl_system {
                     let func = self.func;
 
                     func(#({
-                        // SAFETY: All parameters are compatible so this is fine
-                        let this = unsafe {
-                            &mut *addr_of_mut!(*self)
-                        };
-
-                        // SAFETY: All parameters are compatible so this is fine
-                        let world_data = unsafe {
-                            &mut *addr_of_mut!(*world_data)
-                        };
-
-                        TParam~I::new(this, world_data)
+                        TParam~I::new(self, world_data)
                     },)*);
                 }
 
@@ -104,7 +96,17 @@ macro_rules! impl_system {
                     TypeErased {
                         data: Box::new(self),
                         func: Box::new(|data, world_data| {
-                            let me = data.downcast_mut::<Self>().unwrap();
+                            // 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.run(world_data);
                         }),
@@ -112,12 +114,17 @@ macro_rules! impl_system {
                 }
 
                 fn get_local_component_mut<LocalComponent: Component>(
-                    &mut self,
-                ) -> Option<&mut LocalComponent>
+                    &self,
+                ) -> Option<ComponentRefMut<LocalComponent>>
                 {
-                    self.local_components
-                        .get_mut(&TypeId::of::<LocalComponent>())?
-                        .downcast_mut()
+                    let local_component = self.local_components
+                        .get(&TypeId::of::<LocalComponent>())?
+                        .borrow_mut();
+
+                    Some(ComponentRefMut::new(RefMut::filter_map(
+                        local_component,
+                        |local_component| local_component.downcast_mut()
+                    ).ok()?))
                 }
 
                 fn set_local_component<LocalComponent: Component>(
@@ -126,8 +133,10 @@ macro_rules! impl_system {
                 )
                 {
                     self.local_components
-                        .insert(TypeId::of::<LocalComponent>(),
-                        Box::new(local_component));
+                        .insert(
+                            TypeId::of::<LocalComponent>(),
+                            RefCell::new(Box::new(local_component))
+                        );
                 }
             }