Create an extension trait and iterator

Author: deontologician

crates/bevy_ecs/src/query/iter.rs

@@ -1,7 +1,9 @@
 use crate::{
     archetype::{ArchetypeId, Archetypes},
+    entity::Entity,
     query::{Fetch, FilterFetch, QueryState, ReadOnlyFetch, WorldQuery},
     storage::{TableId, Tables},
+    system::Query,
     world::World,
 };
 use std::mem::MaybeUninit;
@@ -550,3 +552,63 @@ where
         }
     }
 }
+
+pub struct WithQuery<'w, I, Q, F>
+where
+    Q: WorldQuery,
+    Q::Fetch: ReadOnlyFetch,
+    F: WorldQuery,
+    F::Fetch: FilterFetch,
+{
+    iter: I,
+    query: &'w Query<'w, Q, F>,
+}
+
+impl<'w, I, Q, F> Iterator for WithQuery<'w, I, Q, F>
+where
+    I: Iterator<Item = Entity>,
+    Q: WorldQuery,
+    Q::Fetch: ReadOnlyFetch,
+    F: WorldQuery,
+    F::Fetch: FilterFetch,
+{
+    type Item = <<Q as WorldQuery>::Fetch as Fetch<'w>>::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // if the entity iterator is done, this iterator is done
+        let mut entity = self.iter.next()?;
+        loop {
+            match self.query.get(entity).ok() {
+                Some(item) => return Some(item),
+                None => {
+                    entity = self.iter.next()?;
+                }
+            }
+        }
+    }
+}
+
+pub trait WithQueryExt: Sized {
+    fn with_query<'w, Q, F>(self, query: &'w Query<'w, Q, F>) -> WithQuery<'w, Self, Q, F>
+    where
+        Q: WorldQuery,
+        Q::Fetch: ReadOnlyFetch,
+        F: WorldQuery,
+        F::Fetch: FilterFetch;
+}
+
+impl<I> WithQueryExt for I
+where
+    I: Iterator<Item = Entity> + Sized,
+{
+    fn with_query<'w, Q, F>(self, query: &'w Query<'w, Q, F>) -> WithQuery<'w, Self, Q, F>
+    where
+        Self: Sized,
+        Q: WorldQuery,
+        Q::Fetch: ReadOnlyFetch,
+        F: WorldQuery,
+        F::Fetch: FilterFetch,
+    {
+        WithQuery { iter: self, query }
+    }
+}

crates/bevy_transform/src/components/children.rs

@@ -1,7 +1,5 @@
 use bevy_ecs::{
     entity::{Entity, EntityMap, MapEntities, MapEntitiesError},
-    prelude::Query,
-    query::{Fetch, FilterFetch, ReadOnlyFetch, WorldQuery},
     reflect::{ReflectComponent, ReflectMapEntities},
 };
 use bevy_reflect::Reflect;
@@ -31,35 +29,6 @@ impl Children {
     pub fn swap(&mut self, a_index: usize, b_index: usize) {
         self.0.swap(a_index, b_index);
     }
-
-    /// Returns an iterator over the query result for each child
-    ///
-    /// Example:
-    /// ```
-    /// fn example_system(
-    ///    parents: Query<&Children, With<Sprite>>,
-    ///    transforms: Query<&Transform, With<Sprite>>,
-    /// ) {
-    ///     for children in parents.iter() {
-    ///         for transform in children.with_query(transforms) {
-    ///             println!("transform: {:?}", transform);
-    ///         }
-    ///     }
-    /// }
-    /// ```
-    pub fn with_query<'a, Q, F>(
-        &'a self,
-        q: &'a Query<Q, F>,
-    ) -> impl Iterator<Item = <<Q as WorldQuery>::Fetch as Fetch<'a>>::Item>
-    where
-        Q: WorldQuery,
-        Q::Fetch: ReadOnlyFetch,
-        F: WorldQuery,
-        F::Fetch: FilterFetch,
-    {
-        self.iter()
-            .filter_map(move |child_id| q.get(*child_id).ok())
-    }
 }
 
 impl Deref for Children {