Remove children during on_remove

src/hierarchy.rs

@@ -29,21 +29,42 @@ use bevy_hierarchy::BuildWorldChildren;
 /// fn spawn_hierarchy(mut commands: Commands) {
 ///   commands.spawn(
 ///    (A, // Parent
-///     WithChild( // This component is removed on spawn
+///     WithChild::new(
 ///       (A, B(3)) // Child
 ///     )
 ///   ));
 /// }
 /// ```
-#[derive(Debug, Clone, Default)]
-pub struct WithChild<B: Bundle>(pub B);
+#[derive(Debug, Clone)]
+pub struct WithChild<B: Bundle> {
+    /// The bundle to use when spawning the child entity.
+    /// This is an `Option` to allow for moving the bundle out of the component
+    /// and is stored on the heap to minimize size overhead.
+    bundle: Option<Box<B>>,
+}
+
+impl<B: Bundle + Default> Default for WithChild<B> {
+    fn default() -> Self {
+        Self::new(B::default())
+    }
+}
+
+impl<B: Bundle> WithChild<B> {
+    /// Create a new `WithChild` component with the given bundle.
+    pub fn new(bundle: B) -> Self {
+        Self {
+            bundle: Some(Box::new(bundle)),
+        }
+    }
+}
 
 impl<B: Bundle> Component for WithChild<B> {
     /// This is a sparse set component as it's only ever added and removed, never iterated over.
     const STORAGE_TYPE: StorageType = StorageType::SparseSet;
 
     fn register_component_hooks(hooks: &mut ComponentHooks) {
         hooks.on_add(with_child_hook::<B>);
+        hooks.on_remove(with_remove_child_hook::<B>);
     }
 }
 
@@ -62,6 +83,27 @@ fn with_child_hook<B: Bundle>(
     });
 }
 
+/// A hook that runs whenever [`WithChild`] is added to an entity.
+///
+/// Generates a [`WithChildCommand`].
+fn with_remove_child_hook<B: Bundle>(
+    mut world: DeferredWorld<'_>,
+    entity: Entity,
+    _component_id: ComponentId,
+) {
+    // Component hooks can't perform structural changes, so we need to rely on commands.
+    world.commands().add(WithChildRemoveCommand {
+        parent_entity: entity,
+        _phantom: PhantomData::<B>,
+    });
+}
+
+#[derive(Component)]
+struct WithChildEntity<B> {
+    child_entity: Entity,
+    _phantom: PhantomData<B>,
+}
+
 struct WithChildCommand<B> {
     parent_entity: Entity,
     _phantom: PhantomData<B>,
@@ -77,16 +119,45 @@ impl<B: Bundle> Command for WithChildCommand<B> {
             return;
         };
 
-        let Some(with_child_component) = entity_mut.take::<WithChild<B>>() else {
+        let Some(mut with_child_component) = entity_mut.get_mut::<WithChild<B>>() else {
             #[cfg(debug_assertions)]
             panic!("WithChild component not found");
 
             #[cfg(not(debug_assertions))]
             return;
         };
 
-        let child_entity = world.spawn(with_child_component.0).id();
-        world.entity_mut(self.parent_entity).add_child(child_entity);
+        let bundle = with_child_component.bundle.take().unwrap();
+        let child_entity = world.spawn(*bundle).id();
+
+        world
+            .entity_mut(self.parent_entity)
+            .insert(WithChildEntity {
+                child_entity,
+                _phantom: PhantomData::<B>,
+            })
+            .add_child(child_entity);
+    }
+}
+
+struct WithChildRemoveCommand<B> {
+    parent_entity: Entity,
+    _phantom: PhantomData<B>,
+}
+
+impl<B: Bundle> Command for WithChildRemoveCommand<B> {
+    fn apply(self, world: &mut World) {
+        let Some(entity_mut) = world.get_entity(self.parent_entity) else {
+            #[cfg(debug_assertions)]
+            panic!("Parent entity not found");
+
+            #[cfg(not(debug_assertions))]
+            return;
+        };
+
+        if let Some(with_child_component) = entity_mut.get::<WithChildEntity<B>>() {
+            world.despawn(with_child_component.child_entity);
+        };
     }
 }
 
@@ -228,11 +299,10 @@ mod tests {
     fn with_child() {
         let mut world = World::default();
 
-        let parent = world.spawn(WithChild((A, B(3)))).id();
+        let parent = world.spawn(WithChild::new((A, B(3)))).id();
         // FIXME: this should not be needed!
         world.flush();
 
-        assert!(!world.entity(parent).contains::<WithChild<(A, B)>>());
         assert!(!world.entity(parent).contains::<A>());
         assert!(!world.entity(parent).contains::<B>());
 
@@ -242,6 +312,13 @@ mod tests {
         let child_entity = children[0];
         assert_eq!(world.get::<A>(child_entity), Some(&A));
         assert_eq!(world.get::<B>(child_entity), Some(&B(3)));
+
+        world.entity_mut(parent).remove::<WithChild<(A, B)>>();
+        // FIXME: this should not be needed!
+        world.flush();
+
+        assert!(world.get::<A>(child_entity).is_none());
+        assert!(world.get::<B>(child_entity).is_none());
     }
 
     #[test]
@@ -286,7 +363,7 @@ mod tests {
     fn with_distinct_children() {
         let mut world = World::default();
 
-        let parent = world.spawn((WithChild(A), WithChild(B(1)))).id();
+        let parent = world.spawn((WithChild::new(A), WithChild::new(B(1)))).id();
         // FIXME: this should not be needed!
         world.flush();
 
@@ -296,7 +373,7 @@ mod tests {
         assert_eq!(world.get::<B>(children[1]), Some(&B(1)));
 
         // Ordering should matter
-        let parent = world.spawn((WithChild(B(1)), WithChild(A))).id();
+        let parent = world.spawn((WithChild::new(B(1)), WithChild::new(A))).id();
         // FIXME: this should not be needed!
         world.flush();
 
@@ -310,7 +387,9 @@ mod tests {
     fn grandchildren() {
         let mut world = World::default();
 
-        let parent = world.spawn(WithChild((A, WithChild((A, B(3)))))).id();
+        let parent = world
+            .spawn(WithChild::new((A, WithChild::new((A, B(3))))))
+            .id();
         // FIXME: this should not be needed!
         world.flush();
 
@@ -335,14 +414,13 @@ mod tests {
         let parent = world
             .spawn(HierarchicalBundle {
                 a: A,
-                child: WithChild(ABBundle { a: A, b: B(17) }),
+                child: WithChild::new(ABBundle { a: A, b: B(17) }),
             })
             .id();
 
         // FIXME: this should not be needed!
         world.flush();
 
-        assert!(!world.entity(parent).contains::<WithChild<ABBundle>>());
         assert!(world.entity(parent).contains::<A>());
         assert!(!world.entity(parent).contains::<B>());
 
@@ -357,13 +435,12 @@ mod tests {
     #[test]
     fn command_form() {
         fn spawn_with_child(mut commands: Commands) -> Entity {
-            commands.spawn((A, WithChild(B(5)))).id()
+            commands.spawn((A, WithChild::new(B(5)))).id()
         }
 
         let mut world = World::new();
         let parent = world.run_system_once(spawn_with_child);
 
-        assert!(!world.entity(parent).contains::<WithChild<B>>());
         assert!(world.entity(parent).contains::<A>());
         assert!(!world.entity(parent).contains::<B>());