bevy_hierarchy/
query_extension.rs

1use alloc::collections::VecDeque;
2
3use bevy_ecs::{
4    entity::Entity,
5    query::{QueryData, QueryFilter, WorldQuery},
6    system::Query,
7};
8use smallvec::SmallVec;
9
10use crate::{Children, Parent};
11
12/// An extension trait for [`Query`] that adds hierarchy related methods.
13pub trait HierarchyQueryExt<'w, 's, D: QueryData, F: QueryFilter> {
14    /// Returns the parent [`Entity`] of the given `entity`, if any.
15    fn parent(&'w self, entity: Entity) -> Option<Entity>
16    where
17        D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>;
18
19    /// Returns a slice over the [`Children`] of the given `entity`.
20    ///
21    /// This may be empty if the `entity` has no children.
22    fn children(&'w self, entity: Entity) -> &'w [Entity]
23    where
24        D::ReadOnly: WorldQuery<Item<'w> = &'w Children>;
25
26    /// Returns the topmost ancestor of the given `entity`.
27    ///
28    /// This may be the entity itself if it has no parent.
29    fn root_ancestor(&'w self, entity: Entity) -> Entity
30    where
31        D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>;
32
33    /// Returns an [`Iterator`] of [`Entity`]s over the leaves of the hierarchy that are underneath this `entity`.
34    ///
35    /// Only entities which have no children are considered leaves.
36    /// This will not include the entity itself, and will not include any entities which are not descendants of the entity,
37    /// even if they are leaves in the same hierarchical tree.
38    ///
39    /// Traverses the hierarchy depth-first.
40    fn iter_leaves(&'w self, entity: Entity) -> impl Iterator<Item = Entity> + 'w
41    where
42        D::ReadOnly: WorldQuery<Item<'w> = &'w Children>;
43
44    /// Returns an [`Iterator`] of [`Entity`]s over the `entity`s immediate siblings, who share the same parent.
45    ///
46    /// The entity itself is not included in the iterator.
47    fn iter_siblings(&'w self, entity: Entity) -> impl Iterator<Item = Entity>
48    where
49        D::ReadOnly: WorldQuery<Item<'w> = (Option<&'w Parent>, Option<&'w Children>)>;
50
51    /// Returns an [`Iterator`] of [`Entity`]s over all of `entity`s descendants.
52    ///
53    /// Can only be called on a [`Query`] of [`Children`] (i.e. `Query<&Children>`).
54    ///
55    /// Traverses the hierarchy breadth-first and does not include the entity itself.
56    ///
57    /// # Examples
58    /// ```
59    /// # use bevy_ecs::prelude::*;
60    /// # use bevy_hierarchy::prelude::*;
61    /// # #[derive(Component)]
62    /// # struct Marker;
63    /// fn system(entity: Single<Entity, With<Marker>>, children_query: Query<&Children>) {
64    ///     for descendant in children_query.iter_descendants(*entity) {
65    ///         // Do something!
66    ///     }
67    /// }
68    /// # bevy_ecs::system::assert_is_system(system);
69    /// ```
70    fn iter_descendants(&'w self, entity: Entity) -> DescendantIter<'w, 's, D, F>
71    where
72        D::ReadOnly: WorldQuery<Item<'w> = &'w Children>;
73
74    /// Returns an [`Iterator`] of [`Entity`]s over all of `entity`s descendants.
75    ///
76    /// Can only be called on a [`Query`] of [`Children`] (i.e. `Query<&Children>`).
77    ///
78    /// This is a depth-first alternative to [`HierarchyQueryExt::iter_descendants`].
79    fn iter_descendants_depth_first(
80        &'w self,
81        entity: Entity,
82    ) -> DescendantDepthFirstIter<'w, 's, D, F>
83    where
84        D::ReadOnly: WorldQuery<Item<'w> = &'w Children>;
85
86    /// Returns an [`Iterator`] of [`Entity`]s over all of `entity`s ancestors.
87    ///
88    /// Does not include the entity itself.
89    /// Can only be called on a [`Query`] of [`Parent`] (i.e. `Query<&Parent>`).
90    ///
91    /// # Examples
92    /// ```
93    /// # use bevy_ecs::prelude::*;
94    /// # use bevy_hierarchy::prelude::*;
95    /// # #[derive(Component)]
96    /// # struct Marker;
97    /// fn system(entity: Single<Entity, With<Marker>>, parent_query: Query<&Parent>) {
98    ///     for ancestor in parent_query.iter_ancestors(*entity) {
99    ///         // Do something!
100    ///     }
101    /// }
102    /// # bevy_ecs::system::assert_is_system(system);
103    /// ```
104    fn iter_ancestors(&'w self, entity: Entity) -> AncestorIter<'w, 's, D, F>
105    where
106        D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>;
107}
108
109impl<'w, 's, D: QueryData, F: QueryFilter> HierarchyQueryExt<'w, 's, D, F> for Query<'w, 's, D, F> {
110    fn parent(&'w self, entity: Entity) -> Option<Entity>
111    where
112        <D as QueryData>::ReadOnly: WorldQuery<Item<'w> = &'w Parent>,
113    {
114        self.get(entity).map(Parent::get).ok()
115    }
116
117    fn children(&'w self, entity: Entity) -> &'w [Entity]
118    where
119        <D as QueryData>::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
120    {
121        self.get(entity)
122            .map_or(&[] as &[Entity], |children| children)
123    }
124
125    fn root_ancestor(&'w self, entity: Entity) -> Entity
126    where
127        <D as QueryData>::ReadOnly: WorldQuery<Item<'w> = &'w Parent>,
128    {
129        // Recursively search up the tree until we're out of parents
130        match self.get(entity) {
131            Ok(parent) => self.root_ancestor(parent.get()),
132            Err(_) => entity,
133        }
134    }
135
136    fn iter_leaves(&'w self, entity: Entity) -> impl Iterator<Item = Entity>
137    where
138        <D as QueryData>::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
139    {
140        self.iter_descendants_depth_first(entity).filter(|entity| {
141            self.get(*entity)
142                // These are leaf nodes if they have the `Children` component but it's empty
143                .map(|children| children.is_empty())
144                // Or if they don't have the `Children` component at all
145                .unwrap_or(true)
146        })
147    }
148
149    fn iter_siblings(&'w self, entity: Entity) -> impl Iterator<Item = Entity>
150    where
151        D::ReadOnly: WorldQuery<Item<'w> = (Option<&'w Parent>, Option<&'w Children>)>,
152    {
153        self.get(entity)
154            .ok()
155            .and_then(|(maybe_parent, _)| maybe_parent.map(Parent::get))
156            .and_then(|parent| self.get(parent).ok())
157            .and_then(|(_, maybe_children)| maybe_children)
158            .into_iter()
159            .flat_map(move |children| children.iter().filter(move |child| **child != entity))
160            .copied()
161    }
162
163    fn iter_descendants(&'w self, entity: Entity) -> DescendantIter<'w, 's, D, F>
164    where
165        D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
166    {
167        DescendantIter::new(self, entity)
168    }
169
170    fn iter_descendants_depth_first(
171        &'w self,
172        entity: Entity,
173    ) -> DescendantDepthFirstIter<'w, 's, D, F>
174    where
175        D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
176    {
177        DescendantDepthFirstIter::new(self, entity)
178    }
179
180    fn iter_ancestors(&'w self, entity: Entity) -> AncestorIter<'w, 's, D, F>
181    where
182        D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>,
183    {
184        AncestorIter::new(self, entity)
185    }
186}
187
188/// An [`Iterator`] of [`Entity`]s over the descendants of an [`Entity`].
189///
190/// Traverses the hierarchy breadth-first.
191pub struct DescendantIter<'w, 's, D: QueryData, F: QueryFilter>
192where
193    D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
194{
195    children_query: &'w Query<'w, 's, D, F>,
196    vecdeque: VecDeque<Entity>,
197}
198
199impl<'w, 's, D: QueryData, F: QueryFilter> DescendantIter<'w, 's, D, F>
200where
201    D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
202{
203    /// Returns a new [`DescendantIter`].
204    pub fn new(children_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
205        DescendantIter {
206            children_query,
207            vecdeque: children_query
208                .get(entity)
209                .into_iter()
210                .flatten()
211                .copied()
212                .collect(),
213        }
214    }
215}
216
217impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for DescendantIter<'w, 's, D, F>
218where
219    D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
220{
221    type Item = Entity;
222
223    fn next(&mut self) -> Option<Self::Item> {
224        let entity = self.vecdeque.pop_front()?;
225
226        if let Ok(children) = self.children_query.get(entity) {
227            self.vecdeque.extend(children);
228        }
229
230        Some(entity)
231    }
232}
233
234/// An [`Iterator`] of [`Entity`]s over the descendants of an [`Entity`].
235///
236/// Traverses the hierarchy depth-first.
237pub struct DescendantDepthFirstIter<'w, 's, D: QueryData, F: QueryFilter>
238where
239    D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
240{
241    children_query: &'w Query<'w, 's, D, F>,
242    stack: SmallVec<[Entity; 8]>,
243}
244
245impl<'w, 's, D: QueryData, F: QueryFilter> DescendantDepthFirstIter<'w, 's, D, F>
246where
247    D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
248{
249    /// Returns a new [`DescendantDepthFirstIter`].
250    pub fn new(children_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
251        DescendantDepthFirstIter {
252            children_query,
253            stack: children_query
254                .get(entity)
255                .map_or(SmallVec::new(), |children| {
256                    children.iter().rev().copied().collect()
257                }),
258        }
259    }
260}
261
262impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for DescendantDepthFirstIter<'w, 's, D, F>
263where
264    D::ReadOnly: WorldQuery<Item<'w> = &'w Children>,
265{
266    type Item = Entity;
267
268    fn next(&mut self) -> Option<Self::Item> {
269        let entity = self.stack.pop()?;
270
271        if let Ok(children) = self.children_query.get(entity) {
272            self.stack.extend(children.iter().rev().copied());
273        }
274
275        Some(entity)
276    }
277}
278
279/// An [`Iterator`] of [`Entity`]s over the ancestors of an [`Entity`].
280pub struct AncestorIter<'w, 's, D: QueryData, F: QueryFilter>
281where
282    D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>,
283{
284    parent_query: &'w Query<'w, 's, D, F>,
285    next: Option<Entity>,
286}
287
288impl<'w, 's, D: QueryData, F: QueryFilter> AncestorIter<'w, 's, D, F>
289where
290    D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>,
291{
292    /// Returns a new [`AncestorIter`].
293    pub fn new(parent_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
294        AncestorIter {
295            parent_query,
296            next: Some(entity),
297        }
298    }
299}
300
301impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for AncestorIter<'w, 's, D, F>
302where
303    D::ReadOnly: WorldQuery<Item<'w> = &'w Parent>,
304{
305    type Item = Entity;
306
307    fn next(&mut self) -> Option<Self::Item> {
308        self.next = self.parent_query.get(self.next?).ok().map(Parent::get);
309        self.next
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use bevy_ecs::{
316        prelude::Component,
317        system::{Query, SystemState},
318        world::World,
319    };
320
321    use crate::{query_extension::HierarchyQueryExt, BuildChildren, Children, Parent};
322
323    #[derive(Component, PartialEq, Debug)]
324    struct A(usize);
325
326    #[test]
327    fn descendant_iter() {
328        let world = &mut World::new();
329
330        let [a0, a1, a2, a3] = core::array::from_fn(|i| world.spawn(A(i)).id());
331
332        world.entity_mut(a0).add_children(&[a1, a2]);
333        world.entity_mut(a1).add_children(&[a3]);
334
335        let mut system_state = SystemState::<(Query<&Children>, Query<&A>)>::new(world);
336        let (children_query, a_query) = system_state.get(world);
337
338        let result: Vec<_> = a_query
339            .iter_many(children_query.iter_descendants(a0))
340            .collect();
341
342        assert_eq!([&A(1), &A(2), &A(3)], result.as_slice());
343    }
344
345    #[test]
346    fn descendant_depth_first_iter() {
347        let world = &mut World::new();
348
349        let [a0, a1, a2, a3] = core::array::from_fn(|i| world.spawn(A(i)).id());
350
351        world.entity_mut(a0).add_children(&[a1, a2]);
352        world.entity_mut(a1).add_children(&[a3]);
353
354        let mut system_state = SystemState::<(Query<&Children>, Query<&A>)>::new(world);
355        let (children_query, a_query) = system_state.get(world);
356
357        let result: Vec<_> = a_query
358            .iter_many(children_query.iter_descendants_depth_first(a0))
359            .collect();
360
361        assert_eq!([&A(1), &A(3), &A(2)], result.as_slice());
362    }
363
364    #[test]
365    fn ancestor_iter() {
366        let world = &mut World::new();
367
368        let [a0, a1, a2] = core::array::from_fn(|i| world.spawn(A(i)).id());
369
370        world.entity_mut(a0).add_children(&[a1]);
371        world.entity_mut(a1).add_children(&[a2]);
372
373        let mut system_state = SystemState::<(Query<&Parent>, Query<&A>)>::new(world);
374        let (parent_query, a_query) = system_state.get(world);
375
376        let result: Vec<_> = a_query.iter_many(parent_query.iter_ancestors(a2)).collect();
377
378        assert_eq!([&A(1), &A(0)], result.as_slice());
379    }
380
381    #[test]
382    fn root_ancestor() {
383        let world = &mut World::new();
384
385        let [a0, a1, a2] = core::array::from_fn(|i| world.spawn(A(i)).id());
386
387        world.entity_mut(a0).add_children(&[a1]);
388        world.entity_mut(a1).add_children(&[a2]);
389
390        let mut system_state = SystemState::<Query<&Parent>>::new(world);
391        let parent_query = system_state.get(world);
392
393        assert_eq!(a0, parent_query.root_ancestor(a2));
394        assert_eq!(a0, parent_query.root_ancestor(a1));
395        assert_eq!(a0, parent_query.root_ancestor(a0));
396    }
397
398    #[test]
399    fn leaf_iter() {
400        let world = &mut World::new();
401
402        let [a0, a1, a2, a3] = core::array::from_fn(|i| world.spawn(A(i)).id());
403
404        world.entity_mut(a0).add_children(&[a1, a2]);
405        world.entity_mut(a1).add_children(&[a3]);
406
407        let mut system_state = SystemState::<(Query<&Children>, Query<&A>)>::new(world);
408        let (children_query, a_query) = system_state.get(world);
409
410        let result: Vec<_> = a_query.iter_many(children_query.iter_leaves(a0)).collect();
411
412        assert_eq!([&A(3), &A(2)], result.as_slice());
413    }
414
415    #[test]
416    fn siblings() {
417        let world = &mut World::new();
418
419        let [a0, a1, a2, a3, a4] = core::array::from_fn(|i| world.spawn(A(i)).id());
420
421        world.entity_mut(a0).add_children(&[a1, a2, a3]);
422        world.entity_mut(a2).add_children(&[a4]);
423
424        let mut system_state =
425            SystemState::<(Query<(Option<&Parent>, Option<&Children>)>, Query<&A>)>::new(world);
426        let (hierarchy_query, a_query) = system_state.get(world);
427
428        let result: Vec<_> = a_query
429            .iter_many(hierarchy_query.iter_siblings(a1))
430            .collect();
431
432        assert_eq!([&A(2), &A(3)], result.as_slice());
433    }
434}
bevy_hierarchy/query_extension.rs

bevy_hierarchy/
query_extension.rs
