RemoveChildren command

Author: James Higgins

crates/bevy_transform/src/hierarchy/child_builder.rs

@@ -46,11 +46,6 @@ pub struct PushChildren {
     children: SmallVec<[Entity; 8]>,
 }
 
-pub struct ChildBuilder<'a, 'b> {
-    commands: &'b mut Commands<'a>,
-    push_children: PushChildren,
-}
-
 impl Command for PushChildren {
     fn write(self: Box<Self>, world: &mut World) {
         for child in self.children.iter() {
@@ -77,6 +72,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::<Parent>() {
+            if child_parent.0 == parent {
+                remove_parent = true;
+            }
+        }
+        if remove_parent {
+            if let Some(parent) = child.remove::<Parent>() {
+                child.insert(PreviousParent(parent.0));
+            }
+        }
+    }
+    // Remove the children from the parents.
+    if let Some(mut parent_children) = world.get_mut::<Children>(parent) {
+        parent_children
+            .0
+            .retain(|parent_child| !children.contains(parent_child));
+    }
+}
+
+impl Command for RemoveChildren {
+    fn write(self: Box<Self>, world: &mut World) {
+        // Remove any matching Parent components from the children
+        remove_children(self.parent, &self.children, world);
+    }
+}
+
+pub struct ChildBuilder<'a, 'b> {
+    commands: &'b mut Commands<'a>,
+    push_children: PushChildren,
+}
+
 impl<'a, 'b> ChildBuilder<'a, 'b> {
     pub fn spawn_bundle(&mut self, bundle: impl Bundle) -> EntityCommands<'a, '_> {
         let e = self.commands.spawn_bundle(bundle);
@@ -104,6 +139,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<'a, 'b> BuildChildren for EntityCommands<'a, 'b> {
@@ -143,6 +179,15 @@ impl<'a, 'b> BuildChildren for EntityCommands<'a, 'b> {
         });
         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)]
@@ -202,6 +247,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> {
@@ -262,6 +308,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::<Parent>() {
+                if child_parent.0 == parent {
+                    remove_parent = true;
+                }
+            }
+            if remove_parent {
+                if let Some(parent) = child.remove::<Parent>() {
+                    child.insert(PreviousParent(parent.0));
+                }
+            }
+        }
+        // Remove the children from the parents.
+        if let Some(mut parent_children) = world.get_mut::<Children>(parent) {
+            parent_children
+                .0
+                .retain(|parent_child| !children.contains(parent_child));
+        }
+        self
+    }
 }
 
 impl<'w> BuildWorldChildren for WorldChildBuilder<'w> {
@@ -321,6 +394,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)]
@@ -367,7 +449,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
@@ -425,10 +507,33 @@ mod tests {
             *world.get::<PreviousParent>(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::<Children>(parent).unwrap().0.clone(),
+            expected_children
+        );
+        assert!(world.get::<Parent>(child1).is_none());
+        assert!(world.get::<Parent>(child4).is_none());
+        assert_eq!(
+            *world.get::<PreviousParent>(child1).unwrap(),
+            PreviousParent(parent)
+        );
+        assert_eq!(
+            *world.get::<PreviousParent>(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
@@ -476,5 +581,23 @@ mod tests {
             *world.get::<PreviousParent>(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::<Children>(parent).unwrap().0.clone(),
+            expected_children
+        );
+        assert!(world.get::<Parent>(child1).is_none());
+        assert!(world.get::<Parent>(child4).is_none());
+        assert_eq!(
+            *world.get::<PreviousParent>(child1).unwrap(),
+            PreviousParent(parent)
+        );
+        assert_eq!(
+            *world.get::<PreviousParent>(child4).unwrap(),
+            PreviousParent(parent)
+        );
     }
 }