add depth to iter_related

Author: jrobsonchase

crates/bevy_hierarchy/src/query_extension.rs

@@ -109,7 +109,7 @@ where
     C: Component + IterEntities,
 {
     related_query: &'w Query<'w, 's, D, F>,
-    vecdeque: VecDeque<Entity>,
+    vecdeque: VecDeque<(usize, Entity)>,
 }
 
 impl<'w, 's, C, D, F> RelatedIter<'w, 's, C, D, F>
@@ -120,16 +120,32 @@ where
     C: Component + IterEntities,
 {
     /// Create a new [`RelatedIter`].
-    pub fn new(related_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
+    fn new(related_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
         RelatedIter {
             related_query,
             vecdeque: related_query
                 .get(entity)
-                .iter()
+                .into_iter()
                 .flat_map(|t| t.iter_entities())
+                .map(|e| (1, e))
                 .collect(),
         }
     }
+
+    fn next(&mut self) -> Option<(usize, Entity)> {
+        let (depth, entity) = self.vecdeque.pop_front()?;
+
+        if let Ok(children) = self.related_query.get(entity) {
+            self.vecdeque
+                .extend(children.iter_entities().map(|e| (depth + 1, e)));
+        }
+
+        Some((depth, entity))
+    }
+
+    pub fn with_depth(self) -> RelatedDepthIter<'w, 's, C, D, F> {
+        RelatedDepthIter(self)
+    }
 }
 
 impl<'w, 's, C, D, F> Iterator for RelatedIter<'w, 's, C, D, F>
@@ -142,13 +158,28 @@ where
     type Item = Entity;
 
     fn next(&mut self) -> Option<Self::Item> {
-        let entity = self.vecdeque.pop_front()?;
+        RelatedIter::next(self).map(|(_, e)| e)
+    }
+}
 
-        if let Ok(children) = self.related_query.get(entity) {
-            self.vecdeque.extend(children.iter_entities());
-        }
+pub struct RelatedDepthIter<'w, 's, C, D, F>(RelatedIter<'w, 's, C, D, F>)
+where
+    D: QueryData,
+    F: QueryFilter,
+    D::ReadOnly: WorldQuery<Item<'w> = &'w C>,
+    C: Component + IterEntities;
 
-        Some(entity)
+impl<'w, 's, C, D, F> Iterator for RelatedDepthIter<'w, 's, C, D, F>
+where
+    D: QueryData,
+    F: QueryFilter,
+    D::ReadOnly: WorldQuery<Item<'w> = &'w C>,
+    C: Component + IterEntities,
+{
+    type Item = (usize, Entity);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        RelatedIter::next(&mut self.0)
     }
 }
 
@@ -213,9 +244,13 @@ mod tests {
         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_related(a)).collect();
+        let result: Vec<_> = children_query
+            .iter_related(a)
+            .with_depth()
+            .filter_map(|(d, e)| Some((d, a_query.get(e).ok()?)))
+            .collect();
 
-        assert_eq!([&A(1), &A(2), &A(3)], result.as_slice());
+        assert_eq!([(1, &A(1)), (1, &A(2)), (2, &A(3))], result.as_slice());
     }
 
     #[test]