feat(ecs): add read-only query iterating

4 files changed, 396 insertions, 50 deletions

diff --git a/ecs/src/component.rs b/ecs/src/component.rs
index 5b1e7f5..0506346 100644
--- a/ecs/src/component.rs
+++ b/ecs/src/component.rs
@@ -3,8 +3,8 @@ use std::fmt::Debug;
 
 use seq_macro::seq;
 
-use crate::lock::WriteGuard;
-use crate::system::{ComponentRefMut, Input as SystemInput};
+use crate::lock::{ReadGuard, WriteGuard};
+use crate::system::{ComponentRef, ComponentRefMut, Input as SystemInput};
 use crate::type_name::TypeName;
 use crate::{EntityComponent, World};
 
@@ -23,6 +23,10 @@ pub trait Component: SystemInput + Any + TypeName
     where
         Self: Sized;
 
+    type Ref<'component>
+    where
+        Self: Sized;
+
     /// Returns the ID of this component's type.
     fn id(&self) -> Id;
 
@@ -77,6 +81,7 @@ where
     ComponentT: Component,
 {
     type Component = ComponentT;
+    type Ref<'component> = Option<ComponentRef<'component, ComponentT>>;
     type RefMut<'component> = Option<ComponentRefMut<'component, ComponentT>>;
 
     fn id(&self) -> Id
@@ -133,6 +138,10 @@ impl Id
 /// A sequence of components.
 pub trait Sequence
 {
+    type MutRefs<'component>
+    where
+        Self: 'component;
+
     type Refs<'component>
     where
         Self: 'component;
@@ -141,9 +150,20 @@ pub trait Sequence
 
     fn metadata() -> Vec<Metadata>;
 
+    fn from_components_mut<'component>(
+        components: impl Iterator<Item = &'component EntityComponent>,
+        world: &'component World,
+        lock_component: fn(
+            entity_component: &EntityComponent,
+        ) -> WriteGuard<'_, Box<dyn Component>>,
+    ) -> Self::MutRefs<'component>;
+
     fn from_components<'component>(
         components: impl Iterator<Item = &'component EntityComponent>,
         world: &'component World,
+        lock_component: fn(
+            entity_component: &EntityComponent,
+        ) -> ReadGuard<'_, Box<dyn Component>>,
     ) -> Self::Refs<'component>;
 }
 
@@ -200,10 +220,18 @@ pub fn is_optional<ComponentT: Component>() -> bool
     false
 }
 
+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<WriteGuard<'comp, Box<dyn Component>>>,
+        optional_component: Option<ReadGuard<'comp, Box<dyn Component>>>,
         world: &'comp World,
     ) -> Self;
 }
@@ -213,9 +241,13 @@ macro_rules! inner {
         seq!(I in 0..=$c {
             impl<#(Comp~I: Component,)*> Sequence for (#(Comp~I,)*)
                 where
-                    #(for<'comp> Comp~I::RefMut<'comp>: FromOptional<'comp>,)*
+                    #(for<'comp> Comp~I::RefMut<'comp>: FromOptionalMut<'comp>,)*
+                    #(for<'comp> Comp~I::Ref<'comp>: FromOptional<'comp>,)*
             {
-                type Refs<'component> = (#(Comp~I::RefMut<'component>,)*)
+                type MutRefs<'component> = (#(Comp~I::RefMut<'component>,)*)
+                    where Self: 'component;
+
+                type Refs<'component> = (#(Comp~I::Ref<'component>,)*)
                     where Self: 'component;
 
                 fn into_vec(self) -> Vec<Box<dyn Component>>
@@ -235,13 +267,43 @@ macro_rules! inner {
                     ]
                 }
 
+                fn from_components_mut<'component>(
+                    components: impl Iterator<Item = &'component EntityComponent>,
+                    world: &'component World,
+                    lock_component: fn(
+                        entity_component: &EntityComponent,
+                    ) -> WriteGuard<'_, Box<dyn Component>>,
+                ) -> Self::MutRefs<'component>
+                {
+                    #(
+                        let mut comp_~I: Option<WriteGuard<Box<dyn Component>>> = None;
+                    )*
+
+                    for comp in components {
+                        #(
+                            if comp.id == Id::of::<Comp~I::Component>() {
+                                comp_~I = Some(lock_component(comp));
+                                continue;
+                            }
+                        )*
+                    }
+
+                    (#(
+                        Comp~I::RefMut::from_optional_mut_component(comp_~I, world),
+                    )*)
+                }
+
                 fn from_components<'component>(
                     components: impl Iterator<Item = &'component EntityComponent>,
                     world: &'component World,
