From 0a150b0d22711443a943d37a98b7afb7e9862130 Mon Sep 17 00:00:00 2001 From: Alice Cecile Date: Mon, 7 Oct 2024 11:24:57 -0400 Subject: [PATCH] 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 Co-authored-by: poopy Co-authored-by: Christian Hughes <9044780+ItsDoot@users.noreply.github.com> --- crates/bevy_hierarchy/src/query_extension.rs | 250 ++++++++++++++++++- 1 file changed, 241 insertions(+), 9 deletions(-) diff --git a/crates/bevy_hierarchy/src/query_extension.rs b/crates/bevy_hierarchy/src/query_extension.rs index 36bf790cec1bd..5cd8631310762 100644 --- a/crates/bevy_hierarchy/src/query_extension.rs +++ b/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 + where + D::ReadOnly: WorldQuery = &'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 = &'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 = &'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 + 'w + where + D::ReadOnly: WorldQuery = &'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 + where + D::ReadOnly: WorldQuery = (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 = &'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 = &'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,6 +109,59 @@ 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 + where + ::ReadOnly: WorldQuery = &'w Parent>, + { + self.get(entity).map(Parent::get).ok() + } + + fn children(&'w self, entity: Entity) -> &'w [Entity] + where + ::ReadOnly: WorldQuery = &'w Children>, + { + self.get(entity) + .map_or(&[] as &[Entity], |children| children) + } + + fn root_ancestor(&'w self, entity: Entity) -> Entity + where + ::ReadOnly: WorldQuery = &'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 + where + ::ReadOnly: WorldQuery = &'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 + where + D::ReadOnly: WorldQuery = (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 = &'w Children>, @@ -65,6 +169,16 @@ impl<'w, 's, D: QueryData, F: QueryFilter> HierarchyQueryExt<'w, 's, D, F> for Q DescendantIter::new(self, entity) } + fn iter_descendants_depth_first( + &'w self, + entity: Entity, + ) -> DescendantDepthFirstIter<'w, 's, D, F> + where + D::ReadOnly: WorldQuery = &'w Children>, + { + DescendantDepthFirstIter::new(self, entity) + } + fn iter_ancestors(&'w self, entity: Entity) -> AncestorIter<'w, 's, D, F> where D::ReadOnly: WorldQuery = &'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 = &'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 = &'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 = &'w Children>, +{ + type Item = Entity; + + fn next(&mut self) -> Option { + 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::>::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()); + } }