Add add_child, set_parent and remove_parent to EntityMut (#6926)

# Objective
Align the hierarchy API between `EntityCommands` and `EntityMut`.

Added missing methods to `EntityMut`.
Replaced the duplicate `Command` implementations with the ones on `EntityMut` (e.g. The `AddChild` command is now just `world.entity_mut(..).add_child(..)`)

Fixed `update_old_parents` not sending `ChildAdded` events.

This PR does not add `add_children` to `EntityMut` as I would like to remove it from `EntityCommands` instead in #6942.

## Changelog
* Added `add_child`, `set_parent` and `remove_parent` to `EntityMut`
* Fixed missing `ChildAdded` events


Co-authored-by: devil-ira <[email protected]>

Author: tim-blackbird

crates/bevy_hierarchy/src/child_builder.rs

@@ -28,6 +28,9 @@ fn update_parent(world: &mut World, child: Entity, new_parent: Entity) -> Option
     }
 }
 
+/// Remove child from the parent's [`Children`] component.
+///
+/// Removes the [`Children`] component from the parent if it's empty.
 fn remove_from_children(world: &mut World, parent: Entity, child: Entity) {
     let mut parent = world.entity_mut(parent);
     if let Some(mut children) = parent.get_mut::<Children>() {
@@ -38,24 +41,61 @@ fn remove_from_children(world: &mut World, parent: Entity, child: Entity) {
     }
 }
 
+/// Update the [`Parent`] component of the `child`.
+/// Removes the `child` from the previous parent's [`Children`].
+///
+/// Does not update the new parents [`Children`] component.
+///
+/// Does nothing if `child` was already a child of `parent`.
+///
+/// Sends [`HierarchyEvent`]'s.
+fn update_old_parent(world: &mut World, child: Entity, parent: Entity) {
+    let previous = update_parent(world, child, parent);
+    if let Some(previous_parent) = previous {
+        // Do nothing if the child was already parented to this entity.
+        if previous_parent == parent {
+            return;
+        }
+        remove_from_children(world, previous_parent, child);
+
+        world.send_event(HierarchyEvent::ChildMoved {
+            child,
+            previous_parent,
+            new_parent: parent,
+        });
+    } else {
+        world.send_event(HierarchyEvent::ChildAdded { child, parent });
+    }
+}
+
+/// Update the [`Parent`] components of the `children`.
+/// Removes the `children` from their previous parent's [`Children`].
+///
+/// Does not update the new parents [`Children`] component.
+///
+/// Does nothing for a child if it was already a child of `parent`.
+///
+/// Sends [`HierarchyEvent`]'s.
 fn update_old_parents(world: &mut World, parent: Entity, children: &[Entity]) {
-    let mut moved: SmallVec<[HierarchyEvent; 8]> = SmallVec::with_capacity(children.len());
-    for child in children {
-        if let Some(previous) = update_parent(world, *child, parent) {
+    let mut events: SmallVec<[HierarchyEvent; 8]> = SmallVec::with_capacity(children.len());
+    for &child in children {
+        if let Some(previous) = update_parent(world, child, parent) {
             // Do nothing if the entity already has the correct parent.
             if parent == previous {
                 continue;
             }
 
-            remove_from_children(world, previous, *child);
-            moved.push(HierarchyEvent::ChildMoved {
-                child: *child,
+            remove_from_children(world, previous, child);
+            events.push(HierarchyEvent::ChildMoved {
+                child,
                 previous_parent: previous,
                 new_parent: parent,
             });
+        } else {
+            events.push(HierarchyEvent::ChildAdded { child, parent });
         }
     }
-    world.send_event_batch(moved);
+    world.send_event_batch(events);
 }
 
 fn remove_children(parent: Entity, children: &[Entity], world: &mut World) {
@@ -99,30 +139,7 @@ pub struct AddChild {
 
 impl Command for AddChild {
     fn write(self, world: &mut World) {
-        let previous = update_parent(world, self.child, self.parent);
-        if let Some(previous) = previous {
-            if previous == self.parent {
-                return;
-            }
-            remove_from_children(world, previous, self.child);
-            world.send_event(HierarchyEvent::ChildMoved {
-                child: self.child,
-                previous_parent: previous,
-                new_parent: self.parent,
-            });
-        }
-        world.send_event(HierarchyEvent::ChildAdded {
-            child: self.child,
-            parent: self.parent,
-        });
-        let mut parent = world.entity_mut(self.parent);
-        if let Some(mut children) = parent.get_mut::<Children>() {
-            if !children.contains(&self.child) {
-                children.0.push(self.child);
-            }
-        } else {
-            parent.insert(Children(smallvec::smallvec![self.child]));
-        }
+        world.entity_mut(self.parent).add_child(self.child);
     }
 }
 
@@ -136,14 +153,9 @@ pub struct InsertChildren {
 
 impl Command for InsertChildren {
     fn write(self, world: &mut World) {
-        update_old_parents(world, self.parent, &self.children);
-        let mut parent = world.entity_mut(self.parent);
-        if let Some(mut children) = parent.get_mut::<Children>() {
-            children.0.retain(|value| !self.children.contains(value));
-            children.0.insert_from_slice(self.index, &self.children);
-        } else {
-            parent.insert(Children(self.children));
-        }
+        world
+            .entity_mut(self.parent)
+            .insert_children(self.index, &self.children);
     }
 }
 
@@ -155,15 +167,8 @@ pub struct PushChildren {
 }
 
 impl Command for PushChildren {
-    fn write(mut self, world: &mut World) {
-        update_old_parents(world, self.parent, &self.children);
-        let mut parent = world.entity_mut(self.parent);
-        if let Some(mut children) = parent.get_mut::<Children>() {
-            children.0.retain(|child| !self.children.contains(child));
-            children.0.append(&mut self.children);
-        } else {
-            parent.insert(Children(self.children));
-        }
+    fn write(self, world: &mut World) {
+        world.entity_mut(self.parent).push_children(&self.children);
     }
 }
 
@@ -186,15 +191,7 @@ pub struct RemoveParent {
 
 impl Command for RemoveParent {
     fn write(self, world: &mut World) {
-        if let Some(parent) = world.get::<Parent>(self.child) {
-            let parent_entity = parent.get();
-            remove_from_children(world, parent_entity, self.child);
-            world.entity_mut(self.child).remove::<Parent>();
-            world.send_event(HierarchyEvent::ChildRemoved {
-                child: self.child,
-                parent: parent_entity,
-            });
-        }
+        world.entity_mut(self.child).remove_parent();
     }
 }
 
@@ -368,12 +365,28 @@ impl<'w> WorldChildBuilder<'w> {
 pub trait BuildWorldChildren {
     /// Creates a [`WorldChildBuilder`] with the given children built in the given closure
     fn with_children(&mut self, spawn_children: impl FnOnce(&mut WorldChildBuilder)) -> &mut Self;
+
+    /// Adds a single child
+    ///
+    /// If the children were previously children of another parent, that parent's [`Children`] component
+    /// will have those children removed from its list. Removing all children from a parent causes its
+    /// [`Children`] component to be removed from the entity.
+    fn add_child(&mut self, child: Entity) -> &mut Self;
+
     /// Pushes children to the back of the builder's children
     fn push_children(&mut self, children: &[Entity]) -> &mut Self;
     /// Inserts children at the given index
     fn insert_children(&mut self, index: usize, children: &[Entity]) -> &mut Self;
     /// Removes the given children
     fn remove_children(&mut self, children: &[Entity]) -> &mut Self;
+
+    /// Set the `parent` of this entity. This entity will be added to the end of the `parent`'s list of children.
+    ///
+    /// If this entity already had a parent it will be removed from it.
+    fn set_parent(&mut self, parent: Entity) -> &mut Self;
+
+    /// Remove the parent from this entity.
+    fn remove_parent(&mut self) -> &mut Self;
 }
 
 impl<'w> BuildWorldChildren for EntityMut<'w> {
@@ -385,6 +398,20 @@ impl<'w> BuildWorldChildren for EntityMut<'w> {
         self
     }
 
+    fn add_child(&mut self, child: Entity) -> &mut Self {
+        let parent = self.id();
+        self.world_scope(|world| {
+            update_old_parent(world, child, parent);
+        });
+        if let Some(mut children_component) = self.get_mut::<Children>() {
+            children_component.0.retain(|value| child != *value);
+            children_component.0.push(child);
+        } else {
+            self.insert(Children::from_entities(&[child]));
+        }
+        self
+    }
+
     fn push_children(&mut self, children: &[Entity]) -> &mut Self {
         let parent = self.id();
         self.world_scope(|world| {
@@ -424,21 +451,172 @@ impl<'w> BuildWorldChildren for EntityMut<'w> {
         });
         self
     }
+
+    fn set_parent(&mut self, parent: Entity) -> &mut Self {
+        let child = self.id();
+        self.world_scope(|world| {
+            world.entity_mut(parent).add_child(child);
+        });
+        self
+    }
+
+    fn remove_parent(&mut self) -> &mut Self {
+        let child = self.id();
+        if let Some(parent) = self.remove::<Parent>().map(|p| p.get()) {
+            self.world_scope(|world| {
+                remove_from_children(world, parent, child);
+                world.send_event(HierarchyEvent::ChildRemoved { child, parent });
+            });
+        }
+        self
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use super::{BuildChildren, BuildWorldChildren};
-    use crate::prelude::{Children, Parent};
+    use crate::{
+        components::{Children, Parent},
+        HierarchyEvent::{self, ChildAdded, ChildMoved, ChildRemoved},
+    };
     use smallvec::{smallvec, SmallVec};
 
     use bevy_ecs::{
         component::Component,
         entity::Entity,
+        event::Events,
         system::{CommandQueue, Commands},
         world::World,
     };
 
+    /// Assert the (non)existence and state of the child's [`Parent`] component.
+    fn assert_parent(world: &mut World, child: Entity, parent: Option<Entity>) {
+        assert_eq!(world.get::<Parent>(child).map(|p| p.get()), parent);
+    }
+
+    /// Assert the (non)existence and state of the parent's [`Children`] component.
+    fn assert_children(world: &mut World, parent: Entity, children: Option<&[Entity]>) {
+        assert_eq!(world.get::<Children>(parent).map(|c| &**c), children);
+    }
+
+    /// Used to omit a number of events that are not relevant to a particular test.
+    fn omit_events(world: &mut World, number: usize) {
+        let mut events_resource = world.resource_mut::<Events<HierarchyEvent>>();
+        let mut events: Vec<_> = events_resource.drain().collect();
+        events_resource.extend(events.drain(number..));
+    }
+
+    fn assert_events(world: &mut World, expected_events: &[HierarchyEvent]) {
+        let events: Vec<_> = world
+            .resource_mut::<Events<HierarchyEvent>>()
+            .drain()
+            .collect();
+        assert_eq!(events, expected_events);
+    }
+
+    #[test]
+    fn add_child() {
+        let world = &mut World::new();
+        world.insert_resource(Events::<HierarchyEvent>::default());
+
+        let [a, b, c, d] = std::array::from_fn(|_| world.spawn_empty().id());
+
+        world.entity_mut(a).add_child(b);
+
+        assert_parent(world, b, Some(a));
+        assert_children(world, a, Some(&[b]));
+        assert_events(
+            world,
+            &[ChildAdded {
+                child: b,
+                parent: a,
+            }],
+        );
+
+        world.entity_mut(a).add_child(c);
+
+        assert_children(world, a, Some(&[b, c]));
+        assert_parent(world, c, Some(a));
+        assert_events(
+            world,
+            &[ChildAdded {
+                child: c,
+                parent: a,
+            }],
+        );
+        // Children component should be removed when it's empty.
+        world.entity_mut(d).add_child(b).add_child(c);
+        assert_children(world, a, None);
+    }
+
+    #[test]
+    fn set_parent() {
+        let world = &mut World::new();
+        world.insert_resource(Events::<HierarchyEvent>::default());
+
+        let [a, b, c] = std::array::from_fn(|_| world.spawn_empty().id());
+
+        world.entity_mut(a).set_parent(b);
+
+        assert_parent(world, a, Some(b));
+        assert_children(world, b, Some(&[a]));
+        assert_events(
+            world,
+            &[ChildAdded {
+                child: a,
+                parent: b,
+            }],
+        );
+
+        world.entity_mut(a).set_parent(c);
+
+        assert_parent(world, a, Some(c));
+        assert_children(world, b, None);
+        assert_children(world, c, Some(&[a]));
+        assert_events(
+            world,
+            &[ChildMoved {
+                child: a,
+                previous_parent: b,
+                new_parent: c,
+            }],
+        );
+    }
+
+    #[test]
+    fn remove_parent() {
+        let world = &mut World::new();
+        world.insert_resource(Events::<HierarchyEvent>::default());
+
+        let [a, b, c] = std::array::from_fn(|_| world.spawn_empty().id());
+
+        world.entity_mut(a).push_children(&[b, c]);
+        world.entity_mut(b).remove_parent();
+
+        assert_parent(world, b, None);
+        assert_parent(world, c, Some(a));
+        assert_children(world, a, Some(&[c]));
+        omit_events(world, 2); // Omit ChildAdded events.
+        assert_events(
+            world,
+            &[ChildRemoved {
+                child: b,
+                parent: a,
+            }],
+        );
+
+        world.entity_mut(c).remove_parent();
+        assert_parent(world, c, None);
+        assert_children(world, a, None);
+        assert_events(
+            world,
+            &[ChildRemoved {
+                child: c,
+                parent: a,
+            }],
+        );
+    }
+
     #[derive(Component)]
     struct C(u32);
 

crates/bevy_hierarchy/src/events.rs

@@ -3,7 +3,7 @@ use bevy_ecs::prelude::Entity;
 /// An [`Event`] that is fired whenever there is a change in the world's hierarchy.
 ///
 /// [`Event`]: bevy_ecs::event::Event
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub enum HierarchyEvent {
     /// Fired whenever an [`Entity`] is added as a child to a parent.
     ChildAdded {