Add more tools for traversing hierarchies (#15627)

# Objective

- Working with hierarchies in Bevy is far too tedious due to a lack of
helper functions.
- This is the first half of #15609. 

## Solution

Extend
[`HierarchyQueryExt`](https://docs.rs/bevy/latest/bevy/hierarchy/trait.HierarchyQueryExt)
with the following methods:

- `parent`
- `children`
- `root_parent`
- `iter_leaves`
- `iter_siblings`
- `iter_descendants_depth_first`

I've opted to make both `iter_leaves` and `iter_siblings` collect the
list of matching Entities for now, rather that operate by reference like
the existing `iter_descendants`. This was simpler, and in the case of
`iter_siblings` especially, the number of matching entities is likely to
be much smaller.

I've kept the generics in the type signature however, so we can go back
and optimize that freely without a breaking change whenever we want.

## Testing

I've added some basic testing, but they're currently failing. If you'd
like to help, I'd welcome suggestions or a PR to my PR over the weekend
<3

---------

Co-authored-by: Viktor Gustavsson <villor94@gmail.com>
Co-authored-by: poopy <gonesbird@gmail.com>
Co-authored-by: Christian Hughes <9044780+ItsDoot@users.noreply.github.com>

crates/bevy_hierarchy/src/query_extension.rs

@@ -5,16 +5,54 @@ use bevy_ecs::{
  query::{QueryData, QueryFilter, WorldQuery},
  system::Query,
 };
+use smallvec::SmallVec;
 
 use crate::{Children, Parent};
 
 /// An extension trait for [`Query`] that adds hierarchy related methods.
 pub trait HierarchyQueryExt<'w, 's, D: QueryData, F: QueryFilter> {
+ /// Returns the parent [`Entity`] of the given `entity`, if any.
+ fn parent(&'w self, entity: Entity) -> Option<Entity>
+ where
+ D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>;
+
+ /// Returns a slice over the [`Children`] of the given `entity`.
+ ///
+ /// This may be empty if the `entity` has no children.
+ fn children(&'w self, entity: Entity) -> &'w [Entity]
+ where
+ D::ReadOnly: WorldQuery<Item<'w> = &'w Children>;
+
+ /// Returns the topmost ancestor of the given `entity`.
+ ///
+ /// This may be the entity itself if it has no parent.
+ fn root_ancestor(&'w self, entity: Entity) -> Entity
+ where
+ D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>;
+
+ /// Returns an [`Iterator`] of [`Entity`]s over the leaves of the hierarchy that are underneath this `entity`.
+ ///
+ /// Only entities which have no children are considered leaves.
+ /// This will not include the entity itself, and will not include any entities which are not descendants of the entity,
+ /// even if they are leaves in the same hierarchical tree.
+ ///
+ /// Traverses the hierarchy depth-first.
+ fn iter_leaves(&'w self, entity: Entity) -> impl Iterator<Item = Entity> + 'w
+ where
+ D::ReadOnly: WorldQuery<Item<'w> = &'w Children>;
+
+ /// Returns an [`Iterator`] of [`Entity`]s over the `entity`s immediate siblings, who share the same parent.
+ ///
+ /// The entity itself is not included in the iterator.
+ fn iter_siblings(&'w self, entity: Entity) -> impl Iterator<Item = Entity>
+ where
+ D::ReadOnly: WorldQuery<Item<'w> = (Option<&'w Parent>, Option<&'w Children>)>;
+
  /// 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.
+ /// Traverses the hierarchy breadth-first and does not include the entity itself.
  ///
  /// # Examples
  /// ```
@@ -34,8 +72,21 @@ pub trait HierarchyQueryExt<'w, 's, D: QueryData, F: QueryFilter> {
  where
  D::ReadOnly: WorldQuery<Item<'w> = &'w Children>;
 
+ /// 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>`).
+ ///
+ /// This is a depth-first alternative to [`HierarchyQueryExt::iter_descendants`].
+ fn iter_descendants_depth_first(
+ &'w self,
+ entity: Entity,
+ ) -> DescendantDepthFirstIter<'w, 's, D, F>
+ where
+ D::ReadOnly: WorldQuery<Item<'w> = &'w Children>;
+
  /// Returns an [`Iterator`] of [`Entity`]s over all of `entity`s ancestors.
  ///
+ /// Does not include the entity itself.
  /// Can only be called on a [`Query`] of [`Parent`] (i.e. `Query<&Parent>`).
  ///
  /// # Examples
@@ -58,13 +109,76 @@ pub trait HierarchyQueryExt<'w, 's, D: QueryData, F: QueryFilter> {
 }
 
 impl<'w, 's, D: QueryData, F: QueryFilter> HierarchyQueryExt<'w, 's, D, F> for Query<'w, 's, D, F> {
+ fn parent(&'w self, entity: Entity) -> Option<Entity>
+ where
+ <D as QueryData>::ReadOnly: WorldQuery<Item<'w> = &'w Parent>,
+ {
+ self.get(entity).map(Parent::get).ok()
+ }
+
+ fn children(&'w self, entity: Entity) -> &'w [Entity]
+ where
+ <D as QueryData>::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
+ {
+ self.get(entity)
+ .map_or(&[] as &[Entity], |children| children)
+ }
+
+ fn root_ancestor(&'w self, entity: Entity) -> Entity
+ where
+ <D as QueryData>::ReadOnly: WorldQuery<Item<'w> = &'w Parent>,
+ {
+ // Recursively search up the tree until we're out of parents
+ match self.get(entity) {
+ Ok(parent) => self.root_ancestor(parent.get()),
+ Err(_) => entity,
+ }
+ }
+
+ fn iter_leaves(&'w self, entity: Entity) -> impl Iterator<Item = Entity>
+ where
+ <D as QueryData>::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
+ {
+ self.iter_descendants_depth_first(entity).filter(|entity| {
+ self.get(*entity)
+ // These are leaf nodes if they have the `Children` component but it's empty
+ .map(|children| children.is_empty())
+ // Or if they don't have the `Children` component at all
+ .unwrap_or(true)
+ })
+ }
+
+ fn iter_siblings(&'w self, entity: Entity) -> impl Iterator<Item = Entity>
+ where
+ D::ReadOnly: WorldQuery<Item<'w> = (Option<&'w Parent>, Option<&'w Children>)>,
+ {
+ self.get(entity)
+ .ok()
+ .and_then(|(maybe_parent, _)| maybe_parent.map(Parent::get))
+ .and_then(|parent| self.get(parent).ok())
+ .and_then(|(_, maybe_children)| maybe_children)
+ .into_iter()
+ .flat_map(move |children| children.iter().filter(move |child| **child != entity))
+ .copied()
+ }
+
  fn iter_descendants(&'w self, entity: Entity) -> DescendantIter<'w, 's, D, F>
  where
  D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
  {
  DescendantIter::new(self, entity)
  }
 
+ fn iter_descendants_depth_first(
+ &'w self,
+ entity: Entity,
+ ) -> DescendantDepthFirstIter<'w, 's, D, F>
+ where
+ D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
+ {
+ DescendantDepthFirstIter::new(self, entity)
+ }
+
  fn iter_ancestors(&'w self, entity: Entity) -> AncestorIter<'w, 's, D, F>
  where
  D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>,
@@ -119,6 +233,51 @@ where
  }
 }
 
+/// An [`Iterator`] of [`Entity`]s over the descendants of an [`Entity`].
+///
+/// Traverses the hierarchy depth-first.
+pub struct DescendantDepthFirstIter<'w, 's, D: QueryData, F: QueryFilter>
+where
+ D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
+{
+ children_query: &'w Query<'w, 's, D, F>,
+ stack: SmallVec<[Entity; 8]>,
+}
+
+impl<'w, 's, D: QueryData, F: QueryFilter> DescendantDepthFirstIter<'w, 's, D, F>
+where
+ D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
+{
+ /// Returns a new [`DescendantDepthFirstIter`].
+ pub fn new(children_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
+ DescendantDepthFirstIter {
+ children_query,
+ stack: children_query
+ .get(entity)
+ .map_or(SmallVec::new(), |children| {
+ children.iter().rev().copied().collect()
+ }),
+ }
+ }
+}
+
+impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for DescendantDepthFirstIter<'w, 's, D, F>
+where
+ D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
+{
+ type Item = Entity;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ let entity = self.stack.pop()?;
+
+ if let Ok(children) = self.children_query.get(entity) {
+ self.stack.extend(children.iter().rev().copied());
+ }
+
+ Some(entity)
+ }
+}
+
 /// An [`Iterator`] of [`Entity`]s over the ancestors of an [`Entity`].
 pub struct AncestorIter<'w, 's, D: QueryData, F: QueryFilter>
 where
@@ -170,35 +329,108 @@ mod tests {
  fn descendant_iter() {
  let world = &mut World::new();
 
- let [a, b, c, d] = core::array::from_fn(|i| world.spawn(A(i)).id());
+ let [a0, a1, a2, a3] = core::array::from_fn(|i| world.spawn(A(i)).id());
 
- world.entity_mut(a).add_children(&[b, c]);
- world.entity_mut(c).add_children(&[d]);
+ world.entity_mut(a0).add_children(&[a1, a2]);
+ world.entity_mut(a1).add_children(&[a3]);
 
  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))
+ .iter_many(children_query.iter_descendants(a0))
  .collect();
 
  assert_eq!([&A(1), &A(2), &A(3)], result.as_slice());
  }
 
+ #[test]
+ fn descendant_depth_first_iter() {
+ let world = &mut World::new();
+
+ let [a0, a1, a2, a3] = core::array::from_fn(|i| world.spawn(A(i)).id());
+
+ world.entity_mut(a0).add_children(&[a1, a2]);
+ world.entity_mut(a1).add_children(&[a3]);
+
+ 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_depth_first(a0))
+ .collect();
+
+ assert_eq!([&A(1), &A(3), &A(2)], result.as_slice());
+ }
+
  #[test]
  fn ancestor_iter() {
  let world = &mut World::new();
 
- let [a, b, c] = core::array::from_fn(|i| world.spawn(A(i)).id());
+ let [a0, a1, a2] = core::array::from_fn(|i| world.spawn(A(i)).id());
 
- world.entity_mut(a).add_children(&[b]);
- world.entity_mut(b).add_children(&[c]);
+ world.entity_mut(a0).add_children(&[a1]);
+ world.entity_mut(a1).add_children(&[a2]);
 
  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();
+ let result: Vec<_> = a_query.iter_many(parent_query.iter_ancestors(a2)).collect();
 
  assert_eq!([&A(1), &A(0)], result.as_slice());
  }
+
+ #[test]
+ fn root_ancestor() {
+ let world = &mut World::new();
+
+ let [a0, a1, a2] = core::array::from_fn(|i| world.spawn(A(i)).id());
+
+ world.entity_mut(a0).add_children(&[a1]);
+ world.entity_mut(a1).add_children(&[a2]);
+
+ let mut system_state = SystemState::<Query<&Parent>>::new(world);
+ let parent_query = system_state.get(world);
+
+ assert_eq!(a0, parent_query.root_ancestor(a2));
+ assert_eq!(a0, parent_query.root_ancestor(a1));
+ assert_eq!(a0, parent_query.root_ancestor(a0));
+ }
+
+ #[test]
+ fn leaf_iter() {
+ let world = &mut World::new();
+
+ let [a0, a1, a2, a3] = core::array::from_fn(|i| world.spawn(A(i)).id());
+
+ world.entity_mut(a0).add_children(&[a1, a2]);
+ world.entity_mut(a1).add_children(&[a3]);
+
+ 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_leaves(a0)).collect();
+
+ assert_eq!([&A(3), &A(2)], result.as_slice());
+ }
+
+ #[test]
+ fn siblings() {
+ let world = &mut World::new();
+
+ let [a0, a1, a2, a3, a4] = core::array::from_fn(|i| world.spawn(A(i)).id());
+
+ world.entity_mut(a0).add_children(&[a1, a2, a3]);
+ world.entity_mut(a2).add_children(&[a4]);
+
+ let mut system_state =
+ SystemState::<(Query<(Option<&Parent>, Option<&Children>)>, Query<&A>)>::new(world);
+ let (hierarchy_query, a_query) = system_state.get(world);
+
+ let result: Vec<_> = a_query
+ .iter_many(hierarchy_query.iter_siblings(a1))
+ .collect();
+
+ assert_eq!([&A(2), &A(3)], result.as_slice());
+ }
 }