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

6 files changed, 207 insertions, 182 deletions

diff --git a/ecs/src/component.rs b/ecs/src/component.rs
index 1abb23c..4d37cb8 100644
--- a/ecs/src/component.rs
+++ b/ecs/src/component.rs
@@ -8,7 +8,7 @@ use crate::event::component::{
     Kind as ComponentEventKind,
     Removed as ComponentRemovedEvent,
 };
-use crate::lock::{ReadGuard, WriteGuard};
+use crate::lock::{Error as LockError, Lock, ReadGuard, WriteGuard};
 use crate::system::{ComponentRef, ComponentRefMut, Input as SystemInput};
 use crate::type_name::TypeName;
 use crate::uid::Uid;
@@ -26,11 +26,11 @@ pub trait Component: SystemInput + Any + TypeName
     where
         Self: Sized;
 
-    type RefMut<'component>: FromOptionalMut<'component>
+    type RefMut<'component>: FromOptionalMut<'component> + FromLockedOptional<'component>
     where
         Self: Sized;
 
-    type Ref<'component>: FromOptional<'component>
+    type Ref<'component>: FromOptional<'component> + FromLockedOptional<'component>
     where
         Self: Sized;
 
@@ -163,14 +163,6 @@ impl<ComponentT> SystemInput for Option<ComponentT> where ComponentT: Component
 /// A sequence of components.
 pub trait Sequence
 {
-    type MutRefs<'component>
-    where
-        Self: 'component;
-
-    type Refs<'component>
-    where
-        Self: 'component;
-
     /// The number of components in this component sequence.
     const COUNT: usize;
 
@@ -181,20 +173,43 @@ pub trait Sequence
     fn added_event_ids() -> Vec<Uid>;
 
     fn removed_event_ids() -> Vec<Uid>;
+}
 
-    fn from_components_mut<'component>(
-        components: &'component [EntityComponent],
-        component_index_lookup: impl Fn(Uid) -> Option<usize>,
-        world: &'component World,
-        lock_component: fn(&EntityComponent) -> WriteGuard<'_, Box<dyn Component>>,
-    ) -> Self::MutRefs<'component>;
+/// A sequence of references (immutable or mutable) to components.
+pub trait RefSequence
+{
+    type Handles<'component>;
+
+    fn metadata() -> impl Array<Metadata>;
 
     fn from_components<'component>(
         components: &'component [EntityComponent],
         component_index_lookup: impl Fn(Uid) -> Option<usize>,
         world: &'component World,
-        lock_component: fn(&EntityComponent) -> ReadGuard<'_, Box<dyn Component>>,
-    ) -> Self::Refs<'component>;
+    ) -> Self::Handles<'component>;
+}
+
+/// A mutable or immutable reference to a component.
+pub trait Ref
+{
+    type Component: Component;
+    type Handle<'component>: FromLockedOptional<'component>;
+}
+
+impl<ComponentT> Ref for &ComponentT
+where
+    ComponentT: Component,
+{
+    type Component = ComponentT;
+    type Handle<'component> = ComponentT::Ref<'component>;
+}
+
+impl<ComponentT> Ref for &mut ComponentT
+where
+    ComponentT: Component,
+{
+    type Component = ComponentT;
+    type Handle<'component> = ComponentT::RefMut<'component>;
 }
 
 /// [`Component`] metadata.
@@ -261,6 +276,14 @@ pub trait FromOptional<'comp>
     ) -> Self;
 }
 
+pub trait FromLockedOptional<'comp>: Sized
+{
+    fn from_locked_optional_component(
+        optional_component: Option<&'comp Lock<Box<dyn Component>>>,
+        world: &'comp World,
+    ) -> Result<Self, LockError>;
+}
+
 macro_rules! inner {
     ($c: tt) => {
         seq!(I in 0..=$c {
@@ -269,12 +292,6 @@ macro_rules! inner {
                     #(for<'comp> Comp~I::RefMut<'comp>: FromOptionalMut<'comp>,)*
                     #(for<'comp> Comp~I::Ref<'comp>: FromOptional<'comp>,)*
             {
-                type MutRefs<'component> = (#(Comp~I::RefMut<'component>,)*)
-                    where Self: 'component;
-
-                type Refs<'component> = (#(Comp~I::Ref<'component>,)*)
-                    where Self: 'component;
-
                 const COUNT: usize = $c + 1;
 
                 fn into_vec(self) -> Vec<Box<dyn Component>>
@@ -307,46 +324,40 @@ macro_rules! inner {
                         #(ComponentRemovedEvent::<Comp~I>::id(),)*
                     ]
                 }
+            }
 
