Add methods to Query<&Children> and Query<&Parent> to iterate over descendants and ancestors (#6185)

# Objective
Add methods to `Query<&Children>` and `Query<&Parent>` to iterate over descendants and ancestors, respectively.

## Changelog

* Added extension trait for `Query` in `bevy_hierarchy`, `HierarchyQueryExt`
* Added method `iter_descendants` to `Query<&Children>` via `HierarchyQueryExt` for iterating over the descendants of an entity.
* Added method `iter_ancestors` to `Query<&Parent>` via `HierarchyQueryExt` for iterating over the ancestors of an entity.

Co-authored-by: devil-ira <justthecooldude@gmail.com>

Author: tim-blackbird

crates/bevy_hierarchy/src/components/children.rs

@@ -10,7 +10,12 @@ use core::slice;
 use smallvec::SmallVec;
 use std::ops::Deref;
 
-/// Contains references to the child entities of this entity
+/// Contains references to the child entities of this entity.
+///
+/// See [`HierarchyQueryExt`] for hierarchy related methods on [`Query`].
+///
+/// [`HierarchyQueryExt`]: crate::query_extension::HierarchyQueryExt
+/// [`Query`]: bevy_ecs::system::Query
 #[derive(Component, Debug, Reflect)]
 #[reflect(Component, MapEntities)]
 pub struct Children(pub(crate) SmallVec<[Entity; 8]>);

crates/bevy_hierarchy/src/components/parent.rs

@@ -9,6 +9,11 @@ use std::ops::Deref;
 
 /// Holds a reference to the parent entity of this entity.
 /// This component should only be present on entities that actually have a parent entity.
+///
+/// See [`HierarchyQueryExt`] for hierarchy related methods on [`Query`].
+///
+/// [`HierarchyQueryExt`]: crate::query_extension::HierarchyQueryExt
+/// [`Query`]: bevy_ecs::system::Query
 #[derive(Component, Debug, Eq, PartialEq, Reflect)]
 #[reflect(Component, MapEntities, PartialEq)]
 pub struct Parent(pub(crate) Entity);

crates/bevy_hierarchy/src/lib.rs

@@ -19,11 +19,15 @@ pub use events::*;
 mod valid_parent_check_plugin;
 pub use valid_parent_check_plugin::*;
 
+mod query_extension;
+pub use query_extension::*;
+
 #[doc(hidden)]
 pub mod prelude {
     #[doc(hidden)]
     pub use crate::{
-        child_builder::*, components::*, hierarchy::*, HierarchyPlugin, ValidParentCheckPlugin,
+        child_builder::*, components::*, hierarchy::*, query_extension::*, HierarchyPlugin,
+        ValidParentCheckPlugin,
     };
 }
 

crates/bevy_hierarchy/src/query_extension.rs

@@ -0,0 +1,206 @@
+use std::collections::VecDeque;
+
+use bevy_ecs::{
+    entity::Entity,
+    query::{ReadOnlyWorldQuery, WorldQuery, WorldQueryGats},
+    system::Query,
+};
+
+use crate::{Children, Parent};
+
+/// An extension trait for [`Query`] that adds hierarchy related methods.
+pub trait HierarchyQueryExt<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> {
+    /// Returns an [`Iterator`] of [`Entity`]s over all of `entity`s descendants.
+    ///
+    /// Can only be called on a [`Query`] of [`Children`] (i.e. `Query<&Children>`).
+    ///
+    /// Traverses the hierarchy breadth-first.
+    ///
+    /// # Examples
+    /// ```
+    /// # use bevy_ecs::prelude::*;
+    /// # use bevy_hierarchy::prelude::*;
+    /// # #[derive(Component)]
+    /// # struct Marker;
+    /// fn system(query: Query<Entity, With<Marker>>, children_query: Query<&Children>) {
+    ///     let entity = query.single();
+    ///     for descendant in children_query.iter_descendants(entity) {
+    ///         // Do something!
+    ///     }
+    /// }
+    /// # bevy_ecs::system::assert_is_system(system);
+    /// ```
+    fn iter_descendants(&'w self, entity: Entity) -> DescendantIter<'w, 's, Q, F>
+    where
+        Q::ReadOnly: WorldQueryGats<'w, Item = &'w Children>;
+
+    /// Returns an [`Iterator`] of [`Entity`]s over all of `entity`s ancestors.
+    ///
+    /// Can only be called on a [`Query`] of [`Parent`] (i.e. `Query<&Parent>`).
+    ///
+    /// # 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_ancestors(entity) {
+    ///         // Do something!
+    ///     }
+    /// }
+    /// # bevy_ecs::system::assert_is_system(system);
+    /// ```
+    fn iter_ancestors(&'w self, entity: Entity) -> AncestorIter<'w, 's, Q, F>
+    where
+        Q::ReadOnly: WorldQueryGats<'w, Item = &'w Parent>;
+}
+
+impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> HierarchyQueryExt<'w, 's, Q, F>
+    for Query<'w, 's, Q, F>
+{
+    fn iter_descendants(&'w self, entity: Entity) -> DescendantIter<'w, 's, Q, F>
+    where
+        Q::ReadOnly: WorldQueryGats<'w, Item = &'w Children>,
+    {
+        DescendantIter::new(self, entity)
+    }
+
+    fn iter_ancestors(&'w self, entity: Entity) -> AncestorIter<'w, 's, Q, F>
+    where
+        Q::ReadOnly: WorldQueryGats<'w, Item = &'w Parent>,
+    {
+        AncestorIter::new(self, entity)
+    }
+}
+
+/// An [`Iterator`] of [`Entity`]s over the descendants of an [`Entity`].
+///
+/// Traverses the hierarchy breadth-first.
+pub struct DescendantIter<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery>
+where
+    Q::ReadOnly: WorldQueryGats<'w, Item = &'w Children>,
+{
+    children_query: &'w Query<'w, 's, Q, F>,
+    vecdeque: VecDeque<Entity>,
+}
+
+impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> DescendantIter<'w, 's, Q, F>
+where
+    Q::ReadOnly: WorldQueryGats<'w, Item = &'w Children>,
+{
+    /// Returns a new [`DescendantIter`].
+    pub fn new(children_query: &'w Query<'w, 's, Q, F>, entity: Entity) -> Self {
+        DescendantIter {
+            children_query,
+            vecdeque: children_query
+                .get(entity)
+                .into_iter()
+                .flatten()
+                .copied()
+                .collect(),
+        }
+    }
+}
+
+impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> Iterator for DescendantIter<'w, 's, Q, F>
+where
+    Q::ReadOnly: WorldQueryGats<'w, Item = &'w Children>,
+{
+    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);
+        }
+
+        Some(entity)
+    }
+}
+
+/// An [`Iterator`] of [`Entity`]s over the ancestors of an [`Entity`].
+pub struct AncestorIter<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery>
+where
+    Q::ReadOnly: WorldQueryGats<'w, Item = &'w Parent>,
+{
+    parent_query: &'w Query<'w, 's, Q, F>,
+    next: Option<Entity>,
+}
+
+impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> AncestorIter<'w, 's, Q, F>
+where
+    Q::ReadOnly: WorldQueryGats<'w, Item = &'w Parent>,
+{
+    /// Returns a new [`AncestorIter`].
+    pub fn new(parent_query: &'w Query<'w, 's, Q, F>, entity: Entity) -> Self {
+        AncestorIter {
+            parent_query,
+            next: Some(entity),
+        }
+    }
+}
+
+impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> Iterator for AncestorIter<'w, 's, Q, F>
+where
+    Q::ReadOnly: WorldQueryGats<'w, Item = &'w Parent>,
+{
+    type Item = Entity;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.next = self.parent_query.get(self.next?).ok().map(|p| p.get());
+        self.next
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use bevy_ecs::{
+        prelude::Component,
+        system::{Query, SystemState},
+        world::World,
+    };
+
+    use crate::{query_extension::HierarchyQueryExt, BuildWorldChildren, Children, Parent};
+
+    #[derive(Component, PartialEq, Debug)]
+    struct A(usize);
+
+    #[test]
+    fn descendant_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).push_children(&[b, c]);
+        world.entity_mut(c).push_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_descendants(a))
+            .collect();
+
+        assert_eq!([&A(1), &A(2), &A(3)], result.as_slice());
+    }
+
+    #[test]
+    fn ancestor_iter() {
+        let world = &mut World::new();
+
+        let [a, b, c] = std::array::from_fn(|i| world.spawn(A(i)).id());
+
+        world.entity_mut(a).push_children(&[b]);
+        world.entity_mut(b).push_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_ancestors(c)).collect();
+
+        assert_eq!([&A(1), &A(0)], result.as_slice());
+    }
+}

examples/3d/update_gltf_scene.rs

@@ -51,7 +51,7 @@ fn move_scene_entities(
 ) {
     for moved_scene_entity in &moved_scene {
         let mut offset = 0.;
-        iter_hierarchy(moved_scene_entity, &children, &mut |entity| {
+        for entity in children.iter_descendants(moved_scene_entity) {
             if let Ok(mut transform) = transforms.get_mut(entity) {
                 transform.translation = Vec3::new(
                     offset * time.elapsed_seconds().sin() / 20.,
@@ -60,15 +60,6 @@ fn move_scene_entities(
                 );
                 offset += 1.0;
             }
-        });
-    }
-}
-
-fn iter_hierarchy(entity: Entity, children_query: &Query<&Children>, f: &mut impl FnMut(Entity)) {
-    (f)(entity);
-    if let Ok(children) = children_query.get(entity) {
-        for child in children.iter().copied() {
-            iter_hierarchy(child, children_query, f);
         }
     }
 }