diff --git a/crates/bevy_hierarchy/src/components/children.rs b/crates/bevy_hierarchy/src/components/children.rs index e9e3b72654709..541b48d68a07f 100644 --- a/crates/bevy_hierarchy/src/components/children.rs +++ b/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]>); diff --git a/crates/bevy_hierarchy/src/components/parent.rs b/crates/bevy_hierarchy/src/components/parent.rs index d4a5e2f3bec70..d22fa67d3fce7 100644 --- a/crates/bevy_hierarchy/src/components/parent.rs +++ b/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); diff --git a/crates/bevy_hierarchy/src/lib.rs b/crates/bevy_hierarchy/src/lib.rs index 1f8c08d9c76c4..0959fa8130a7b 100644 --- a/crates/bevy_hierarchy/src/lib.rs +++ b/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, }; } diff --git a/crates/bevy_hierarchy/src/query_extension.rs b/crates/bevy_hierarchy/src/query_extension.rs new file mode 100644 index 0000000000000..6282c81992514 --- /dev/null +++ b/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>, 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>, 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, +} + +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 { + 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, +} + +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.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()); + } +} diff --git a/examples/3d/update_gltf_scene.rs b/examples/3d/update_gltf_scene.rs index 265cf09f375b9..07f312c793995 100644 --- a/examples/3d/update_gltf_scene.rs +++ b/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.seconds_since_startup().sin() as f32 / 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); } } }