diff --git a/src/hierarchy.rs b/src/hierarchy.rs index cae886f..ce88188 100644 --- a/src/hierarchy.rs +++ b/src/hierarchy.rs @@ -29,14 +29,34 @@ 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(pub B); +#[derive(Debug, Clone)] +pub struct WithChild { + /// 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>, +} + +impl Default for WithChild { + fn default() -> Self { + Self::new(B::default()) + } +} + +impl WithChild { + /// Create a new `WithChild` component with the given bundle. + pub fn new(bundle: B) -> Self { + Self { + bundle: Some(Box::new(bundle)), + } + } +} impl Component for WithChild { /// This is a sparse set component as it's only ever added and removed, never iterated over. @@ -44,6 +64,7 @@ impl Component for WithChild { fn register_component_hooks(hooks: &mut ComponentHooks) { hooks.on_add(with_child_hook::); + hooks.on_remove(with_remove_child_hook::); } } @@ -62,6 +83,27 @@ fn with_child_hook( }); } +/// A hook that runs whenever [`WithChild`] is added to an entity. +/// +/// Generates a [`WithChildCommand`]. +fn with_remove_child_hook( + 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::, + }); +} + +#[derive(Component)] +struct WithChildEntity { + child_entity: Entity, + _phantom: PhantomData, +} + struct WithChildCommand { parent_entity: Entity, _phantom: PhantomData, @@ -77,7 +119,7 @@ impl Command for WithChildCommand { return; }; - let Some(with_child_component) = entity_mut.take::>() else { + let Some(mut with_child_component) = entity_mut.get_mut::>() else { #[cfg(debug_assertions)] panic!("WithChild component not found"); @@ -85,8 +127,37 @@ impl Command for WithChildCommand { 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::, + }) + .add_child(child_entity); + } +} + +struct WithChildRemoveCommand { + parent_entity: Entity, + _phantom: PhantomData, +} + +impl Command for WithChildRemoveCommand { + 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::>() { + 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::>()); assert!(!world.entity(parent).contains::()); assert!(!world.entity(parent).contains::()); @@ -242,6 +312,13 @@ mod tests { let child_entity = children[0]; assert_eq!(world.get::(child_entity), Some(&A)); assert_eq!(world.get::(child_entity), Some(&B(3))); + + world.entity_mut(parent).remove::>(); + // FIXME: this should not be needed! + world.flush(); + + assert!(world.get::(child_entity).is_none()); + assert!(world.get::(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::(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::>()); assert!(world.entity(parent).contains::()); assert!(!world.entity(parent).contains::()); @@ -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::>()); assert!(world.entity(parent).contains::()); assert!(!world.entity(parent).contains::());