add iter_related hierarchy query extension

Author: jrobsonchase

crates/bevy_hierarchy/src/query_extension.rs

@@ -1,7 +1,8 @@
 use std::collections::VecDeque;
 
 use bevy_ecs::{
-    entity::Entity,
+    component::Component,
+    entity::{Entity, IterEntities},
     query::{QueryData, QueryFilter, WorldQuery},
     system::Query,
 };
@@ -30,9 +31,12 @@ pub trait HierarchyQueryExt<'w, 's, D: QueryData, F: QueryFilter> {
     /// }
     /// # bevy_ecs::system::assert_is_system(system);
     /// ```
-    fn iter_descendants(&'w self, entity: Entity) -> DescendantIter<'w, 's, D, F>
+    fn iter_descendants(&'w self, entity: Entity) -> RelatedIter<'w, 's, Children, D, F>
     where
-        D::ReadOnly: WorldQuery<Item<'w> = &'w Children>;
+        D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
+    {
+        self.iter_related(entity)
+    }
 
     /// Returns an [`Iterator`] of [`Entity`]s over all of `entity`s ancestors.
     ///
@@ -52,102 +56,102 @@ pub trait HierarchyQueryExt<'w, 's, D: QueryData, F: QueryFilter> {
     /// }
     /// # bevy_ecs::system::assert_is_system(system);
     /// ```
-    fn iter_ancestors(&'w self, entity: Entity) -> AncestorIter<'w, 's, D, F>
-    where
-        D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>;
-}
-
-impl<'w, 's, D: QueryData, F: QueryFilter> HierarchyQueryExt<'w, 's, D, F> for Query<'w, 's, D, F> {
-    fn iter_descendants(&'w self, entity: Entity) -> DescendantIter<'w, 's, D, F>
+    fn iter_ancestors(&'w self, entity: Entity) -> RelatedIter<'w, 's, Parent, D, F>
     where
-        D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
+        D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>,
     {
-        DescendantIter::new(self, entity)
+        self.iter_related(entity)
     }
 
-    fn iter_ancestors(&'w self, entity: Entity) -> AncestorIter<'w, 's, D, F>
+    /// Returns an [`Iterator`] of [`Entity`]'s over all of `entity`'s `C` relations.
+    ///
+    /// Can only be called on a [`Query`] of `C` (i.e. `Query<&C>`), where `C`
+    /// is some type implementing [`IterEntities`].
+    /// # Examples
+    /// ```
+    /// # use bevy_ecs::prelude::*;
+    /// # use bevy_hierarchy::prelude::*;
+    /// # #[derive(Component)]
+    /// # struct Marker;
+    /// fn system(query: Query<Entity, With<Marker>>, parent_query: Query<&Parent>) {
+    ///     let entity = query.single();
+    ///     for ancestor in parent_query.iter_related(entity) {
+    ///         // Do something!
+    ///     }
+    /// }
+    /// # bevy_ecs::system::assert_is_system(system);
+    /// ```
+    fn iter_related<C>(&'w self, entity: Entity) -> RelatedIter<'w, 's, C, D, F>
+    where
+        D::ReadOnly: WorldQuery<Item<'w> = &'w C>,
+        C: Component + IterEntities;
+}
+
+impl<'w, 's, D: QueryData, F: QueryFilter> HierarchyQueryExt<'w, 's, D, F> for Query<'w, 's, D, F> {
+    fn iter_related<C>(&'w self, entity: Entity) -> RelatedIter<'w, 's, C, D, F>
     where
-        D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>,
+        <D as QueryData>::ReadOnly: WorldQuery<Item<'w> = &'w C>,
+        C: Component + IterEntities,
     {
-        AncestorIter::new(self, entity)
+        RelatedIter::new(self, entity)
     }
 }
 
-/// An [`Iterator`] of [`Entity`]s over the descendants of an [`Entity`].
+/// An iterator over entities related via some component `C` that implements
+/// [`IterEntities`].
 ///
 /// Traverses the hierarchy breadth-first.
-pub struct DescendantIter<'w, 's, D: QueryData, F: QueryFilter>
+pub struct RelatedIter<'w, 's, C, D, F>
 where
-    D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
+    D: QueryData,
+    F: QueryFilter,
+    D::ReadOnly: WorldQuery<Item<'w> = &'w C>,
+    C: Component + IterEntities,
 {
-    children_query: &'w Query<'w, 's, D, F>,
+    related_query: &'w Query<'w, 's, D, F>,
     vecdeque: VecDeque<Entity>,
 }
 
-impl<'w, 's, D: QueryData, F: QueryFilter> DescendantIter<'w, 's, D, F>
+impl<'w, 's, C, D, F> RelatedIter<'w, 's, C, D, F>
 where
-    D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
+    D: QueryData,
+    F: QueryFilter,
+    D::ReadOnly: WorldQuery<Item<'w> = &'w C>,
+    C: Component + IterEntities,
 {
-    /// Returns a new [`DescendantIter`].
-    pub fn new(children_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
-        DescendantIter {
-            children_query,
-            vecdeque: children_query.get(entity).into_iter().flatten().collect(),
+    /// Create a new [`RelatedIter`].
+    pub fn new(related_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
+        RelatedIter {
+            related_query,
+            vecdeque: related_query
+                .get(entity)
+                .iter()
+                .flat_map(|t| t.iter_entities())
+                .collect(),
         }
     }
 }
 
-impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for DescendantIter<'w, 's, D, F>
+impl<'w, 's, C, D, F> Iterator for RelatedIter<'w, 's, C, D, F>
 where
-    D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
+    D: QueryData,
+    F: QueryFilter,
+    D::ReadOnly: WorldQuery<Item<'w> = &'w C>,
+    C: Component + IterEntities,
 {
     type Item = Entity;
 
     fn next(&mut self) -> Option<Self::Item> {
         let entity = self.vecdeque.pop_front()?;
 
-        if let Ok(children) = self.children_query.get(entity) {
-            self.vecdeque.extend(children);
+        if let Ok(children) = self.related_query.get(entity) {
+            self.vecdeque.extend(children.iter_entities());
         }
 
         Some(entity)
     }
 }
 
-/// An [`Iterator`] of [`Entity`]s over the ancestors of an [`Entity`].
-pub struct AncestorIter<'w, 's, D: QueryData, F: QueryFilter>
-where
-    D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>,
-{
-    parent_query: &'w Query<'w, 's, D, F>,
-    next: Option<Entity>,
-}
-
-impl<'w, 's, D: QueryData, F: QueryFilter> AncestorIter<'w, 's, D, F>
-where
-    D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>,
-{
-    /// Returns a new [`AncestorIter`].
-    pub fn new(parent_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
-        AncestorIter {
-            parent_query,
-            next: Some(entity),
-        }
-    }
-}
-
-impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for AncestorIter<'w, 's, D, F>
-where
-    D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>,
-{
-    type Item = Entity;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        self.next = self.parent_query.get(self.next?).ok().map(Parent::get);
-        self.next
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use bevy_ecs::{
@@ -196,4 +200,38 @@ mod tests {
 
         assert_eq!([&A(1), &A(0)], result.as_slice());
     }
+
+    #[test]
+    fn related_children_iter() {
+        let world = &mut World::new();
+
+        let [a, b, c, d] = std::array::from_fn(|i| world.spawn(A(i)).id());
+
+        world.entity_mut(a).add_children(&[b, c]);
+        world.entity_mut(c).add_children(&[d]);
+
+        let mut system_state = SystemState::<(Query<&Children>, Query<&A>)>::new(world);
+        let (children_query, a_query) = system_state.get(world);
+
+        let result: Vec<_> = a_query.iter_many(children_query.iter_related(a)).collect();
+
+        assert_eq!([&A(1), &A(2), &A(3)], result.as_slice());
+    }
+
+    #[test]
+    fn related_parent_iter() {
+        let world = &mut World::new();
+
+        let [a, b, c] = std::array::from_fn(|i| world.spawn(A(i)).id());
+
+        world.entity_mut(a).add_children(&[b]);
+        world.entity_mut(b).add_children(&[c]);
+
+        let mut system_state = SystemState::<(Query<&Parent>, Query<&A>)>::new(world);
+        let (parent_query, a_query) = system_state.get(world);
+
+        let result: Vec<_> = a_query.iter_many(parent_query.iter_related(c)).collect();
+
+        assert_eq!([&A(1), &A(0)], result.as_slice());
+    }
 }