From caef866b4e3f87fd6ae2dd5b979a1fe1a1f3e5f2 Mon Sep 17 00:00:00 2001
From: HampusM <hampus@hampusmat.com>
Date: Sun, 3 Nov 2024 21:12:57 +0100
Subject: feat(ecs): add read-only query iterating

---
 ecs/src/query.rs | 95 ++++++++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 89 insertions(+), 6 deletions(-)

(limited to 'ecs/src/query.rs')

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;
 
@@ -38,6 +42,28 @@ where
     Comps: ComponentSequence,
     OptionsT: Options,
 {
+    /// 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>>
@@ -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
 {
-- 
cgit v1.2.3-18-g5258