Add insert and remove recursive methods on EntityWorldMut and EntityCommands (#17463)
				
					
				
			# 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 <zac@harrold.com.au>
This commit is contained in:
		
							parent
							
								
									42b928b90e
								
							
						
					
					
						commit
						85eceb022d
					
				| @ -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>()); | ||||
|         } | ||||
|     } | ||||
| } | ||||
|  | ||||
		Loading…
	
		Reference in New Issue
	
	Block a user
	 Alice Cecile
						Alice Cecile