diff --git a/crates/bevy_ecs/src/relationship/related_methods.rs b/crates/bevy_ecs/src/relationship/related_methods.rs index 4b42709384d1d..ea8bfb979fd77 100644 --- a/crates/bevy_ecs/src/relationship/related_methods.rs +++ b/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( + &mut self, + bundle: impl Bundle + Clone, + ) -> &mut Self { + self.insert(bundle.clone()); + if let Some(relationship_target) = self.get::() { + let related_vec: Vec = relationship_target.iter().collect(); + for related in related_vec { + self.world_scope(|world| { + world + .entity_mut(related) + .insert_recursive::(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(&mut self) -> &mut Self { + self.remove::(); + if let Some(relationship_target) = self.get::() { + let related_vec: Vec = relationship_target.iter().collect(); + for related in related_vec { + self.world_scope(|world| { + world.entity_mut(related).remove_recursive::(); + }); + } + } + + 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( + &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::(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(&mut self) -> &mut Self { + let id = self.id(); + self.commands.queue(move |world: &mut World| { + world.entity_mut(id).remove_recursive::(); + }); + 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::(TestComponent); + + for entity in [a, b, c, d] { + assert!(world.entity(entity).contains::()); + } + + world + .entity_mut(b) + .remove_recursive::(); + + // Parent + assert!(world.entity(a).contains::()); + // Target + assert!(!world.entity(b).contains::()); + // Sibling + assert!(world.entity(c).contains::()); + // Child + assert!(!world.entity(d).contains::()); + + world + .entity_mut(a) + .remove_recursive::(); + + for entity in [a, b, c, d] { + assert!(!world.entity(entity).contains::()); + } + } +}