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

1 files changed, 40 insertions, 1 deletions

diff --git a/ecs/src/relationship.rs b/ecs/src/relationship.rs
index 7088613..d136db4 100644
--- a/ecs/src/relationship.rs
+++ b/ecs/src/relationship.rs
@@ -6,10 +6,11 @@ use ecs_macros::Component;
 use crate::component::storage::Storage as ComponentStorage;
 use crate::component::{
     Component,
+    FromLockedOptional as FromLockedOptionalComponent,
     FromOptional as FromOptionalComponent,
     FromOptionalMut as FromOptionalMutComponent,
 };
-use crate::lock::ReadGuard;
+use crate::lock::{Error as LockError, Lock, ReadGuard};
 use crate::system::{ComponentRef, ComponentRefMut};
 use crate::uid::{Kind as UidKind, Uid};
 use crate::World;
@@ -103,6 +104,25 @@ where
     }
 }
 
+impl<'rel_comp, Kind, ComponentT> FromLockedOptionalComponent<'rel_comp>
+    for RelationMut<'rel_comp, Kind, ComponentT>
+where
+    ComponentT: Component,
+{
+    fn from_locked_optional_component(
+        optional_component: Option<&'rel_comp crate::lock::Lock<Box<dyn Component>>>,
+        world: &'rel_comp World,
+    ) -> Result<Self, LockError>
+    {
+        Ok(Self::from_optional_mut_component(
+            optional_component
+                .map(|lock| lock.write_nonblock())
+                .transpose()?,
+            world,
+        ))
+    }
+}
+
 impl<'rel_comp, Kind, ComponentT> RelationMut<'rel_comp, Kind, ComponentT>
 where
     ComponentT: Component,
@@ -307,6 +327,25 @@ where
     }
 }
 
+impl<'rel_comp, Kind, ComponentT> FromLockedOptionalComponent<'rel_comp>
+    for Relation<'rel_comp, Kind, ComponentT>
+where
+    ComponentT: Component,
+{
+    fn from_locked_optional_component(
+        optional_component: Option<&'rel_comp Lock<Box<dyn Component>>>,
+        world: &'rel_comp World,
+    ) -> Result<Self, LockError>
+    {
+        Ok(Self::from_optional_component(
+            optional_component
+                .map(|lock| lock.read_nonblock())
+                .transpose()?,
+            world,
+        ))
+    }
+}
+
 impl<'rel_comp, Kind, ComponentT> Relation<'rel_comp, Kind, ComponentT>
 where
     ComponentT: Component,