feat(ecs): add Query iter_with_extra_comps function

1 files changed, 37 insertions, 1 deletions

diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index 253398d..f3318bd 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -8,7 +8,11 @@ use crate::component::storage::{
     EntityIter,
     Storage as ComponentStorage,
 };
-use crate::component::{Component, Sequence as ComponentSequence};
+use crate::component::{
+    Component,
+    Metadata as ComponentMetadata,
+    Sequence as ComponentSequence,
+};
 use crate::lock::{ReadGuard, WriteGuard};
 use crate::query::options::Options;
 use crate::system::{
@@ -78,6 +82,38 @@ where
         }
     }
 
+    /// Iterates over the entities matching this query and has the provided extra
+    /// component.
+    #[must_use]
+    pub fn iter_with_extra_comps(
+        &'world self,
+        extra_components: impl IntoIterator<Item = ComponentMetadata>,
+    ) -> ComponentIter<'world, Comps, QueryEntityIter<'world>>
+    {
+        #[cfg(feature = "debug")]
+        tracing::debug!(
+            "Searching for {} + extra components",
+            std::any::type_name::<Comps>()
+        );
+
+        #[allow(clippy::map_flatten)]
+        ComponentIter {
+            world: self.world,
+            entities: self
+                .component_storage
+                .find_entities(
+                    Comps::metadata()
+                        .into_iter()
+                        .chain(extra_components)
+                        .collect::<Vec<_>>(),
+                )
+                .map(Archetype::entities as ComponentIterMapFn)
+                .flatten()
+                .filter(|entity| OptionsT::entity_filter(entity.components())),
+            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>