mirror of
https://github.com/bevyengine/bevy
synced 2024-11-21 20:23:28 +00:00
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>
This commit is contained in:
Alice Cecile
GitHub
parent
584d14808a
commit
0a150b0d22
1 changed files with 241 additions and 9 deletions
|
|
@ -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,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<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>,
|
||||
|
|
@ -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<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());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue