From 3d6d2431fe52039548153f8441b5fca6f535d144 Mon Sep 17 00:00:00 2001
From: Nathan Stocks <cleancut@github.com>
Date: Tue, 1 Sep 2020 17:39:22 -0600
Subject: [PATCH] Remove child from parent when it is despawned (#386)

Remove child from parent when it is despawned
---
 .../bevy_transform/src/hierarchy/hierarchy.rs | 57 +++++++++++++++----
 1 file changed, 45 insertions(+), 12 deletions(-)

diff --git a/crates/bevy_transform/src/hierarchy/hierarchy.rs b/crates/bevy_transform/src/hierarchy/hierarchy.rs
index 354c08d614312..b13f665e963e0 100644
--- a/crates/bevy_transform/src/hierarchy/hierarchy.rs
+++ b/crates/bevy_transform/src/hierarchy/hierarchy.rs
@@ -1,4 +1,4 @@
-use crate::components::Children;
+use crate::components::{Children, Parent};
 use bevy_ecs::{Commands, Entity, Query, World, WorldWriter};
 
 pub fn run_on_hierarchy<T, S>(
@@ -44,6 +44,18 @@ pub struct DespawnRecursive {
 }
 
 fn despawn_with_children_recursive(world: &mut World, entity: Entity) {
+    // first, make the entity's own parent forget about it
+    if let Ok(parent) = world.get::<Parent>(entity) {
+        if let Ok(mut children) = world.get_mut::<Children>(parent.0) {
+            children.retain(|c| *c != entity);
+        }
+    }
+    // then despawn the entity and all of its children
+    despawn_with_children_recursive_inner(world, entity);
+}
+
+// Should only be called by `despawn_with_children_recursive`!
+fn despawn_with_children_recursive_inner(world: &mut World, entity: Entity) {
     if let Some(children) = world
         .get::<Children>(entity)
         .ok()
@@ -78,32 +90,41 @@ impl DespawnRecursiveExt for Commands {
 #[cfg(test)]
 mod tests {
     use super::DespawnRecursiveExt;
-    use crate::hierarchy::BuildChildren;
-    use bevy_ecs::{Commands, Entity, Resources, World};
+    use crate::{components::Children, hierarchy::BuildChildren};
+    use bevy_ecs::{Commands, Resources, World};
 
     #[test]
     fn despawn_recursive() {
         let mut world = World::default();
         let mut resources = Resources::default();
         let mut command_buffer = Commands::default();
-        let parent_entity = Entity::new();
 
         command_buffer.spawn((0u32, 0u64)).with_children(|parent| {
             parent.spawn((0u32, 0u64));
         });
 
-        command_buffer
-            .spawn_as_entity(parent_entity, (1u32, 2u64))
-            .with_children(|parent| {
-                parent.spawn((1u32, 2u64)).with_children(|parent| {
-                    parent.spawn((1u32, 2u64));
+        // Create a grandparent entity which will _not_ be deleted
+        command_buffer.spawn((1u32, 1u64));
+        let grandparent_entity = command_buffer.current_entity().unwrap();
+
+        command_buffer.with_children(|parent| {
+            // Add a child to the grandparent (the "parent"), which will get deleted
+            parent.spawn((2u32, 2u64));
+            // All descendents of the "parent" should also be deleted.
+            parent.with_children(|parent| {
+                parent.spawn((3u32, 3u64)).with_children(|parent| {
+                    // child
+                    parent.spawn((4u32, 4u64));
                 });
-                parent.spawn((1u32, 2u64));
+                parent.spawn((5u32, 5u64));
             });
+        });
 
         command_buffer.spawn((0u32, 0u64));
         command_buffer.apply(&mut world, &mut resources);
 
+        let parent_entity = world.get::<Children>(grandparent_entity).unwrap()[0];
+
         command_buffer.despawn_recursive(parent_entity);
         command_buffer.apply(&mut world, &mut resources);
 
@@ -113,8 +134,20 @@ mod tests {
             .map(|(a, b)| (*a, *b))
             .collect::<Vec<_>>();
 
+        {
+            let children = world.get::<Children>(grandparent_entity).unwrap();
+            assert_eq!(
+                children.iter().any(|&i| i == parent_entity),
+                false,
+                "grandparent should no longer know about its child which has been removed"
+            );
+        }
+
         // parent_entity and its children should be deleted,
-        // the (0, 0) tuples remaining.
-        assert_eq!(results, vec![(0u32, 0u64), (0u32, 0u64), (0u32, 0u64)]);
+        // the grandparent tuple (1, 1) and (0, 0) tuples remaining.
+        assert_eq!(
+            results,
+            vec![(0u32, 0u64), (0u32, 0u64), (0u32, 0u64), (1u32, 1u64)]
+        );
     }
 }