-                #[inline]
-                fn from_components_mut<'component>(
-                    components: &'component [EntityComponent],
-                    component_index_lookup: impl Fn(Uid) -> Option<usize>,
-                    world: &'component World,
-                    lock_component:
-                        fn(&EntityComponent) -> WriteGuard<'_, Box<dyn Component>>,
-                ) -> Self::MutRefs<'component>
+            impl<#(CompRef~I: Ref,)*> RefSequence for (#(CompRef~I,)*)
+            {
+                type Handles<'component> = (#(CompRef~I::Handle<'component>,)*);
+
+                fn metadata() -> impl Array<Metadata>
                 {
-                    (#(
-                        Comp~I::RefMut::from_optional_mut_component(
-                            component_index_lookup(Comp~I::Component::id())
-                                .and_then(|component_index| {
-                                    components.get(component_index)
-                                        .map(lock_component)
-                                }),
-                            world
-                        ),
-                    )*)
+                    [#(
+                        Metadata {
+                            id: CompRef~I::Component::id(),
+                            is_optional: CompRef~I::Component::is_optional(),
+                        },
+                    )*]
                 }
 
-                #[inline]
                 fn from_components<'component>(
                     components: &'component [EntityComponent],
                     component_index_lookup: impl Fn(Uid) -> Option<usize>,
                     world: &'component World,
-                    lock_component:
-                        fn(&EntityComponent) -> ReadGuard<'_, Box<dyn Component>>,
-                ) -> Self::Refs<'component>
+                ) -> Self::Handles<'component>
                 {
                     (#(
-                        Comp~I::Ref::from_optional_component(
-                            component_index_lookup(Comp~I::Component::id())
-                                .and_then(|component_index| {
-                                    components.get(component_index)
-                                        .map(lock_component)
-                                }),
-                            world
-                        ),
+                        CompRef~I::Handle::from_locked_optional_component(
+                            component_index_lookup(CompRef~I::Component::id())
+                                .and_then(|component_index| components.get(component_index))
+                                .map(|component| &component.component),
+                            world,
+                        ).unwrap_or_else(|err| {
+                            panic!(
+                                "Taking component {} lock failed: {err}",
+                                type_name::<CompRef~I::Component>()
+                            );
+                        }),
                     )*)
                 }
             }
@@ -360,9 +371,6 @@ seq!(C in 0..=16 {
 
 impl Sequence for ()
 {
-    type MutRefs<'component> = ();
-    type Refs<'component> = ();
-
     const COUNT: usize = 0;
 
     fn into_vec(self) -> Vec<Box<dyn Component>>
@@ -384,25 +392,22 @@ impl Sequence for ()
     {
         Vec::new()
     }
+}
 
-    #[inline]
-    fn from_components_mut<'component>(
-        _components: &'component [EntityComponent],
-        _component_index_lookup: impl Fn(Uid) -> Option<usize>,
-        _world: &'component World,
-        _lock_component: fn(&EntityComponent) -> WriteGuard<'_, Box<dyn Component>>,
-    ) -> Self::MutRefs<'component>
+impl RefSequence for ()
+{
+    type Handles<'component> = ();
+
+    fn metadata() -> impl Array<Metadata>
     {
-        ()
+        []
     }
 
-    #[inline]
     fn from_components<'component>(
         _components: &'component [EntityComponent],
         _component_index_lookup: impl Fn(Uid) -> Option<usize>,
         _world: &'component World,
-        _lock_component: fn(&EntityComponent) -> ReadGuard<'_, Box<dyn Component>>,
-    ) -> Self::Refs<'component>
+    ) -> Self::Handles<'component>
     {
         ()
     }
diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index ae5fd19..d217525 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -15,6 +15,7 @@ use crate::component::{
     Component,
     IsOptional as ComponentIsOptional,
     Metadata as ComponentMetadata,
+    RefSequence as ComponentRefSequence,
     Sequence as ComponentSequence,
 };
 use crate::entity::CREATE_STATIC_ENTITIES;
@@ -175,7 +176,7 @@ impl World
 
     pub fn query<Comps, OptionsT>(&self) -> Query<Comps, OptionsT>
     where
