feat(ecs): add Query function to iterate with entity UIDs

1 files changed, 52 insertions, 1 deletions

diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index 8c1ede5..3ee7e5b 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -29,7 +29,8 @@ where
     Comps: ComponentRefSequence,
     OptionsT: Options,
 {
-    /// Iterates over the entities matching this query.
+    /// Iterates over the entities matching this query, the iterator item being the entity
+    /// components.
     #[must_use]
     pub fn iter<'query>(
         &'query self,
@@ -45,6 +46,23 @@ where
         }
     }
 
+    /// Iterates over the entities matching this query, the iterator item being the entity
+    /// [`Uid`] and the entity components.
+    #[must_use]
+    pub fn iter_with_euids<'query>(
+        &'query self,
+    ) -> ComponentAndEuidIter<'query, 'world, Comps, FlexibleQueryIter<'query>>
+    {
+        #[cfg(feature = "debug")]
+        tracing::debug!("Searching for {}", std::any::type_name::<Comps>());
+
+        ComponentAndEuidIter {
+            world: self.world,
+            iter: self.inner.iter::<OptionsT>(),
+            comps_pd: PhantomData,
+        }
+    }
+
     /// Iterates over the entities matching this query using the iterator returned by
     /// `func`.
     ///
@@ -165,3 +183,36 @@ where
         ))
     }
 }
+
+pub struct ComponentAndEuidIter<'query, 'world, Comps, EntityHandleIter>
+where
+    EntityHandleIter: Iterator<Item = EntityHandle<'query>>,
+{
+    world: &'world World,
+    iter: EntityHandleIter,
+    comps_pd: PhantomData<Comps>,
+}
+
+impl<'query, 'world, Comps, EntityHandleIter> Iterator
+    for ComponentAndEuidIter<'query, 'world, Comps, EntityHandleIter>
+where
+    Comps: ComponentRefSequence + 'world,
+    EntityHandleIter: Iterator<Item = EntityHandle<'query>>,
+    'world: 'query,
+{
+    type Item = (Uid, Comps::Handles<'query>);
+
+    fn next(&mut self) -> Option<Self::Item>
+    {
+        let entity_handle = self.iter.next()?;
+
+        Some((
+            entity_handle.uid(),
+            Comps::from_components(
+                entity_handle.components(),
+                |component_uid| entity_handle.get_component_index(component_uid),
+                self.world,
+            ),
+        ))
+    }
+}