bevy_ecs/relationship/
relationship_query.rs

1use crate::{
2    entity::Entity,
3    query::{QueryData, QueryFilter},
4    relationship::{Relationship, RelationshipTarget},
5    system::Query,
6};
7use alloc::collections::VecDeque;
8use smallvec::SmallVec;
9
10use super::SourceIter;
11
12impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
13    /// If the given `entity` contains the `R` [`Relationship`] component, returns the
14    /// target entity of that relationship.
15    pub fn related<R: Relationship>(&'w self, entity: Entity) -> Option<Entity>
16    where
17        <D as QueryData>::ReadOnly: QueryData<Item<'w> = &'w R>,
18    {
19        self.get(entity).map(R::get).ok()
20    }
21
22    /// If the given `entity` contains the `S` [`RelationshipTarget`] component, returns the
23    /// source entities stored on that component.
24    pub fn relationship_sources<S: RelationshipTarget>(
25        &'w self,
26        entity: Entity,
27    ) -> impl Iterator<Item = Entity> + 'w
28    where
29        <D as QueryData>::ReadOnly: QueryData<Item<'w> = &'w S>,
30    {
31        self.get(entity)
32            .into_iter()
33            .flat_map(RelationshipTarget::iter)
34    }
35
36    /// Recursively walks up the tree defined by the given `R` [`Relationship`] until
37    /// there are no more related entities, returning the "root entity" of the relationship hierarchy.
38    ///
39    /// # Warning
40    ///
41    /// For relationship graphs that contain loops, this could loop infinitely.
42    /// If your relationship is not a tree (like Bevy's hierarchy), be sure to stop if you encounter a duplicate entity.
43    pub fn root_ancestor<R: Relationship>(&'w self, entity: Entity) -> Entity
44    where
45        <D as QueryData>::ReadOnly: QueryData<Item<'w> = &'w R>,
46    {
47        // Recursively search up the tree until we're out of parents
48        match self.get(entity) {
49            Ok(parent) => self.root_ancestor(parent.get()),
50            Err(_) => entity,
51        }
52    }
53
54    /// Iterates all "leaf entities" as defined by the [`RelationshipTarget`] hierarchy.
55    ///
56    /// # Warning
57    ///
58    /// For relationship graphs that contain loops, this could loop infinitely.
59    /// If your relationship is not a tree (like Bevy's hierarchy), be sure to stop if you encounter a duplicate entity.
60    pub fn iter_leaves<S: RelationshipTarget>(
61        &'w self,
62        entity: Entity,
63    ) -> impl Iterator<Item = Entity> + 'w
64    where
65        <D as QueryData>::ReadOnly: QueryData<Item<'w> = &'w S>,
66        SourceIter<'w, S>: DoubleEndedIterator,
67    {
68        self.iter_descendants_depth_first(entity).filter(|entity| {
69            self.get(*entity)
70                // These are leaf nodes if they have the `Children` component but it's empty
71                .map(|children| children.len() == 0)
72                // Or if they don't have the `Children` component at all
73                .unwrap_or(true)
74        })
75    }
76
77    /// Iterates all sibling entities that also have the `R` [`Relationship`] with the same target entity.
78    pub fn iter_siblings<R: Relationship>(
79        &'w self,
80        entity: Entity,
81    ) -> impl Iterator<Item = Entity> + 'w
82    where
83        D::ReadOnly: QueryData<Item<'w> = (Option<&'w R>, Option<&'w R::RelationshipTarget>)>,
84    {
85        self.get(entity)
86            .ok()
87            .and_then(|(maybe_parent, _)| maybe_parent.map(R::get))
88            .and_then(|parent| self.get(parent).ok())
89            .and_then(|(_, maybe_children)| maybe_children)
90            .into_iter()
91            .flat_map(move |children| children.iter().filter(move |child| *child != entity))
92    }
93
94    /// Iterates all descendant entities as defined by the given `entity`'s [`RelationshipTarget`] and their recursive
95    /// [`RelationshipTarget`].
96    ///
97    /// # Warning
98    ///
99    /// For relationship graphs that contain loops, this could loop infinitely.
100    /// If your relationship is not a tree (like Bevy's hierarchy), be sure to stop if you encounter a duplicate entity.
101    pub fn iter_descendants<S: RelationshipTarget>(
102        &'w self,
103        entity: Entity,
104    ) -> DescendantIter<'w, 's, D, F, S>
105    where
106        D::ReadOnly: QueryData<Item<'w> = &'w S>,
107    {
108        DescendantIter::new(self, entity)
109    }
110
111    /// Iterates all descendant entities as defined by the given `entity`'s [`RelationshipTarget`] and their recursive
112    /// [`RelationshipTarget`] in depth-first order.
113    ///
114    /// # Warning
115    ///
116    /// For relationship graphs that contain loops, this could loop infinitely.
117    /// If your relationship is not a tree (like Bevy's hierarchy), be sure to stop if you encounter a duplicate entity.
118    pub fn iter_descendants_depth_first<S: RelationshipTarget>(
119        &'w self,
120        entity: Entity,
121    ) -> DescendantDepthFirstIter<'w, 's, D, F, S>
122    where
123        D::ReadOnly: QueryData<Item<'w> = &'w S>,
124        SourceIter<'w, S>: DoubleEndedIterator,
125    {
126        DescendantDepthFirstIter::new(self, entity)
127    }
128
129    /// Iterates all ancestors of the given `entity` as defined by the `R` [`Relationship`].
130    ///
131    /// # Warning
132    ///
133    /// For relationship graphs that contain loops, this could loop infinitely.
134    /// If your relationship is not a tree (like Bevy's hierarchy), be sure to stop if you encounter a duplicate entity.
135    pub fn iter_ancestors<R: Relationship>(
136        &'w self,
137        entity: Entity,
138    ) -> AncestorIter<'w, 's, D, F, R>
139    where
140        D::ReadOnly: QueryData<Item<'w> = &'w R>,
141    {
142        AncestorIter::new(self, entity)
143    }
144}
145
146/// An [`Iterator`] of [`Entity`]s over the descendants of an [`Entity`].
147///
148/// Traverses the hierarchy breadth-first.
149pub struct DescendantIter<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget>
150where
151    D::ReadOnly: QueryData<Item<'w> = &'w S>,
152{
153    children_query: &'w Query<'w, 's, D, F>,
154    vecdeque: VecDeque<Entity>,
155}
156
157impl<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget> DescendantIter<'w, 's, D, F, S>
158where
159    D::ReadOnly: QueryData<Item<'w> = &'w S>,
160{
161    /// Returns a new [`DescendantIter`].
162    pub fn new(children_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
163        DescendantIter {
164            children_query,
165            vecdeque: children_query
166                .get(entity)
167                .into_iter()
168                .flat_map(RelationshipTarget::iter)
169                .collect(),
170        }
171    }
172}
173
174impl<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget> Iterator
175    for DescendantIter<'w, 's, D, F, S>
176where
177    D::ReadOnly: QueryData<Item<'w> = &'w S>,
178{
179    type Item = Entity;
180
181    fn next(&mut self) -> Option<Self::Item> {
182        let entity = self.vecdeque.pop_front()?;
183
184        if let Ok(children) = self.children_query.get(entity) {
185            self.vecdeque.extend(children.iter());
186        }
187
188        Some(entity)
189    }
190}
191
192/// An [`Iterator`] of [`Entity`]s over the descendants of an [`Entity`].
193///
194/// Traverses the hierarchy depth-first.
195pub struct DescendantDepthFirstIter<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget>
196where
197    D::ReadOnly: QueryData<Item<'w> = &'w S>,
198{
199    children_query: &'w Query<'w, 's, D, F>,
200    stack: SmallVec<[Entity; 8]>,
201}
202
203impl<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget>
204    DescendantDepthFirstIter<'w, 's, D, F, S>
205where
206    D::ReadOnly: QueryData<Item<'w> = &'w S>,
207    SourceIter<'w, S>: DoubleEndedIterator,
208{
209    /// Returns a new [`DescendantDepthFirstIter`].
210    pub fn new(children_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
211        DescendantDepthFirstIter {
212            children_query,
213            stack: children_query
214                .get(entity)
215                .map_or(SmallVec::new(), |children| children.iter().rev().collect()),
216        }
217    }
218}
219
220impl<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget> Iterator
221    for DescendantDepthFirstIter<'w, 's, D, F, S>
222where
223    D::ReadOnly: QueryData<Item<'w> = &'w S>,
224    SourceIter<'w, S>: DoubleEndedIterator,
225{
226    type Item = Entity;
227
228    fn next(&mut self) -> Option<Self::Item> {
229        let entity = self.stack.pop()?;
230
231        if let Ok(children) = self.children_query.get(entity) {
232            self.stack.extend(children.iter().rev());
233        }
234
235        Some(entity)
236    }
237}
238
239/// An [`Iterator`] of [`Entity`]s over the ancestors of an [`Entity`].
240pub struct AncestorIter<'w, 's, D: QueryData, F: QueryFilter, R: Relationship>
241where
242    D::ReadOnly: QueryData<Item<'w> = &'w R>,
243{
244    parent_query: &'w Query<'w, 's, D, F>,
245    next: Option<Entity>,
246}
247
248impl<'w, 's, D: QueryData, F: QueryFilter, R: Relationship> AncestorIter<'w, 's, D, F, R>
249where
250    D::ReadOnly: QueryData<Item<'w> = &'w R>,
251{
252    /// Returns a new [`AncestorIter`].
253    pub fn new(parent_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
254        AncestorIter {
255            parent_query,
256            next: Some(entity),
257        }
258    }
259}
260
261impl<'w, 's, D: QueryData, F: QueryFilter, R: Relationship> Iterator
262    for AncestorIter<'w, 's, D, F, R>
263where
264    D::ReadOnly: QueryData<Item<'w> = &'w R>,
265{
266    type Item = Entity;
267
268    fn next(&mut self) -> Option<Self::Item> {
269        self.next = self.parent_query.get(self.next?).ok().map(R::get);
270        self.next
271    }
272}
bevy_ecs/relationship/relationship_query.rs

bevy_ecs/relationship/
relationship_query.rs