-        Comps: ComponentSequence,
+        Comps: ComponentRefSequence,
         OptionsT: QueryOptions,
     {
         Query::new(self)
@@ -236,7 +237,7 @@ impl World
     fn query_and_run_systems(&self, phase_euid: Uid)
     {
         let system_comps_query =
-            self.query::<(SystemComponent, Relationship<DependsOn, Phase>), ()>();
+            self.query::<(&SystemComponent, &Relationship<DependsOn, Phase>), ()>();
 
         let system_iter = system_comps_query.iter().filter(|(_, phase_rel)| {
             phase_rel
@@ -254,7 +255,7 @@ impl World
 
     fn perform_child_phases(&self, parent_phase_euid: Uid)
     {
-        let phase_query = self.query::<(Phase, Relationship<ChildOf, Phase>), ()>();
+        let phase_query = self.query::<(&Phase, &Relationship<ChildOf, Phase>), ()>();
 
         for (index, (_, phase_rel)) in phase_query.iter().enumerate() {
             if !phase_rel
@@ -277,7 +278,7 @@ impl World
     fn perform_phases(&self)
     {
         let phase_query =
-            self.query::<(Phase,), Not<With<Relationship<ChildOf, Phase>>>>();
+            self.query::<(&Phase,), Not<With<Relationship<ChildOf, Phase>>>>();
 
         for (index, (_,)) in phase_query.iter().enumerate() {
             let child_phase_euid = phase_query
@@ -440,7 +441,7 @@ impl World
     fn emit_event_by_id(&self, event_id: Uid)
     {
         for (system,) in self
-            .query::<(SystemComponent,), ()>()
+            .query::<(&SystemComponent,), ()>()
             .iter_with_extra_comps([ComponentMetadata {
                 id: event_id,
                 is_optional: ComponentIsOptional::No,
diff --git a/ecs/src/lock.rs b/ecs/src/lock.rs
index fbc6842..135f654 100644
--- a/ecs/src/lock.rs
+++ b/ecs/src/lock.rs
@@ -28,7 +28,7 @@ where
     pub fn read_nonblock(&self) -> Result<ReadGuard<Value>, Error>
     {
         let guard = self.inner.try_read().or_else(|err| match err {
-            TryLockError::WouldBlock => Err(Error::Unavailable),
+            TryLockError::WouldBlock => Err(Error::ReadUnavailable),
             TryLockError::Poisoned(poison_err) => Ok(poison_err.into_inner()),
         })?;
 
@@ -45,7 +45,7 @@ where
     pub fn write_nonblock(&self) -> Result<WriteGuard<Value>, Error>
     {
         let guard = self.inner.try_write().or_else(|err| match err {
-            TryLockError::WouldBlock => Err(Error::Unavailable),
+            TryLockError::WouldBlock => Err(Error::WriteUnavailable),
             TryLockError::Poisoned(poison_err) => Ok(poison_err.into_inner()),
         })?;
 
@@ -69,8 +69,11 @@ where
 #[derive(Debug, thiserror::Error)]
 pub enum Error
 {
-    #[error("Lock is unavailable")]
-    Unavailable,
+    #[error("Lock is unavailable for reading")]
+    ReadUnavailable,
+
+    #[error("Lock is unavailable for writing")]
+    WriteUnavailable,
 }
 
 #[derive(Debug)]
diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index 68964eb..53d4b74 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -9,22 +9,21 @@ use crate::component::storage::{
     Storage as ComponentStorage,
 };
 use crate::component::{
-    Component,
     Metadata as ComponentMetadata,
-    Sequence as ComponentSequence,
+    RefSequence as ComponentRefSequence,
 };
-use crate::lock::{ReadGuard, WriteGuard};
+use crate::lock::ReadGuard;
 use crate::query::options::Options;
 use crate::system::{Param as SystemParam, System};
 use crate::uid::Uid;
-use crate::{EntityComponent, World};
+use crate::World;
 
 pub mod options;
 
 #[derive(Debug)]
 pub struct Query<'world, Comps, OptionsT = ()>
 where
-    Comps: ComponentSequence,
+    Comps: ComponentRefSequence,
 {
     world: &'world World,
     component_storage: ReadGuard<'world, ComponentStorage>,
@@ -33,38 +32,11 @@ where
 
 impl<'world, Comps, OptionsT> Query<'world, Comps, OptionsT>
 where
-    Comps: ComponentSequence,
+    Comps: ComponentRefSequence,
     OptionsT: Options,
 {
     /// Iterates over the entities matching this query.
     #[must_use]
-    pub fn iter_mut<'query>(
-        &'query self,
-    ) -> ComponentIterMut<'query, 'world, Comps, QueryEntityIter<'query>>
-    {
-        #[cfg(feature = "debug")]
-        tracing::debug!("Searching for {}", std::any::type_name::<Comps>());
-
-        #[allow(clippy::map_flatten)]
-        ComponentIterMut {
-            world: self.world,
-            entities: self
-                .component_storage
-                .iter_archetypes_with_comps(Comps::metadata())
-                .map(
-                    (|archetype| {
-                        repeat_n(archetype, archetype.entity_cnt())
-                            .zip(archetype.entities())
-                    }) as ComponentIterMapFn,
-                )
-                .flatten()
-                .filter(|(_, entity)| OptionsT::entity_filter(entity.components())),
-            comps_pd: PhantomData,
-        }
-    }
-
-    /// Iterates over the entities matching this query.
-    #[must_use]
     pub fn iter<'query>(
         &'query self,
     ) -> ComponentIter<'query, 'world, Comps, QueryEntityIter<'query>>
@@ -158,21 +130,21 @@ where
 impl<'query, 'world, Comps, OptionsT> IntoIterator
     for &'query Query<'world, Comps, OptionsT>
 where
-    Comps: ComponentSequence + 'world,
+    Comps: ComponentRefSequence + 'world,
     OptionsT: Options,
 {
-    type IntoIter = ComponentIterMut<'query, 'world, Comps, QueryEntityIter<'query>>;
-    type Item = Comps::MutRefs<'query>;
+    type IntoIter = ComponentIter<'query, 'world, Comps, QueryEntityIter<'query>>;
+    type Item = Comps::Handles<'query>;
 
     fn into_iter(self) -> Self::IntoIter
     {
-        self.iter_mut()
+        self.iter()
     }
 }
 
 impl<'world, Comps, OptionsT> SystemParam<'world> for Query<'world, Comps, OptionsT>
 where
-    Comps: ComponentSequence,
+    Comps: ComponentRefSequence,
     OptionsT: Options,
 {
     type Input = ();
@@ -204,52 +176,6 @@ type QueryEntityIter<'query> = Filter<
     ComponentIterFilterFn,
 >;
 
-pub struct ComponentIterMut<'query, 'world, Comps, EntityIter>
-where
-    EntityIter: Iterator<Item = (&'query Archetype, &'query ArchetypeEntity)>,
-{
-    world: &'world World,
-    entities: EntityIter,
-    comps_pd: PhantomData<Comps>,
-}
-
-impl<'query, 'world, Comps, EntityIter> Iterator
-    for ComponentIterMut<'query, 'world, Comps, EntityIter>
-where
-    Comps: ComponentSequence + 'world,
-    EntityIter: Iterator<Item = (&'query Archetype, &'query ArchetypeEntity)>,
-    'world: 'query,
-{
-    type Item = Comps::MutRefs<'query>;
-
-    fn next(&mut self) -> Option<Self::Item>
-    {
-        let (archetype, entity) = self.entities.next()?;
-
-        Some(Comps::from_components_mut(
-            entity.components(),
-            |component_uid| archetype.get_index_for_component(component_uid),
-            self.world,
-            lock_component_rw,
-        ))
-    }
-}
-
-fn lock_component_rw(
-    entity_component: &EntityComponent,
-) -> WriteGuard<'_, Box<dyn Component>>
-{
-    entity_component
-        .component
-        .write_nonblock()
-        .unwrap_or_else(|_| {
-            panic!(
-                "Failed to acquire read-write lock to component {}",
-                entity_component.name
-            );
-        })
-}
-
 pub struct ComponentIter<'query, 'world, Comps, EntityIter>
 where
     EntityIter: Iterator<Item = (&'query Archetype, &'query ArchetypeEntity)>,
@@ -262,11 +188,11 @@ where
 impl<'query, 'world, Comps, EntityIter> Iterator
     for ComponentIter<'query, 'world, Comps, EntityIter>
 where
-    Comps: ComponentSequence + 'world,
+    Comps: ComponentRefSequence + 'world,
     EntityIter: Iterator<Item = (&'query Archetype, &'query ArchetypeEntity)>,
     'world: 'query,
 {
-    type Item = Comps::Refs<'query>;
+    type Item = Comps::Handles<'query>;
 
     fn next(&mut self) -> Option<Self::Item>
     {
@@ -276,22 +202,6 @@ where
             entity.components(),
             |component_uid| archetype.get_index_for_component(component_uid),
             self.world,
-            lock_component_ro,
         ))
     }
 }
-
-fn lock_component_ro(
-    entity_component: &EntityComponent,
-) -> ReadGuard<'_, Box<dyn Component>>
-{
-    entity_component
-        .component
-        .read_nonblock()
-        .unwrap_or_else(|_| {
-            panic!(
-                "Failed to acquire read-write lock to component {}",
-                entity_component.name
-            );
-        })
-}
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,
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;