feat(ecs): add support for custom inner ComponentIter iter

1 files changed, 43 insertions, 9 deletions

diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index 2f8b285..8c1ede5 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -1,7 +1,11 @@
 use std::marker::PhantomData;
 
 use crate::component::RefSequence as ComponentRefSequence;
-use crate::query::flexible::{Iter as FlexibleQueryIter, Query as FlexibleQuery};
+use crate::query::flexible::{
+    EntityHandle,
+    Iter as FlexibleQueryIter,
+    Query as FlexibleQuery,
+};
 use crate::query::options::Options;
 use crate::system::{Param as SystemParam, System};
 use crate::uid::Uid;
@@ -27,12 +31,13 @@ where
 {
     /// Iterates over the entities matching this query.
     #[must_use]
-    pub fn iter<'query>(&'query self) -> ComponentIter<'query, 'world, Comps>
+    pub fn iter<'query>(
+        &'query self,
+    ) -> ComponentIter<'query, 'world, Comps, FlexibleQueryIter<'query>>
     {
         #[cfg(feature = "debug")]
         tracing::debug!("Searching for {}", std::any::type_name::<Comps>());
 
-        #[allow(clippy::map_flatten)]
         ComponentIter {
             world: self.world,
             iter: self.inner.iter::<OptionsT>(),
@@ -40,6 +45,29 @@ where
         }
     }
 
+    /// Iterates over the entities matching this query using the iterator returned by
+    /// `func`.
+    ///
+    /// This function exists so that a custom [`EntityHandle`] iterator can be given to
+    /// [`ComponentIter`] without giving the user access to a reference to the [`World`].
+    #[must_use]
+    pub fn iter_with<'query, OutIter>(
+        &'query self,
+        func: impl FnOnce(FlexibleQueryIter<'query>) -> OutIter,
+    ) -> ComponentIter<'query, 'world, Comps, OutIter>
+    where
+        OutIter: Iterator<Item = EntityHandle<'query>>,
+    {
+        #[cfg(feature = "debug")]
+        tracing::debug!("Searching for {}", std::any::type_name::<Comps>());
+
+        ComponentIter {
+            world: self.world,
+            iter: func(self.inner.iter::<OptionsT>()),
+            comps_pd: PhantomData,
+        }
+    }
+
     /// Returns the UID of the entity at the given query iteration index.
     #[must_use]
     pub fn get_entity_uid(&self, entity_index: usize) -> Option<Uid>
@@ -63,7 +91,7 @@ where
     Comps: ComponentRefSequence + 'world,
     OptionsT: Options,
 {
-    type IntoIter = ComponentIter<'query, 'world, Comps>;
+    type IntoIter = ComponentIter<'query, 'world, Comps, FlexibleQueryIter<'query>>;
     type Item = Comps::Handles<'query>;
 
     fn into_iter(self) -> Self::IntoIter
@@ -95,27 +123,33 @@ where
     }
 }
 
-pub struct ComponentIter<'query, 'world, Comps>
+pub struct ComponentIter<'query, 'world, Comps, EntityHandleIter>
+where
+    EntityHandleIter: Iterator<Item = EntityHandle<'query>>,
 {
     world: &'world World,
-    iter: FlexibleQueryIter<'query>,
+    iter: EntityHandleIter,
     comps_pd: PhantomData<Comps>,
 }
 
-impl<'query, 'world, Comps> ComponentIter<'query, 'world, Comps>
+impl<'query, 'world, Comps, EntityHandleIter>
+    ComponentIter<'query, 'world, Comps, EntityHandleIter>
 where
     Comps: ComponentRefSequence + 'world,
+    EntityHandleIter: Iterator<Item = EntityHandle<'query>>,
     'world: 'query,
 {
-    pub(crate) fn new(world: &'world World, iter: FlexibleQueryIter<'query>) -> Self
+    pub(crate) fn new(world: &'world World, iter: EntityHandleIter) -> Self
     {
         Self { world, iter, comps_pd: PhantomData }
     }
 }
 
-impl<'query, 'world, Comps> Iterator for ComponentIter<'query, 'world, Comps>
+impl<'query, 'world, Comps, EntityHandleIter> Iterator
+    for ComponentIter<'query, 'world, Comps, EntityHandleIter>
 where
     Comps: ComponentRefSequence + 'world,
+    EntityHandleIter: Iterator<Item = EntityHandle<'query>>,
     'world: 'query,
 {
     type Item = Comps::Handles<'query>;