feat(ecs): add support for optional query components

1 files changed, 34 insertions, 2 deletions

diff --git a/ecs/src/system.rs b/ecs/src/system.rs
index 3440a57..3c44148 100644
--- a/ecs/src/system.rs
+++ b/ecs/src/system.rs
@@ -1,4 +1,4 @@
-use std::any::Any;
+use std::any::{type_name, Any};
 use std::convert::Infallible;
 use std::fmt::Debug;
 use std::marker::PhantomData;
@@ -8,7 +8,7 @@ use std::ptr::addr_of;
 
 use seq_macro::seq;
 
-use crate::component::Component;
+use crate::component::{Component, FromOptionalComponent};
 use crate::lock::WriteGuard;
 use crate::system::util::check_params_are_compatible;
 use crate::tuple::{FilterElement as TupleFilterElement, With as TupleWith};
@@ -208,6 +208,38 @@ impl<'a, ComponentT: Component> ComponentRefMut<'a, ComponentT>
     }
 }
 
+impl<'component, ComponentT: Component> FromOptionalComponent<'component>
+    for ComponentRefMut<'component, ComponentT>
+{
+    fn from_optional_component(
+        inner: Option<WriteGuard<'component, Box<dyn Component>>>,
+    ) -> Self
+    {
+        Self {
+            inner: inner.unwrap_or_else(|| {
+                panic!(
+                    "Component {} was not found in entity",
+                    type_name::<ComponentT>()
+                );
+            }),
+            _ph: PhantomData,
+        }
+    }
+}
+
+impl<'comp, ComponentT> FromOptionalComponent<'comp>
+    for Option<ComponentRefMut<'comp, ComponentT>>
+where
+    ComponentT: Component,
+{
+    fn from_optional_component(
+        optional_component: Option<WriteGuard<'comp, Box<dyn Component>>>,
+    ) -> Self
+    {
+        optional_component.map(|component| ComponentRefMut::new(component))
+    }
+}
+
 impl<'a, ComponentT: Component> Deref for ComponentRefMut<'a, ComponentT>
 {
     type Target = ComponentT;