+                    lock_component: fn(
+                        entity_component: &EntityComponent,
+                    ) -> ReadGuard<'_, Box<dyn Component>>,
                 ) -> Self::Refs<'component>
                 {
+
                     #(
-                        let mut comp_~I: Option<WriteGuard<Box<dyn Component>>> = None;
+                        let mut comp_~I: Option<ReadGuard<Box<dyn Component>>> = None;
                     )*
 
                     for comp in components {
@@ -254,7 +316,7 @@ macro_rules! inner {
                     }
 
                     (#(
-                        Comp~I::RefMut::from_optional_component(comp_~I, world),
+                        Comp~I::Ref::from_optional_component(comp_~I, world),
                     )*)
                 }
             }
@@ -265,18 +327,3 @@ macro_rules! inner {
 seq!(C in 0..=64 {
     inner!(C);
 });
-
-fn lock_component(
-    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
-            );
-        })
-}
diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index d0fa872..e7fd7a6 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -10,16 +10,20 @@ use crate::component::storage::{
     EntityIter,
     Storage as ComponentStorage,
 };
-use crate::component::{Metadata as ComponentMetadata, Sequence as ComponentSequence};
+use crate::component::{
+    Component,
+    Metadata as ComponentMetadata,
+    Sequence as ComponentSequence,
+};
 use crate::entity::Uid as EntityUid;
-use crate::lock::ReadGuard;
+use crate::lock::{ReadGuard, WriteGuard};
 use crate::query::options::Options;
 use crate::system::{
     NoInitParamFlag as NoInitSystemParamFlag,
     Param as SystemParam,
     System,
 };
-use crate::World;
+use crate::{EntityComponent, World};
 
 pub mod options;
 
