1 files changed, 163 insertions, 31 deletions

diff --git a/ecs/src/component.rs b/ecs/src/component.rs
index 265eb9e..46fbf8a 100644
--- a/ecs/src/component.rs
+++ b/ecs/src/component.rs
@@ -1,5 +1,6 @@
 use std::any::{type_name, Any};
 use std::fmt::Debug;
+use std::ops::{Deref, DerefMut};
 
 use seq_macro::seq;
 
@@ -8,7 +9,14 @@ use crate::event::component::{
     Kind as ComponentEventKind,
     Removed as ComponentRemovedEvent,
 };
-use crate::lock::{Error as LockError, Lock, ReadGuard, WriteGuard};
+use crate::lock::{
+    Error as LockError,
+    Lock,
+    MappedReadGuard,
+    MappedWriteGuard,
+    ReadGuard,
+    WriteGuard,
+};
 use crate::system::Input as SystemInput;
 use crate::type_name::TypeName;
 use crate::uid::Uid;
@@ -26,11 +34,11 @@ pub trait Component: SystemInput + Any + TypeName
     where
         Self: Sized;
 
-    type RefMut<'component>: FromOptionalMut<'component> + FromLockedOptional<'component>
+    type HandleMut<'component>: FromLockedOptional<'component>
     where
         Self: Sized;
 
-    type Ref<'component>: FromOptional<'component> + FromLockedOptional<'component>
+    type Handle<'component>: FromLockedOptional<'component>
     where
         Self: Sized;
 
@@ -101,12 +109,12 @@ impl TypeName for Box<dyn Component>
 impl<ComponentT> Component for Option<ComponentT>
 where
     ComponentT: Component,
-    for<'a> Option<ComponentT::Ref<'a>>: FromOptional<'a> + FromLockedOptional<'a>,
-    for<'a> Option<ComponentT::RefMut<'a>>: FromOptionalMut<'a> + FromLockedOptional<'a>,
+    for<'a> Option<ComponentT::Handle<'a>>: FromLockedOptional<'a>,
+    for<'a> Option<ComponentT::HandleMut<'a>>: FromLockedOptional<'a>,
 {
     type Component = ComponentT;
-    type Ref<'component> = Option<ComponentT::Ref<'component>>;
-    type RefMut<'component> = Option<ComponentT::RefMut<'component>>;
+    type Handle<'component> = Option<ComponentT::Handle<'component>>;
+    type HandleMut<'component> = Option<ComponentT::HandleMut<'component>>;
 
     fn id() -> Uid
     {
@@ -182,7 +190,7 @@ where
     ComponentT: Component,
 {
     type Component = ComponentT;
-    type Handle<'component> = ComponentT::Ref<'component>;
+    type Handle<'component> = ComponentT::Handle<'component>;
 }
 
 impl<ComponentT> Ref for &mut ComponentT
@@ -190,7 +198,7 @@ where
     ComponentT: Component,
 {
     type Component = ComponentT;
-    type Handle<'component> = ComponentT::RefMut<'component>;
+    type Handle<'component> = ComponentT::HandleMut<'component>;
 }
 
 /// [`Component`] metadata.
@@ -220,22 +228,6 @@ impl Metadata
     }
 }
 
-pub trait FromOptionalMut<'comp>
-{
-    fn from_optional_mut_component(
-        optional_component: Option<WriteGuard<'comp, Box<dyn Component>>>,
-        world: &'comp World,
-    ) -> Self;
-}
-
-pub trait FromOptional<'comp>
-{
-    fn from_optional_component(
-        optional_component: Option<ReadGuard<'comp, Box<dyn Component>>>,
-        world: &'comp World,
-    ) -> Self;
-}
-
 pub trait FromLockedOptional<'comp>: Sized
 {
     /// Converts a reference to a optional locked boxed component to a instance of `Self`.
@@ -248,18 +240,158 @@ pub trait FromLockedOptional<'comp>: Sized
     ) -> Result<Self, LockError>;
 }
 
