1 files changed, 38 insertions, 59 deletions

diff --git a/ecs/src/component.rs b/ecs/src/component.rs
index 7a997c5..07701c8 100644
--- a/ecs/src/component.rs
+++ b/ecs/src/component.rs
@@ -1,10 +1,16 @@
 use std::any::{Any, TypeId};
+use std::cell::{RefCell, RefMut};
 use std::fmt::Debug;
 use std::ops::{Deref, DerefMut};
 
 use seq_macro::seq;
 
-use crate::system::{Input as SystemInput, Param as SystemParam, System};
+use crate::system::{
+    ComponentRefMut,
+    Input as SystemInput,
+    Param as SystemParam,
+    System,
+};
 use crate::WorldData;
 
 pub trait Component: SystemInput + Any
@@ -49,35 +55,27 @@ pub trait Sequence
     where
         Self: 'component;
 
-    type MutRefs<'component>
-    where
-        Self: 'component;
-
     fn into_vec(self) -> Vec<Box<dyn Component>>;
 
     fn type_ids() -> Vec<TypeId>;
 
-    fn from_components(components: &[Box<dyn Component>]) -> Self::Refs<'_>;
-
-    fn from_components_mut(components: &mut [Box<dyn Component>]) -> Self::MutRefs<'_>;
+    fn from_components(components: &[RefCell<Box<dyn Component>>]) -> Self::Refs<'_>;
 }
 
 macro_rules! inner {
     ($c: tt) => {
         seq!(I in 0..=$c {
             impl<#(Comp~I: Component,)*> Sequence for (#(Comp~I,)*) {
-                type Refs<'component> = (#(&'component Comp~I,)*)
+                type Refs<'component> = (#(ComponentRefMut<'component, Comp~I>,)*)
                     where Self: 'component;
 
-                type MutRefs<'component> = (#(&'component mut Comp~I,)*)
-                    where Self: 'component;
-
-
-                fn into_vec(self) -> Vec<Box<dyn Component>> {
+                fn into_vec(self) -> Vec<Box<dyn Component>>
+                {
                     Vec::from_iter([#(Box::new(self.I) as Box<dyn Component>,)*])
                 }
 
-                fn type_ids() -> Vec<TypeId> {
+                fn type_ids() -> Vec<TypeId>
+                {
                     vec![
                         #(
                             TypeId::of::<Comp~I>(),
@@ -85,46 +83,34 @@ macro_rules! inner {
                     ]
                 }
 
-                fn from_components(components: &[Box<dyn Component>]) -> Self::Refs<'_>
+                fn from_components(
+                    components: &[RefCell<Box<dyn Component>>]
+                ) -> Self::Refs<'_>
                 {
                     #(
-                        let mut comp_~I = None;
+                        let mut comp_~I: Option<RefMut<Box<dyn Component>>> = None;
                     )*
 
                     for comp in components {
-                        #(
-                            if comp.is::<Comp~I>() {
-                                comp_~I = Some(comp);
-                                continue;
-                            }
-                        )*
-                    }
-
-                    (#(
-                        comp_~I.unwrap().downcast_ref::<Comp~I>().unwrap(),
-                    )*)
-
-                }
-
-                fn from_components_mut(
-                    components: &mut [Box<dyn Component>],
-                ) -> Self::MutRefs<'_>
-                {
-                    #(
-                        let mut comp_~I = None;
-                    )*
+                        let Ok(comp_ref) = comp.try_borrow_mut() else {
+                            continue;
+                        };
 
-                    for comp in components.iter_mut() {
                         #(
-                            if comp.is::<Comp~I>() {
-                                comp_~I = Some(comp);
+                            if comp_ref.is::<Comp~I>() {
+                                comp_~I = Some(comp_ref);
                                 continue;
                             }
                         )*
                     }
 
                     (#(
-                        comp_~I.unwrap().downcast_mut::<Comp~I>().unwrap(),
+                        ComponentRefMut::new(
+                            RefMut::filter_map(
+                                comp_~I.unwrap(),
+                                |component| component.downcast_mut::<Comp~I>()
+                            ).expect("Failed to downcast component")
+                        ),
                     )*)
                 }
             }
@@ -140,17 +126,7 @@ seq!(C in 0..=64 {
 #[derive(Debug)]
 pub struct Local<'world, LocalComponent: Component>
 {
-    local_component: &'world mut LocalComponent,
-}
-
-impl<'world, LocalComponent> Local<'world, LocalComponent>
-where
-    LocalComponent: Component,
-{
-    fn new(local_component: &'world mut LocalComponent) -> Self
-    {
-        Self { local_component }
-    }
+    local_component: ComponentRefMut<'world, LocalComponent>,
 }
 
 unsafe impl<'world, LocalComponent> SystemParam<'world> for Local<'world, LocalComponent>
@@ -160,21 +136,24 @@ where
     type Flags = ();
     type Input = LocalComponent;
 
-    fn initialize<SystemImpl>(system: &mut impl System<SystemImpl>, input: Self::Input)
+    fn initialize<SystemImpl>(
+        system: &mut impl System<'world, SystemImpl>,
+        input: Self::Input,
+    )
     {
         system.set_local_component(input);
     }
 
     fn new<SystemImpl>(
-        system: &'world mut impl System<SystemImpl>,
-        _world_data: &'world mut WorldData,
+        system: &'world impl System<'world, SystemImpl>,
+        _world_data: &'world WorldData,
     ) -> Self
     {
         let local_component = system
             .get_local_component_mut::<LocalComponent>()
             .expect("Local component is uninitialized");
 
-        Self::new(local_component)
+        Self { local_component }
     }
 
     fn is_compatible<Other: SystemParam<'world>>() -> bool
@@ -202,7 +181,7 @@ where
 
     fn deref(&self) -> &Self::Target
     {
-        self.local_component
+        &self.local_component
     }
 }
 
@@ -212,6 +191,6 @@ where
 {
     fn deref_mut(&mut self) -> &mut Self::Target
     {
-        self.local_component
+        &mut self.local_component
     }
 }