@@ -40,6 +44,28 @@ where
 {
     /// Iterates over the entities matching this query.
     #[must_use]
+    pub fn iter_mut(
+        &'world self,
+    ) -> ComponentIterMut<'world, Comps, QueryEntityIter<'world>>
+    {
+        #[cfg(feature = "debug")]
+        tracing::debug!("Searching for {}", std::any::type_name::<Comps>());
+
+        #[allow(clippy::map_flatten)]
+        ComponentIterMut {
+            world: self.world,
+            entities: self
+                .component_storage
+                .find_entities(Comps::metadata())
+                .map(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(&'world self) -> ComponentIter<'world, Comps, QueryEntityIter<'world>>
     {
         #[cfg(feature = "debug")]
@@ -91,12 +117,12 @@ where
     Comps: ComponentSequence,
     OptionsT: Options,
 {
-    type IntoIter = ComponentIter<'world, Comps, QueryEntityIter<'world>>;
-    type Item = Comps::Refs<'world>;
+    type IntoIter = ComponentIterMut<'world, Comps, QueryEntityIter<'world>>;
+    type Item = Comps::MutRefs<'world>;
 
     fn into_iter(self) -> Self::IntoIter
     {
-        self.iter()
+        self.iter_mut()
     }
 }
 
@@ -152,6 +178,47 @@ type QueryEntityIter<'world> = Filter<
     ComponentIterFilterFn,
 >;
 
+pub struct ComponentIterMut<'world, Comps, EntityIter>
+where
+    EntityIter: Iterator<Item = &'world ArchetypeEntity>,
+{
+    world: &'world World,
+    entities: EntityIter,
+    comps_pd: PhantomData<Comps>,
+}
+
+impl<'world, Comps, EntityIter> Iterator for ComponentIterMut<'world, Comps, EntityIter>
+where
+    Comps: ComponentSequence + 'world,
+    EntityIter: Iterator<Item = &'world ArchetypeEntity>,
+{
+    type Item = Comps::MutRefs<'world>;
+
+    fn next(&mut self) -> Option<Self::Item>
+    {
+        Some(Comps::from_components_mut(
+            self.entities.next()?.components().iter(),
+            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<'world, Comps, EntityIter>
 where
     EntityIter: Iterator<Item = &'world ArchetypeEntity>,
@@ -173,10 +240,26 @@ where
         Some(Comps::from_components(
             self.entities.next()?.components().iter(),
             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
+            );
+        })
+}
+
 #[derive(Debug)]
 struct QueryComponentIds
 {
diff --git a/ecs/src/relationship.rs b/ecs/src/relationship.rs
index 0fe776e..4db29da 100644
--- a/ecs/src/relationship.rs
+++ b/ecs/src/relationship.rs
@@ -4,12 +4,13 @@ use std::marker::PhantomData;
 use crate::component::storage::Storage as ComponentStorage;
 use crate::component::{
     Component,
-    FromOptional as ComponentFromOptional,
+    FromOptional as FromOptionalComponent,
+    FromOptionalMut as FromOptionalMutComponent,
     Id as ComponentId,
 };
 use crate::entity::Uid as EntityUid;
 use crate::lock::ReadGuard;
-use crate::system::{ComponentRefMut, Input as SystemInput};
+use crate::system::{ComponentRef, ComponentRefMut, Input as SystemInput};
 use crate::type_name::TypeName;
 use crate::World;
 
@@ -52,7 +53,8 @@ where
     ComponentT: Component,
 {
     type Component = Self;
-    type RefMut<'component> = Relation<'component, Kind, ComponentT>;
+    type Ref<'component> = Relation<'component, Kind, ComponentT>;
+    type RefMut<'component> = RelationMut<'component, Kind, ComponentT>;
 
     fn id(&self) -> ComponentId
     {
@@ -87,7 +89,7 @@ where
     }
 }
 
-pub struct Relation<'rel_comp, Kind, ComponentT>
+pub struct RelationMut<'rel_comp, Kind, ComponentT>
 where
     Kind: 'static,
     ComponentT: Component,
@@ -96,12 +98,12 @@ where
     relationship_comp: ComponentRefMut<'rel_comp, Relationship<Kind, ComponentT>>,
 }
 
-impl<'rel_comp, Kind, ComponentT> ComponentFromOptional<'rel_comp>
-    for Relation<'rel_comp, Kind, ComponentT>
+impl<'rel_comp, Kind, ComponentT> FromOptionalMutComponent<'rel_comp>
+    for RelationMut<'rel_comp, Kind, ComponentT>
 where
     ComponentT: Component,
 {
-    fn from_optional_component(
+    fn from_optional_mut_component(
         optional_component: Option<
             crate::lock::WriteGuard<'rel_comp, Box<dyn Component>>,
         >,
@@ -109,7 +111,7 @@ where
     ) -> Self
     {
         let relationship_comp =
-            ComponentRefMut::<Relationship<Kind, ComponentT>>::from_optional_component(
+            ComponentRefMut::<Relationship<Kind, ComponentT>>::from_optional_mut_component(
                 optional_component,
                 world,
             );
@@ -129,7 +131,7 @@ where
     }
 }
 
-impl<'rel_comp, Kind, ComponentT> Relation<'rel_comp, Kind, ComponentT>
+impl<'rel_comp, Kind, ComponentT> RelationMut<'rel_comp, Kind, ComponentT>
 where
     ComponentT: Component,
 {
@@ -235,19 +237,19 @@ where
 
     /// Returns a iterator of the components of the targets of this relationship.
     #[must_use]
-    pub fn iter(&self) -> TargetComponentIter<'_, 'rel_comp, Kind, ComponentT>
+    pub fn iter(&self) -> TargetComponentIterMut<'_, 'rel_comp, Kind, ComponentT>
     {
-        TargetComponentIter { relation: self, index: 0 }
+        TargetComponentIterMut { relation: self, index: 0 }
     }
 }
 
 impl<'relationship, 'rel_comp, Kind, ComponentT> IntoIterator
-    for &'relationship Relation<'rel_comp, Kind, ComponentT>
+    for &'relationship RelationMut<'rel_comp, Kind, ComponentT>
 where
     'relationship: 'rel_comp,
     ComponentT: Component,
 {
-    type IntoIter = TargetComponentIter<'relationship, 'rel_comp, Kind, ComponentT>;
+    type IntoIter = TargetComponentIterMut<'relationship, 'rel_comp, Kind, ComponentT>;
     type Item = ComponentRefMut<'rel_comp, ComponentT>;
 
     fn into_iter(self) -> Self::IntoIter
@@ -257,17 +259,17 @@ where
 }
 
 /// Iterator of the components of the targets of a relationship.
-pub struct TargetComponentIter<'relationship, 'rel_comp, Kind, ComponentT>
+pub struct TargetComponentIterMut<'relationship, 'rel_comp, Kind, ComponentT>
 where
     Kind: 'static,
     ComponentT: Component,
 {
-    relation: &'relationship Relation<'rel_comp, Kind, ComponentT>,
+    relation: &'relationship RelationMut<'rel_comp, Kind, ComponentT>,
     index: usize,
 }
 
 impl<'relationship, 'rel_comp, Kind, ComponentT> Iterator
-    for TargetComponentIter<'relationship, 'rel_comp, Kind, ComponentT>
+    for TargetComponentIterMut<'relationship, 'rel_comp, Kind, ComponentT>
 where
     'relationship: 'rel_comp,
     Kind: 'static,
@@ -291,3 +293,154 @@ enum SingleOrMultiple<Value>
     Single(Value),
     Multiple(Vec<Value>),
 }
+
+pub struct Relation<'rel_comp, Kind, ComponentT>
+where
+    Kind: 'static,
+    ComponentT: Component,
+{
+    component_storage_lock: ReadGuard<'static, ComponentStorage>,
+    relationship_comp: ComponentRef<'rel_comp, Relationship<Kind, ComponentT>>,
+}
+
+impl<'rel_comp, Kind, ComponentT> FromOptionalComponent<'rel_comp>
+    for Relation<'rel_comp, Kind, ComponentT>
+where
+    ComponentT: Component,
+{
+    fn from_optional_component(
+        optional_component: Option<ReadGuard<'rel_comp, Box<dyn Component>>>,
+        world: &'rel_comp World,
+    ) -> Self
+    {
+        let relationship_comp =
+            ComponentRef::<Relationship<Kind, ComponentT>>::from_optional_component(
+                optional_component,
+                world,
+            );
+
+        let component_storage_lock = world
+            .data
+            .component_storage
+            .read_nonblock()
+            .expect("Failed to aquire read-only component storage lock");
+
+        Self {
+            relationship_comp,
+            // SAFETY: The component lock is not used for longer than the original
+            // lifetime
+            component_storage_lock: unsafe { component_storage_lock.upgrade_lifetime() },
+        }
+    }
+}
+
+impl<'rel_comp, Kind, ComponentT> Relation<'rel_comp, Kind, ComponentT>
+where
+    ComponentT: Component,
+{
+    /// Returns the component of the target at the specified index.
+    ///
+    /// # Panics
+    /// Will panic if the entity does not exist in the archetype it belongs to. This
+    /// should hopefully never happend.
+    #[must_use]
+    pub fn get(&self, index: usize) -> Option<ComponentRef<'_, ComponentT>>
+    {
+        let target = self.get_target(index)?;
+
+        let archetype = self.component_storage_lock.get_entity_archetype(*target)?;
+
+        let entity = archetype
+            .get_entity(*target)
+            .expect("Target entity is gone from archetype");
+
+        let component_index =
+            archetype.get_index_for_component(&ComponentId::of::<ComponentT>())?;
+
+        let component = ComponentRef::new(
+            entity
+                .get_component(component_index)?
+                .component
+                .read_nonblock()
+                .unwrap_or_else(|_| {
+                    panic!(
+                        "Failed to aquire read-write lock of component {}",
+                        type_name::<ComponentT>()
+                    )
+                }),
+        );
+
+        Some(component)
+    }
+
+    /// Returns a reference to the target at the specified index.
+    #[must_use]
+    pub fn get_target(&self, index: usize) -> Option<&EntityUid>
+    {
+        match &self.relationship_comp.entity_uid {
+            SingleOrMultiple::Single(entity_uid) if index == 0 => Some(entity_uid),
+            SingleOrMultiple::Multiple(entity_uids) => entity_uids.get(index),
+            SingleOrMultiple::Single(_) => None,
+        }
+    }
+
+    #[must_use]
+    pub fn target_count(&self) -> usize
+    {
+        match &self.relationship_comp.entity_uid {
+            SingleOrMultiple::Single(_) => 1,
+            SingleOrMultiple::Multiple(entity_uids) => entity_uids.len(),
+        }
+    }
+
+    /// Returns a iterator of the components of the targets of this relationship.
+    #[must_use]
+    pub fn iter(&self) -> TargetComponentIter<'_, 'rel_comp, Kind, ComponentT>
+    {
+        TargetComponentIter { relation: self, index: 0 }
+    }
+}
+
+impl<'relationship, 'rel_comp, Kind, ComponentT> IntoIterator
+    for &'relationship Relation<'rel_comp, Kind, ComponentT>
+where
+    'relationship: 'rel_comp,
+    ComponentT: Component,
+{
+    type IntoIter = TargetComponentIter<'relationship, 'rel_comp, Kind, ComponentT>;
+    type Item = ComponentRef<'rel_comp, ComponentT>;
+
+    fn into_iter(self) -> Self::IntoIter
+    {
+        self.iter()
+    }
+}
+
+/// Iterator of the components of the targets of a relationship.
+pub struct TargetComponentIter<'relationship, 'rel_comp, Kind, ComponentT>
+where
+    Kind: 'static,
+    ComponentT: Component,
+{
+    relation: &'relationship Relation<'rel_comp, Kind, ComponentT>,
+    index: usize,
+}
+
+impl<'relationship, 'rel_comp, Kind, ComponentT> Iterator
+    for TargetComponentIter<'relationship, 'rel_comp, Kind, ComponentT>
+where
+    'relationship: 'rel_comp,
+    Kind: 'static,
+    ComponentT: Component,
+{
+    type Item = ComponentRef<'rel_comp, ComponentT>;
+
+    fn next(&mut self) -> Option<Self::Item>
+    {
+        let index = self.index;
+
+        self.index += 1;
+
+        self.relation.get(index)
+    }
+}
diff --git a/ecs/src/system.rs b/ecs/src/system.rs
index 7c0e454..afb5aac 100644
--- a/ecs/src/system.rs
+++ b/ecs/src/system.rs
@@ -7,8 +7,12 @@ use std::panic::{RefUnwindSafe, UnwindSafe};
 
 use seq_macro::seq;
 
-use crate::component::{Component, FromOptional as FromOptionalComponent};
-use crate::lock::WriteGuard;
+use crate::component::{
+    Component,
+    FromOptional as FromOptionalComponent,
+    FromOptionalMut as FromOptionalMutComponent,
+};
+use crate::lock::{ReadGuard, WriteGuard};
 use crate::system::util::check_params_are_compatible;
 use crate::tuple::{ReduceElement as TupleReduceElement, With as TupleWith};
 use crate::World;
@@ -214,10 +218,10 @@ impl<'a, ComponentT: Component> ComponentRefMut<'a, ComponentT>
     }
 }
 
-impl<'component, ComponentT: Component> FromOptionalComponent<'component>
+impl<'component, ComponentT: Component> FromOptionalMutComponent<'component>
     for ComponentRefMut<'component, ComponentT>
 {
-    fn from_optional_component(
+    fn from_optional_mut_component(
         inner: Option<WriteGuard<'component, Box<dyn Component>>>,
         _world: &'component World,
     ) -> Self
@@ -234,12 +238,12 @@ impl<'component, ComponentT: Component> FromOptionalComponent<'component>
     }
 }
 
-impl<'comp, ComponentT> FromOptionalComponent<'comp>
+impl<'comp, ComponentT> FromOptionalMutComponent<'comp>
     for Option<ComponentRefMut<'comp, ComponentT>>
 where
     ComponentT: Component,
 {
-    fn from_optional_component(
+    fn from_optional_mut_component(
         optional_component: Option<WriteGuard<'comp, Box<dyn Component>>>,
         _world: &'comp World,
     ) -> Self
@@ -265,3 +269,62 @@ impl<'a, ComponentT: Component> DerefMut for ComponentRefMut<'a, ComponentT>
         self.inner.downcast_mut().unwrap()
     }
 }
+
+#[derive(Debug)]
+pub struct ComponentRef<'a, ComponentT: Component>
+{
+    inner: ReadGuard<'a, Box<dyn Component>>,
+    _ph: PhantomData<ComponentT>,
+}
+
+impl<'a, ComponentT: Component> ComponentRef<'a, ComponentT>
+{
+    pub(crate) fn new(inner: ReadGuard<'a, Box<dyn Component>>) -> Self
+    {
+        Self { inner, _ph: PhantomData }
+    }
+}
+
+impl<'component, ComponentT: Component> FromOptionalComponent<'component>
+    for ComponentRef<'component, ComponentT>
+{
+    fn from_optional_component(
+        inner: Option<ReadGuard<'component, Box<dyn Component>>>,
+        _world: &'component World,
+    ) -> 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<ComponentRef<'comp, ComponentT>>
+where
+    ComponentT: Component,
+{
+    fn from_optional_component(
+        optional_component: Option<ReadGuard<'comp, Box<dyn Component>>>,
+        _world: &'comp World,
+    ) -> Self
+    {
+        optional_component.map(|component| ComponentRef::new(component))
+    }
+}
+
+impl<'a, ComponentT: Component> Deref for ComponentRef<'a, ComponentT>
+{
+    type Target = ComponentT;
+
+    fn deref(&self) -> &Self::Target
+    {
+        self.inner.downcast_ref().unwrap()
+    }
+}