1 files changed, 145 insertions, 58 deletions

diff --git a/ecs/src/system/stateful.rs b/ecs/src/system/stateful.rs
index b641cf2..9b2f279 100644
--- a/ecs/src/system/stateful.rs
+++ b/ecs/src/system/stateful.rs
@@ -1,67 +1,154 @@
-use std::marker::PhantomData;
+use std::any::{type_name, TypeId};
+use std::collections::HashMap;
+use std::ptr::addr_of_mut;
 
-use crate::component::Local;
-use crate::system::{System, TypeErased};
-use crate::{ComponentStorage, Query};
+use seq_macro::seq;
+
+use crate::component::Component;
+use crate::system::util::check_params_are_compatible;
+use crate::system::{
+    FilteredInputs,
+    InputFilter,
+    Into as IntoSystem,
+    OptionInputs,
+    Param,
+    System,
+    TakeOptionInputResult,
+    TypeErased,
+};
+use crate::ComponentStorage;
 
 /// A stateful system.
-pub struct Stateful<Func, LocalComponent, Comps>
+pub struct Stateful<Func>
 {
     func: Func,
-    local_component: Option<LocalComponent>,
-    _comps_pd: PhantomData<Comps>,
+    local_components: HashMap<TypeId, Box<dyn Component>>,
 }
 
-impl<Func: 'static, Comps: 'static, LocalComponent: 'static>
-    System<fn(Query<Comps>, Local<LocalComponent>)>
-    for Stateful<Func, LocalComponent, Comps>
-where
-    Func: Fn(Query<Comps>, Local<LocalComponent>),
-{
-    type Input = LocalComponent;
-    type Query<'a> = Query<'a, Comps>;
-
-    fn initialize(mut self, input: Self::Input) -> Self
-    {
-        self.local_component = Some(input);
-
-        self
-    }
-
-    fn run(&mut self, component_storage: &mut ComponentStorage)
-    {
-        (self.func)(
-            Query::new(component_storage),
-            Local::new(self.local_component.as_mut().unwrap()),
-        );
-    }
-
-    fn into_type_erased(self) -> TypeErased
-    {
-        TypeErased {
-            data: Box::new(self),
-            func: Box::new(move |data, component_storage| {
-                let this = data.downcast_mut::<Self>().unwrap();
-
-                this.run(component_storage);
-            }),
-        }
-    }
-}
+macro_rules! impl_system {
+    ($c: tt) => {
+        seq!(I in 0..=$c {
+            impl<'world, Func, #(TParam~I,)*> System<fn(&'world (), #(TParam~I,)*)>
+                for Stateful<Func>
+            where
+                Func: Fn(#(TParam~I,)*) + Copy + 'static,
+                #(TParam~I: Param<'world>,)*
+                #(TParam~I::Input: 'static,)*
+                (#(TParam~I::Input,)*): InputFilter
+            {
+                type Input = <(#(TParam~I::Input,)*) as InputFilter>::Filtered;
 
-impl<Func: 'static, Comps, LocalComponent: 'static>
-    crate::system::Into<fn(Query<Comps>, Local<LocalComponent>)> for Func
-where
-    Func: Fn(Query<Comps>, Local<LocalComponent>),
-{
-    type System = Stateful<Func, LocalComponent, Comps>;
-
-    fn into_system(self) -> Self::System
-    {
-        Self::System {
-            func: self,
-            local_component: None,
-            _comps_pd: PhantomData,
-        }
-    }
+                fn initialize(mut self, input: Self::Input) -> Self
+                {
+                    let mut option_input = input.into_in_options();
+
+                    #(
+                        if TypeId::of::<TParam~I::Input>() != TypeId::of::<()>() {
+                            let input = match option_input.take::<TParam~I::Input>() {
+                                TakeOptionInputResult::Found(input) => input,
+                                TakeOptionInputResult::NotFound => {
+                                    panic!(
+                                        "Parameter input {} not found",
+                                        type_name::<TParam~I::Input>()
+                                    );
+                                }
+                                TakeOptionInputResult::AlreadyTaken => {
+                                    panic!(
+                                        concat!(
+                                            "Parameter {} is already initialized. ",
+                                            "System cannot contain multiple inputs with ",
+                                            "the same type",
+                                        ),
+                                        type_name::<TParam~I>()
+                                    );
+
+                                }
+                            };
+
+                            TParam~I::initialize(
+                                &mut self,
+                                input
+                            );
+                        }
+                    )*
+
+                    self
+                }
+
+                fn run(&mut self, component_storage: &mut ComponentStorage)
+                {
+                    #(
+                        check_params_are_compatible!(I, TParam~I, $c);
+                    )*
+
+                    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 component_storage = unsafe {
+                            &mut *addr_of_mut!(*component_storage)
+                        };
+
+                        TParam~I::new(this, component_storage)
+                    },)*);
+                }
+
+                fn into_type_erased(self) -> TypeErased
+                {
+                    TypeErased {
+                        data: Box::new(self),
+                        func: Box::new(|data, component_storage| {
+                            let me = data.downcast_mut::<Self>().unwrap();
+
+                            me.run(component_storage);
+                        }),
+                    }
+                }
+
+                fn get_local_component_mut<LocalComponent: Component>(
+                    &mut self,
+                ) -> Option<&mut LocalComponent>
+                {
+                    self.local_components
+                        .get_mut(&TypeId::of::<LocalComponent>())?
+                        .downcast_mut()
+                }
+
+                fn set_local_component<LocalComponent: Component>(
+                    &mut self,
+                    local_component: LocalComponent,
+                )
+                {
+                    self.local_components
+                        .insert(TypeId::of::<LocalComponent>(),
+                        Box::new(local_component));
+                }
+            }
+
+            impl<Func, #(TParam~I,)*> IntoSystem<fn(#(TParam~I,)*)>
+                for Func
+            where
+                Func: Fn(#(TParam~I,)*) + Copy + 'static,
+            {
+                type System = Stateful<Func>;
+
+                fn into_system(self) -> Self::System
+                {
+                    Self::System {
+                        func: self,
+                        local_components: HashMap::new(),
+                    }
+                }
+            }
+        });
+    };
 }
+
+seq!(C in 0..4 {
+    impl_system!(C);
+});