Add insert and remove recursive methods on EntityWorldMut and EntityCommands

crates/bevy_ecs/src/relationship/related_methods.rs

@@ -5,6 +5,7 @@ use crate::{
     system::{Commands, EntityCommands},
     world::{EntityWorldMut, World},
 };
+use alloc::vec::Vec;
 use core::marker::PhantomData;
 
 impl<'w> EntityWorldMut<'w> {
@@ -45,6 +46,55 @@ impl<'w> EntityWorldMut<'w> {
         }
         self
     }
+
+    /// Inserts a component or bundle of components into the entity and all related entities,
+    /// traversing the relationship tracked in `S` in a breadth-first manner.
+    ///
+    /// # Warning
+    ///
+    /// This method should only be called on relationships that form a tree-like structure.
+    /// Any cycles will cause this method to loop infinitely.
+    // We could keep track of a list of visited entities and track cycles,
+    // but this is not a very well-defined operation (or hard to write) for arbitrary relationships.
+    pub fn insert_recursive<S: RelationshipTarget>(
+        &mut self,
+        bundle: impl Bundle + Clone,
+    ) -> &mut Self {
+        self.insert(bundle.clone());
+        if let Some(relationship_target) = self.get::<S>() {
+            let related_vec: Vec<Entity> = relationship_target.iter().collect();
+            for related in related_vec {
+                self.world_scope(|world| {
+                    world
+                        .entity_mut(related)
+                        .insert_recursive::<S>(bundle.clone());
+                });
+            }
+        }
+
+        self
+    }
+
+    /// Removes a component or bundle of components of type `B` from the entity and all related entities,
+    /// traversing the relationship tracked in `S` in a breadth-first manner.
+    ///
+    /// # Warning
+    ///
+    /// This method should only be called on relationships that form a tree-like structure.
+    /// Any cycles will cause this method to loop infinitely.
+    pub fn remove_recursive<S: RelationshipTarget, B: Bundle>(&mut self) -> &mut Self {
+        self.remove::<B>();
+        if let Some(relationship_target) = self.get::<S>() {
+            let related_vec: Vec<Entity> = relationship_target.iter().collect();
+            for related in related_vec {
+                self.world_scope(|world| {
+                    world.entity_mut(related).remove_recursive::<S, B>();
+                });
+            }
+        }
+
+        self
+    }
 }
 
 impl<'a> EntityCommands<'a> {
@@ -79,6 +129,39 @@ impl<'a> EntityCommands<'a> {
         });
         self
     }
+
+    /// Inserts a component or bundle of components into the entity and all related entities,
+    /// traversing the relationship tracked in `S` in a breadth-first manner.
+    ///
+    /// # Warning
+    ///
+    /// This method should only be called on relationships that form a tree-like structure.
+    /// Any cycles will cause this method to loop infinitely.
+    pub fn insert_recursive<S: RelationshipTarget>(
+        &mut self,
+        bundle: impl Bundle + Clone,
+    ) -> &mut Self {
+        let id = self.id();
+        self.commands.queue(move |world: &mut World| {
+            world.entity_mut(id).insert_recursive::<S>(bundle);
+        });
+        self
+    }
+
+    /// Removes a component or bundle of components of type `B` from the entity and all related entities,
+    /// traversing the relationship tracked in `S` in a breadth-first manner.
+    ///
+    /// # Warning
+    ///
+    /// This method should only be called on relationships that form a tree-like structure.
+    /// Any cycles will cause this method to loop infinitely.
+    pub fn remove_recursive<S: RelationshipTarget, B: Bundle>(&mut self) -> &mut Self {
+        let id = self.id();
+        self.commands.queue(move |world: &mut World| {
+            world.entity_mut(id).remove_recursive::<S, B>();
+        });
+        self
+    }
 }
 
 /// Directly spawns related "source" entities with the given [`Relationship`], targeting
@@ -162,3 +245,52 @@ impl<'w, R: Relationship> RelatedSpawnerCommands<'w, R> {
         &mut self.commands
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate as bevy_ecs;
+    use crate::prelude::{ChildOf, Children, Component};
+
+    #[derive(Component, Clone, Copy)]
+    struct TestComponent;
+
+    #[test]
+    fn insert_and_remove_recursive() {
+        let mut world = World::new();
+
+        let a = world.spawn_empty().id();
+        let b = world.spawn(ChildOf(a)).id();
+        let c = world.spawn(ChildOf(a)).id();
+        let d = world.spawn(ChildOf(b)).id();
+
+        world
+            .entity_mut(a)
+            .insert_recursive::<Children>(TestComponent);
+
+        for entity in [a, b, c, d] {
+            assert!(world.entity(entity).contains::<TestComponent>());
+        }
+
+        world
+            .entity_mut(b)
+            .remove_recursive::<Children, TestComponent>();
+
+        // Parent
+        assert!(world.entity(a).contains::<TestComponent>());
+        // Target
+        assert!(!world.entity(b).contains::<TestComponent>());
+        // Sibling
+        assert!(world.entity(c).contains::<TestComponent>());
+        // Child
+        assert!(!world.entity(d).contains::<TestComponent>());
+
+        world
+            .entity_mut(a)
+            .remove_recursive::<Children, TestComponent>();
+
+        for entity in [a, b, c, d] {
+            assert!(!world.entity(entity).contains::<TestComponent>());
+        }
+    }
+}