feat(ecs): add support for multiple system queries & local components

5 files changed, 474 insertions, 95 deletions

diff --git a/ecs/src/component.rs b/ecs/src/component.rs
index bead3b5..59b737e 100644
--- a/ecs/src/component.rs
+++ b/ecs/src/component.rs
@@ -4,6 +4,9 @@ use std::ops::{Deref, DerefMut};
 
 use seq_macro::seq;
 
+use crate::system::{Input as SystemInput, Param as SystemParam, System};
+use crate::ComponentStorage;
+
 pub trait Component: Any
 {
     #[doc(hidden)]
@@ -114,20 +117,65 @@ seq!(C in 0..=64 {
 
 /// Holds a component which is local to a single system.
 #[derive(Debug)]
-pub struct Local<'world, Value>
+pub struct Local<'world, Value: SystemInput>
 {
     value: &'world mut Value,
 }
 
 impl<'world, Value> Local<'world, Value>
+where
+    Value: SystemInput,
 {
-    pub(crate) fn new(value: &'world mut Value) -> Self
+    fn new(value: &'world mut Value) -> Self
     {
         Self { value }
     }
 }
 
+unsafe impl<'world, Value: 'static> SystemParam<'world> for Local<'world, Value>
+where
+    Value: SystemInput,
+{
+    type Flags = ();
+    type Input = Value;
+
+    fn initialize<SystemImpl>(system: &mut impl System<SystemImpl>, input: Self::Input)
+    {
+        system.set_local_component(input);
+    }
+
+    fn new<SystemImpl>(
+        system: &'world mut impl System<SystemImpl>,
+        _component_storage: &'world mut ComponentStorage,
+    ) -> Self
+    {
+        let local_component = system
+            .get_local_component_mut::<Value>()
+            .expect("Local component is uninitialized");
+
+        Self::new(local_component)
+    }
+
+    fn is_compatible<Other: SystemParam<'world>>() -> bool
+    {
+        let other_comparable = Other::get_comparable();
+
+        let Some(other_type_id) = other_comparable.downcast_ref::<TypeId>() else {
+            return true;
+        };
+
+        TypeId::of::<Value>() != *other_type_id
+    }
+
+    fn get_comparable() -> Box<dyn Any>
+    {
+        Box::new(TypeId::of::<Value>())
+    }
+}
+
 impl<'world, Value> Deref for Local<'world, Value>
+where
+    Value: SystemInput,
 {
     type Target = Value;
 
@@ -138,6 +186,8 @@ impl<'world, Value> Deref for Local<'world, Value>
 }
 
 impl<'world, Value> DerefMut for Local<'world, Value>
+where
+    Value: SystemInput,
 {
     fn deref_mut(&mut self) -> &mut Self::Target
     {
diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index df46a5a..84009e0 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -1,6 +1,6 @@
 #![deny(clippy::all, clippy::pedantic)]
 
-use std::any::TypeId;
+use std::any::{Any, TypeId};
 use std::collections::{HashMap, HashSet};
 use std::fmt::Debug;
 use std::hash::Hash;
@@ -8,7 +8,12 @@ use std::marker::PhantomData;
 use std::slice::IterMut as SliceIterMut;
 
 use crate::component::{Component, Sequence as ComponentSequence};
-use crate::system::{System, TypeErased as TypeErasedSystem};
+use crate::system::{
+    NoInitParamFlag as NoInitSystemParamFlag,
+    Param as SystemParam,
+    System,
+    TypeErased as TypeErasedSystem,
+};
 
 pub mod component;
 pub mod system;
@@ -142,6 +147,70 @@ where
     }
 }
 
+unsafe impl<'world, Comps> SystemParam<'world> for Query<'world, Comps>
+where
+    Comps: ComponentSequence,
+{
+    type Flags = NoInitSystemParamFlag;
+    type Input = ();
+
+    fn initialize<SystemImpl>(_system: &mut impl System<SystemImpl>, _input: Self::Input)
+    {
+    }
+
+    fn new<SystemImpl>(
+        _system: &'world mut impl System<SystemImpl>,
+        component_storage: &'world mut ComponentStorage,
+    ) -> Self
+    {
+        Self::new(component_storage)
+    }
+
+    fn is_compatible<Other: SystemParam<'world>>() -> bool
+    {
+        let other_comparable = Other::get_comparable();
+
+        let Some(other_query_component_ids) =
+            other_comparable.downcast_ref::<QueryComponentIds>()
+        else {
+            return true;
+        };
+
+        !other_query_component_ids.contains_component_in::<Comps>()
+    }
+
+    fn get_comparable() -> Box<dyn Any>
+    {
+        Box::new(QueryComponentIds {
+            component_type_ids: Comps::type_ids(),
+        })
+    }
+}
+
+#[derive(Debug)]
+struct QueryComponentIds
+{
+    component_type_ids: Vec<TypeId>,
+}
+
+impl QueryComponentIds
+{
+    fn contains_component_in<OtherComps>(&self) -> bool
+    where
+        OtherComps: ComponentSequence,
+    {
+        let other_component_type_ids = OtherComps::type_ids()
+            .into_iter()
+            .collect::<HashSet<TypeId>>();
+
+        // TODO: Make this a bit smarter. Queries with a same component can be compatible
+        // if one of the queries have a component the other one does not have
+        self.component_type_ids
+            .iter()
+            .any(|component_type_id| other_component_type_ids.contains(component_type_id))
+    }
+}
+
 pub struct QueryComponentIter<'world, Comps>
 {
     entity_iter: SliceIterMut<'world, Entity>,
diff --git a/ecs/src/system.rs b/ecs/src/system.rs
index 22e2215..d90e0a2 100644
--- a/ecs/src/system.rs
+++ b/ecs/src/system.rs
@@ -1,16 +1,21 @@
-use std::any::Any;
+use std::any::{Any, TypeId};
 use std::convert::Infallible;
 use std::fmt::Debug;
+use std::mem::{transmute_copy, ManuallyDrop};
+use std::ptr::addr_of_mut;
 
-use crate::component::Sequence as ComponentSequence;
-use crate::{ComponentStorage, Query};
+use seq_macro::seq;
+
+use crate::component::Component;
+use crate::system::util::check_params_are_compatible;
+use crate::ComponentStorage;
 
 pub mod stateful;
 
+mod util;
+
 pub trait System<Impl>: 'static
 {
-    type Query<'a>;
-
     type Input;
 
     #[must_use]
@@ -19,39 +24,90 @@ pub trait System<Impl>: 'static
     fn run(&mut self, component_storage: &mut ComponentStorage);
 
     fn into_type_erased(self) -> TypeErased;
-}
 
-impl<Func, Comps> System<fn(Query<Comps>)> for Func
-where
-    Func: Fn(Query<Comps>) + 'static,
-    Comps: ComponentSequence,
-{
-    type Input = Infallible;
-    type Query<'a> = Query<'a, Comps>;
+    fn get_local_component_mut<LocalComponent: Component>(
+        &mut self,
+    ) -> Option<&mut LocalComponent>;
 
-    fn initialize(self, _input: Self::Input) -> Self
-    {
-        self
-    }
-
-    fn run(&mut self, component_storage: &mut ComponentStorage)
-    {
-        self(Query::new(component_storage));
-    }
+    fn set_local_component<LocalComponent: Component>(
+        &mut self,
+        local_component: LocalComponent,
+    );
+}
 
-    fn into_type_erased(self) -> TypeErased
-    {
-        TypeErased {
-            data: Box::new(self),
-            func: Box::new(|data, component_storage| {
-                let me = data.downcast_mut::<Func>().unwrap();
-
-                me.run(component_storage);
-            }),
-        }
-    }
+macro_rules! impl_system {
+    ($c: tt) => {
+        seq!(I in 0..=$c {
+            impl<'world, Func, #(TParam~I,)*> System<fn(#(TParam~I,)*)>
+                for Func
+            where
+                Func: Fn(#(TParam~I,)*) + Copy + 'static,
+                #(TParam~I: Param<'world, Flags = NoInitParamFlag>,)*
+            {
+                type Input = Infallible;
+
+                fn initialize(self, _input: Self::Input) -> Self
+                {
+                    self
+                }
+
+                fn run(&mut self, component_storage: &mut ComponentStorage)
+                {
+                    #(
+                        check_params_are_compatible!(I, TParam~I, $c);
+                    )*
+
+                    let func = *self;
+
+                    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::<Func>().unwrap();
+
+                            me.run(component_storage);
+                        }),
+                    }
+                }
+
+                fn get_local_component_mut<LocalComponent: Component>(
+                    &mut self,
+                ) -> Option<&mut LocalComponent>
+                {
+                    panic!("System does not have any local components");
+                }
+
+                fn set_local_component<LocalComponent: Component>(
+                    &mut self,
+                    _local_component: LocalComponent,
+                ) {
+                    panic!("System does not have any local components");
+                }
+            }
+        });
+    };
 }
 
+seq!(C in 0..=4 {
+    impl_system!(C);
+});
+
 pub trait Into<Impl>
 {
     type System;
@@ -83,3 +139,107 @@ impl Debug for TypeErased
 
 /// Function in [`TypeErased`] used to run the system.
 type TypeErasedFunc = dyn Fn(&mut dyn Any, &mut ComponentStorage);
+
+/// A parameter to a [`System`].
+///
+/// # Safety
+/// The `is_compatible` function is used for safety so it must be implemented properly.
+pub unsafe trait Param<'world>
+{
+    type Input;
+    type Flags;
+
+    fn initialize<SystemImpl>(system: &mut impl System<SystemImpl>, input: Self::Input);
+
+    fn new<SystemImpl>(
+        system: &'world mut impl System<SystemImpl>,
+        component_storage: &'world mut ComponentStorage,
+    ) -> Self;
+
+    fn is_compatible<Other: Param<'world>>() -> bool;
+
+    fn get_comparable() -> Box<dyn Any>;
+}
+
+pub struct NoInitParamFlag {}
+
+/// A type which can be used as input to a [`System`].
+pub trait Input: 'static {}
+
+pub trait InputFilter
+{
+    type Filtered: FilteredInputs;
+}
+
+pub trait FilteredInputs
+{
+    type InOptions: OptionInputs;
+
+    fn into_in_options(self) -> Self::InOptions;
+}
+
+macro_rules! impl_filtered_inputs {
+    ($cnt: tt) => {
+        seq!(I in 0..$cnt {
+            impl<#(Input~I: Input,)*> FilteredInputs for (#(Input~I,)*) {
+                type InOptions = (#(Option<Input~I>,)*);
+
+                fn into_in_options(self) -> Self::InOptions {
+                    #![allow(clippy::unused_unit)]
+                    (#(Some(self.I),)*)
+                }
+            }
+        });
+    };
+}
+
+seq!(N in 0..4 {
+    impl_filtered_inputs!(N);
+});
+
+pub trait OptionInputs
+{
+    fn take<Input: 'static>(&mut self) -> TakeOptionInputResult<Input>;
+}
+
+macro_rules! impl_option_inputs {
+    ($cnt: tt) => {
+        seq!(I in 0..$cnt {
+            impl<#(Input~I: 'static,)*> OptionInputs for (#(Option<Input~I>,)*) {
+                fn take<Input: 'static>(&mut self) -> TakeOptionInputResult<Input> {
+                    #(
+                        if TypeId::of::<Input~I>() == TypeId::of::<Input>() {
+                            let input = match self.I.take() {
+                                Some(input) => ManuallyDrop::new(input),
+                                None => {
+                                    return TakeOptionInputResult::AlreadyTaken;
+                                }
+                            };
+
+                            return TakeOptionInputResult::Found(
+                                // SAFETY: It can be transmuted safely since it is the
+                                // same type and the type is 'static
+                                unsafe { transmute_copy(&input) }
+                            );
+                        }
+                    )*
+
+                    TakeOptionInputResult::NotFound
+                }
+            }
+        });
+    };
+}
+
+seq!(N in 0..4 {
+    impl_option_inputs!(N);
+});
+
+pub enum TakeOptionInputResult<Input>
+{
+    Found(Input),
+    NotFound,
+    AlreadyTaken,
+}
+
+include!(concat!(env!("OUT_DIR"), "/system_input_impls.rs"));
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);
+});
diff --git a/ecs/src/system/util.rs b/ecs/src/system/util.rs
new file mode 100644
index 0000000..9d04f1d
--- /dev/null
+++ b/ecs/src/system/util.rs
@@ -0,0 +1,13 @@
+macro_rules! check_params_are_compatible {
+    ($excluded_index: tt, $param: ident, $cnt: tt) => {
+        seq!(N in 0..=$cnt {
+            if N != $excluded_index {
+                if !$param::is_compatible::<TParam~N>() {
+                    panic!("Atleast two parameters are incompatible");
+                }
+            }
+        })
+    };
+}
+
+pub(crate) use check_params_are_compatible;