From 85eceb022da0326b47ac2b0d9202c9c9f01835bb Mon Sep 17 00:00:00 2001 From: Alice Cecile Date: Mon, 20 Jan 2025 21:57:57 -0500 Subject: [PATCH] Add insert and remove recursive methods on `EntityWorldMut` and `EntityCommands` (#17463) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Objective While being able to quickly add / remove components down a tree is broadly useful (material changing!), it's particularly necessary when combined with the newly added #13120. ## Solution Write four methods: covering both adding and removal on both `EntityWorldMut` and `EntityCommands`. These methods are generic over the `RelationshipTarget`, thanks to the freshly merged relations 🎉 ## Testing I've added a simple unit test for these methods. --------- Co-authored-by: Zachary Harrold --- .../src/relationship/related_methods.rs | 132 ++++++++++++++++++ 1 file changed, 132 insertions(+) diff --git a/crates/bevy_ecs/src/relationship/related_methods.rs b/crates/bevy_ecs/src/relationship/related_methods.rs index 4b42709384..ea8bfb979f 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::()); + } + } +}