1 files changed, 96 insertions, 10 deletions

diff --git a/ecs/src/component.rs b/ecs/src/component.rs
index 7a61f39..604f54d 100644
--- a/ecs/src/component.rs
+++ b/ecs/src/component.rs
@@ -1,4 +1,4 @@
-use std::any::{Any, TypeId};
+use std::any::{type_name, Any, TypeId};
 use std::fmt::Debug;
 
 use seq_macro::seq;
@@ -11,6 +11,15 @@ pub mod local;
 
 pub trait Component: SystemInput + Any + TypeName
 {
+    /// The component type in question. Will usually be `Self`
+    type Component
+    where
+        Self: Sized;
+
+    type RefMut<'component>
+    where
+        Self: Sized;
+
     fn drop_last(&self) -> bool;
 
     #[doc(hidden)]
@@ -54,6 +63,43 @@ impl TypeName for Box<dyn Component>
     }
 }
 
+impl<ComponentT> Component for Option<ComponentT>
+where
+    ComponentT: Component,
+{
+    type Component = ComponentT;
+    type RefMut<'component> = Option<ComponentRefMut<'component, ComponentT>>;
+
+    fn drop_last(&self) -> bool
+    {
+        self.as_ref()
+            .map(|component| component.drop_last())
+            .unwrap_or_default()
+    }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any
+    {
+        self
+    }
+
+    fn as_any(&self) -> &dyn Any
+    {
+        self
+    }
+}
+
+impl<ComponentT> TypeName for Option<ComponentT>
+where
+    ComponentT: Component,
+{
+    fn type_name(&self) -> &'static str
+    {
+        type_name::<Self>()
+    }
+}
+
+impl<ComponentT> SystemInput for Option<ComponentT> where ComponentT: Component {}
+
 /// A sequence of components.
 pub trait Sequence
 {
@@ -63,18 +109,60 @@ pub trait Sequence
 
     fn into_vec(self) -> Vec<Box<dyn Component>>;
 
-    fn type_ids() -> Vec<TypeId>;
+    fn type_ids() -> Vec<(TypeId, IsOptional)>;
 
     fn from_components<'component>(
         components: impl Iterator<Item = &'component Lock<Box<dyn Component>>>,
     ) -> Self::Refs<'component>;
 }
 
+/// Whether or not a `Component` is optional.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum IsOptional
+{
+    Yes,
+    No,
+}
+
+impl From<bool> for IsOptional
+{
+    fn from(is_optional: bool) -> Self
+    {
+        if is_optional {
+            return IsOptional::Yes;
+        }
+
+        IsOptional::No
+    }
+}
+
+/// Returns whether the given component type is a optional component.
+///
+/// Will return `true` if the component is a [`Option`].
+pub fn is_optional<ComponentT: Component>() -> bool
+{
+    if TypeId::of::<ComponentT>() == TypeId::of::<Option<ComponentT::Component>>() {
+        return true;
+    }
+
+    false
+}
+
+pub trait FromOptionalComponent<'comp>
+{
+    fn from_optional_component(
+        optional_component: Option<WriteGuard<'comp, Box<dyn Component>>>,
+    ) -> Self;
+}
+
 macro_rules! inner {
     ($c: tt) => {
         seq!(I in 0..=$c {
-            impl<#(Comp~I: Component,)*> Sequence for (#(Comp~I,)*) {
-                type Refs<'component> = (#(ComponentRefMut<'component, Comp~I>,)*)
+            impl<#(Comp~I: Component,)*> Sequence for (#(Comp~I,)*)
+                where
+                    #(for<'comp> Comp~I::RefMut<'comp>: FromOptionalComponent<'comp>,)*
+            {
+                type Refs<'component> = (#(Comp~I::RefMut<'component>,)*)
                     where Self: 'component;
 
                 fn into_vec(self) -> Vec<Box<dyn Component>>
@@ -82,11 +170,11 @@ macro_rules! inner {
                     Vec::from_iter([#(Box::new(self.I) as Box<dyn Component>,)*])
                 }
 
-                fn type_ids() -> Vec<TypeId>
+                fn type_ids() -> Vec<(TypeId, IsOptional)>
                 {
                     vec![
                         #(
-                            TypeId::of::<Comp~I>(),
+                            (TypeId::of::<Comp~I>(), is_optional::<Comp~I>().into()),
                         )*
                     ]
                 }
@@ -105,7 +193,7 @@ macro_rules! inner {
                             .expect("Failed to acquire read-write component lock");
 
                         #(
-                            if comp_ref.is::<Comp~I>() {
+                            if comp_ref.is::<Comp~I::Component>() {
                                 comp_~I = Some(comp_ref);
                                 continue;
                             }
@@ -113,9 +201,7 @@ macro_rules! inner {
                     }
 
                     (#(
-                        ComponentRefMut::new(
-                            comp_~I.unwrap(),
-                        ),
+                        Comp~I::RefMut::from_optional_component(comp_~I),
                     )*)
                 }
             }