+#[derive(Debug)]
+pub struct Handle<'a, ComponentT: Component>
+{
+    inner: MappedReadGuard<'a, ComponentT>,
+}
+
+impl<'a, ComponentT: Component> Handle<'a, ComponentT>
+{
+    pub(crate) fn new(inner: ReadGuard<'a, Box<dyn Component>>) -> Self
+    {
+        Self {
+            inner: inner.map(|component| {
+                component.downcast_ref::<ComponentT>().unwrap_or_else(|| {
+                    panic!(
+                        "Cannot downcast component {} to type {}",
+                        component.type_name(),
+                        type_name::<ComponentT>()
+                    );
+                })
+            }),
+        }
+    }
+}
+
+impl<'component, ComponentT: Component> FromLockedOptional<'component>
+    for Handle<'component, ComponentT>
+{
+    fn from_locked_optional_component(
+        optional_component: Option<&'component crate::lock::Lock<Box<dyn Component>>>,
+        _world: &'component World,
+    ) -> Result<Self, LockError>
+    {
+        let component = optional_component.unwrap_or_else(|| {
+            panic!(
+                "Component {} was not found in entity",
+                type_name::<ComponentT>()
+            );
+        });
+
+        Ok(Self::new(component.read_nonblock()?))
+    }
+}
+
+impl<'comp, ComponentT> FromLockedOptional<'comp> for Option<Handle<'comp, ComponentT>>
+where
+    ComponentT: Component,
+{
+    fn from_locked_optional_component(
+        optional_component: Option<&'comp Lock<Box<dyn Component>>>,
+        _world: &'comp World,
+    ) -> Result<Self, LockError>
+    {
+        optional_component
+            .map(|lock| Ok(Handle::new(lock.read_nonblock()?)))
+            .transpose()
+    }
+}
+
+impl<ComponentT: Component> Deref for Handle<'_, ComponentT>
+{
+    type Target = ComponentT;
+
+    fn deref(&self) -> &Self::Target
+    {
+        &self.inner
+    }
+}
+
+#[derive(Debug)]
+pub struct HandleMut<'a, ComponentT: Component>
+{
+    inner: MappedWriteGuard<'a, ComponentT>,
+}
+
+impl<'a, ComponentT: Component> HandleMut<'a, ComponentT>
+{
+    pub(crate) fn new(inner: WriteGuard<'a, Box<dyn Component>>) -> Self
+    {
+        Self {
+            inner: inner.map(|component| {
+                let component_type_name = component.type_name();
+
+                component.downcast_mut::<ComponentT>().unwrap_or_else(|| {
+                    panic!(
+                        "Cannot downcast component {component_type_name} to type {}",
+                        type_name::<ComponentT>()
+                    );
+                })
+            }),
+        }
+    }
+}
+
+impl<'component, ComponentT: Component> FromLockedOptional<'component>
+    for HandleMut<'component, ComponentT>
+{
+    fn from_locked_optional_component(
+        optional_component: Option<&'component Lock<Box<dyn Component>>>,
+        _world: &'component World,
+    ) -> Result<Self, LockError>
+    {
+        let component = optional_component.unwrap_or_else(|| {
+            panic!(
+                "Component {} was not found in entity",
+                type_name::<ComponentT>()
+            );
+        });
+
+        Ok(Self::new(component.write_nonblock()?))
+    }
+}
+
+impl<'comp, ComponentT> FromLockedOptional<'comp> for Option<HandleMut<'comp, ComponentT>>
+where
+    ComponentT: Component,
+{
+    fn from_locked_optional_component(
+        optional_component: Option<&'comp Lock<Box<dyn Component>>>,
+        _world: &'comp World,
+    ) -> Result<Self, LockError>
+    {
+        optional_component
+            .map(|lock| Ok(HandleMut::new(lock.write_nonblock()?)))
+            .transpose()
+    }
+}
+
+impl<ComponentT: Component> Deref for HandleMut<'_, ComponentT>
+{
+    type Target = ComponentT;
+
+    fn deref(&self) -> &Self::Target
+    {
+        &self.inner
+    }
+}
+
+impl<ComponentT: Component> DerefMut for HandleMut<'_, ComponentT>
+{
+    fn deref_mut(&mut self) -> &mut Self::Target
+    {
+        &mut self.inner
+    }
+}
+
 macro_rules! inner {
     ($c: tt) => {
         seq!(I in 0..=$c {
             impl<#(IntoComp~I,)*> Sequence for (#(IntoComp~I,)*)
                 where
                 #(
-                    for<'comp> IntoComp~I: Into<
-                        Component: Component<
-                            RefMut<'comp>: FromOptionalMut<'comp>,
-                            Ref<'comp>: FromOptional<'comp>
-                        >
-                    >,
+                    for<'comp> IntoComp~I: Into<Component: Component>,
                 )*
             {
                 const COUNT: usize = $c + 1;