diff --git a/crates/bevy_transform/src/hierarchy/child_builder.rs b/crates/bevy_transform/src/hierarchy/child_builder.rs index b97c3b32dfba8..094a2af1314bd 100644 --- a/crates/bevy_transform/src/hierarchy/child_builder.rs +++ b/crates/bevy_transform/src/hierarchy/child_builder.rs @@ -40,11 +40,6 @@ pub struct PushChildren { children: SmallVec<[Entity; 8]>, } -pub struct ChildBuilder<'w, 's, 'a> { - commands: &'a mut Commands<'w, 's>, - push_children: PushChildren, -} - impl Command for PushChildren { fn write(self, world: &mut World) { for child in self.children.iter() { @@ -71,6 +66,46 @@ impl Command for PushChildren { } } +pub struct RemoveChildren { + parent: Entity, + children: SmallVec<[Entity; 8]>, +} + +fn remove_children(parent: Entity, children: &[Entity], world: &mut World) { + for child in children.iter() { + let mut child = world.entity_mut(*child); + let mut remove_parent = false; + if let Some(child_parent) = child.get_mut::() { + if child_parent.0 == parent { + remove_parent = true; + } + } + if remove_parent { + if let Some(parent) = child.remove::() { + child.insert(PreviousParent(parent.0)); + } + } + } + // Remove the children from the parents. + if let Some(mut parent_children) = world.get_mut::(parent) { + parent_children + .0 + .retain(|parent_child| !children.contains(parent_child)); + } +} + +impl Command for RemoveChildren { + fn write(self, world: &mut World) { + // Remove any matching Parent components from the children + remove_children(self.parent, &self.children, world); + } +} + +pub struct ChildBuilder<'w, 's, 'a> { + commands: &'a mut Commands<'w, 's>, + push_children: PushChildren, +} + impl<'w, 's, 'a> ChildBuilder<'w, 's, 'a> { pub fn spawn_bundle(&mut self, bundle: impl Bundle) -> EntityCommands<'w, 's, '_> { let e = self.commands.spawn_bundle(bundle); @@ -98,6 +133,7 @@ pub trait BuildChildren { fn with_children(&mut self, f: impl FnOnce(&mut ChildBuilder)) -> &mut Self; fn push_children(&mut self, children: &[Entity]) -> &mut Self; fn insert_children(&mut self, index: usize, children: &[Entity]) -> &mut Self; + fn remove_children(&mut self, children: &[Entity]) -> &mut Self; } impl<'w, 's, 'a> BuildChildren for EntityCommands<'w, 's, 'a> { @@ -137,6 +173,15 @@ impl<'w, 's, 'a> BuildChildren for EntityCommands<'w, 's, 'a> { }); self } + + fn remove_children(&mut self, children: &[Entity]) -> &mut Self { + let parent = self.id(); + self.commands().add(RemoveChildren { + children: SmallVec::from(children), + parent, + }); + self + } } #[derive(Debug)] @@ -196,6 +241,7 @@ pub trait BuildWorldChildren { fn with_children(&mut self, spawn_children: impl FnOnce(&mut WorldChildBuilder)) -> &mut Self; fn push_children(&mut self, children: &[Entity]) -> &mut Self; fn insert_children(&mut self, index: usize, children: &[Entity]) -> &mut Self; + fn remove_children(&mut self, children: &[Entity]) -> &mut Self; } impl<'w> BuildWorldChildren for EntityMut<'w> { @@ -260,6 +306,33 @@ impl<'w> BuildWorldChildren for EntityMut<'w> { } self } + + fn remove_children(&mut self, children: &[Entity]) -> &mut Self { + let parent = self.id(); + // SAFE: This doesn't change the parent's location + let world = unsafe { self.world_mut() }; + for child in children.iter() { + let mut child = world.entity_mut(*child); + let mut remove_parent = false; + if let Some(child_parent) = child.get_mut::() { + if child_parent.0 == parent { + remove_parent = true; + } + } + if remove_parent { + if let Some(parent) = child.remove::() { + child.insert(PreviousParent(parent.0)); + } + } + } + // Remove the children from the parents. + if let Some(mut parent_children) = world.get_mut::(parent) { + parent_children + .0 + .retain(|parent_child| !children.contains(parent_child)); + } + self + } } impl<'w> BuildWorldChildren for WorldChildBuilder<'w> { @@ -319,6 +392,15 @@ impl<'w> BuildWorldChildren for WorldChildBuilder<'w> { } self } + + fn remove_children(&mut self, children: &[Entity]) -> &mut Self { + let parent = self + .current_entity + .expect("Cannot remove children without a parent. Try creating an entity first."); + + remove_children(parent, children, self.world); + self + } } #[cfg(test)] @@ -369,7 +451,7 @@ mod tests { } #[test] - fn push_and_insert_children_commands() { + fn push_and_insert_and_remove_children_commands() { let mut world = World::default(); let entities = world @@ -427,10 +509,33 @@ mod tests { *world.get::(child4).unwrap(), PreviousParent(parent) ); + + let remove_children = [child1, child4]; + { + let mut commands = Commands::new(&mut queue, &world); + commands.entity(parent).remove_children(&remove_children); + } + queue.apply(&mut world); + + let expected_children: SmallVec<[Entity; 8]> = smallvec![child3, child2]; + assert_eq!( + world.get::(parent).unwrap().0.clone(), + expected_children + ); + assert!(world.get::(child1).is_none()); + assert!(world.get::(child4).is_none()); + assert_eq!( + *world.get::(child1).unwrap(), + PreviousParent(parent) + ); + assert_eq!( + *world.get::(child4).unwrap(), + PreviousParent(parent) + ); } #[test] - fn push_and_insert_children_world() { + fn push_and_insert_and_remove_children_world() { let mut world = World::default(); let entities = world @@ -478,6 +583,24 @@ mod tests { *world.get::(child4).unwrap(), PreviousParent(parent) ); + + let remove_children = [child1, child4]; + world.entity_mut(parent).remove_children(&remove_children); + let expected_children: SmallVec<[Entity; 8]> = smallvec![child3, child2]; + assert_eq!( + world.get::(parent).unwrap().0.clone(), + expected_children + ); + assert!(world.get::(child1).is_none()); + assert!(world.get::(child4).is_none()); + assert_eq!( + *world.get::(child1).unwrap(), + PreviousParent(parent) + ); + assert_eq!( + *world.get::(child4).unwrap(), + PreviousParent(parent) + ); } #[test]