feat(ecs): allow control over component mutability in query

1 files changed, 68 insertions, 1 deletions

diff --git a/ecs/src/system.rs b/ecs/src/system.rs
index 992fd69..a810741 100644
--- a/ecs/src/system.rs
+++ b/ecs/src/system.rs
@@ -10,10 +10,11 @@ use seq_macro::seq;
 
 use crate::component::{
     Component,
+    FromLockedOptional as FromLockedOptionalComponent,
     FromOptional as FromOptionalComponent,
     FromOptionalMut as FromOptionalMutComponent,
 };
-use crate::lock::{ReadGuard, WriteGuard};
+use crate::lock::{Error as LockError, Lock, ReadGuard, WriteGuard};
 use crate::tuple::{ReduceElement as TupleReduceElement, Tuple};
 use crate::World;
 
@@ -222,6 +223,23 @@ impl<'component, ComponentT: Component> FromOptionalMutComponent<'component>
     }
 }
 
+impl<'component, ComponentT: Component> FromLockedOptionalComponent<'component>
+    for ComponentRefMut<'component, ComponentT>
+{
+    fn from_locked_optional_component(
+        optional_component: Option<&'component crate::lock::Lock<Box<dyn Component>>>,
+        world: &'component World,
+    ) -> Result<Self, LockError>
+    {
+        Ok(Self::from_optional_mut_component(
+            optional_component
+                .map(|lock| lock.write_nonblock())
+                .transpose()?,
+            world,
+        ))
+    }
+}
+
 impl<'comp, ComponentT> FromOptionalMutComponent<'comp>
     for Option<ComponentRefMut<'comp, ComponentT>>
 where
@@ -236,6 +254,22 @@ where
     }
 }
 
+impl<'comp, ComponentT> FromLockedOptionalComponent<'comp>
+    for Option<ComponentRefMut<'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(ComponentRefMut::new(lock.write_nonblock()?)))
+            .transpose()
+    }
+}
+
 impl<'a, ComponentT: Component> Deref for ComponentRefMut<'a, ComponentT>
 {
     type Target = ComponentT;
@@ -289,6 +323,23 @@ impl<'component, ComponentT: Component> FromOptionalComponent<'component>
     }
 }
 
+impl<'component, ComponentT: Component> FromLockedOptionalComponent<'component>
+    for ComponentRef<'component, ComponentT>
+{
+    fn from_locked_optional_component(
+        optional_component: Option<&'component crate::lock::Lock<Box<dyn Component>>>,
+        world: &'component World,
+    ) -> Result<Self, LockError>
+    {
+        Ok(Self::from_optional_component(
+            optional_component
+                .map(|lock| lock.read_nonblock())
+                .transpose()?,
+            world,
+        ))
+    }
+}
+
 impl<'comp, ComponentT> FromOptionalComponent<'comp>
     for Option<ComponentRef<'comp, ComponentT>>
 where
@@ -303,6 +354,22 @@ where
     }
 }
 
+impl<'comp, ComponentT> FromLockedOptionalComponent<'comp>
+    for Option<ComponentRef<'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(ComponentRef::new(lock.read_nonblock()?)))
+            .transpose()
+    }
+}
+
 impl<'a, ComponentT: Component> Deref for ComponentRef<'a, ComponentT>
 {
     type Target = ComponentT;