From b2a48e1f7bb8f841990a865b3443fddc3b4552a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois?= <mockersf@gmail.com>
Date: Thu, 23 Dec 2021 23:36:21 +0000
Subject: [PATCH] fix parenting of scenes (#2410)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

# Objective

Fix #2406

Scene parenting was not done completely, leaving the hierarchy maintenance to the standard system. As scene spawning happens in stage `PreUpdate` and hierarchy maintenance in stage `PostUpdate`, this left the scene in an invalid state parent wise for part of a frame

## Solution

Also add/update the `Children` component when spawning the scene.

I kept the `Children` component as a `SmallVec`, it could be moved to an `HashSet` to guarantee uniqueness


Co-authored-by: François <8672791+mockersf@users.noreply.github.com>
---
 crates/bevy_scene/src/scene_spawner.rs        | 21 +++++++++++---
 .../src/hierarchy/child_builder.rs            | 29 +++++++++++++++++++
 .../hierarchy/hierarchy_maintenance_system.rs |  9 +++---
 3 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/crates/bevy_scene/src/scene_spawner.rs b/crates/bevy_scene/src/scene_spawner.rs
index e6570a20f9ed23..7015195757d79f 100644
--- a/crates/bevy_scene/src/scene_spawner.rs
+++ b/crates/bevy_scene/src/scene_spawner.rs
@@ -4,10 +4,11 @@ use bevy_asset::{AssetEvent, Assets, Handle};
 use bevy_ecs::{
     entity::{Entity, EntityMap},
     reflect::{ReflectComponent, ReflectMapEntities},
+    system::Command,
     world::{Mut, World},
 };
 use bevy_reflect::TypeRegistryArc;
-use bevy_transform::prelude::Parent;
+use bevy_transform::{hierarchy::AddChild, prelude::Parent};
 use bevy_utils::{tracing::error, HashMap};
 use thiserror::Error;
 use uuid::Uuid;
@@ -268,10 +269,22 @@ impl SceneSpawner {
         for (instance_id, parent) in scenes_with_parent {
             if let Some(instance) = self.spawned_instances.get(&instance_id) {
                 for entity in instance.entity_map.values() {
-                    if let Some(mut entity_mut) = world.get_entity_mut(entity) {
-                        if !entity_mut.contains::<Parent>() {
-                            entity_mut.insert(Parent(parent));
+                    // Add the `Parent` component to the scene root, and update the `Children` component of
+                    // the scene parent
+                    if !world
+                        .get_entity(entity)
+                        // This will filter only the scene root entity, as all other from the
+                        // scene have a parent
+                        .map(|entity| entity.contains::<Parent>())
+                        // Default is true so that it won't run on an entity that wouldn't exist anymore
+                        // this case shouldn't happen anyway
+                        .unwrap_or(true)
+                    {
+                        AddChild {
+                            parent,
+                            child: entity,
                         }
+                        .write(world);
                     }
                 }
             } else {
diff --git a/crates/bevy_transform/src/hierarchy/child_builder.rs b/crates/bevy_transform/src/hierarchy/child_builder.rs
index 094a2af1314bd7..56f6a073d8b5bf 100644
--- a/crates/bevy_transform/src/hierarchy/child_builder.rs
+++ b/crates/bevy_transform/src/hierarchy/child_builder.rs
@@ -7,6 +7,28 @@ use bevy_ecs::{
 };
 use smallvec::SmallVec;
 
+#[derive(Debug)]
+pub struct AddChild {
+    pub parent: Entity,
+    pub child: Entity,
+}
+
+impl Command for AddChild {
+    fn write(self, world: &mut World) {
+        world
+            .entity_mut(self.child)
+            // FIXME: don't erase the previous parent (see #1545)
+            .insert_bundle((Parent(self.parent), PreviousParent(self.parent)));
+        if let Some(mut children) = world.get_mut::<Children>(self.parent) {
+            children.0.push(self.child);
+        } else {
+            world
+                .entity_mut(self.parent)
+                .insert(Children(smallvec::smallvec![self.child]));
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct InsertChildren {
     parent: Entity,
@@ -134,6 +156,7 @@ pub trait BuildChildren {
     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;
+    fn add_child(&mut self, child: Entity) -> &mut Self;
 }
 
 impl<'w, 's, 'a> BuildChildren for EntityCommands<'w, 's, 'a> {
@@ -182,6 +205,12 @@ impl<'w, 's, 'a> BuildChildren for EntityCommands<'w, 's, 'a> {
         });
         self
     }
+
+    fn add_child(&mut self, child: Entity) -> &mut Self {
+        let parent = self.id();
+        self.commands().add(AddChild { child, parent });
+        self
+    }
 }
 
 #[derive(Debug)]
diff --git a/crates/bevy_transform/src/hierarchy/hierarchy_maintenance_system.rs b/crates/bevy_transform/src/hierarchy/hierarchy_maintenance_system.rs
index 6fde5d52ebfbbe..75b5dfc1a2f52c 100644
--- a/crates/bevy_transform/src/hierarchy/hierarchy_maintenance_system.rs
+++ b/crates/bevy_transform/src/hierarchy/hierarchy_maintenance_system.rs
@@ -49,11 +49,10 @@ pub fn parent_update_system(
         // `children_additions`).
         if let Ok(mut new_parent_children) = children_query.get_mut(parent.0) {
             // This is the parent
-            debug_assert!(
-                !(*new_parent_children).0.contains(&entity),
-                "children already added"
-            );
-            (*new_parent_children).0.push(entity);
+            // PERF: Ideally we shouldn't need to check for duplicates
+            if !(*new_parent_children).0.contains(&entity) {
+                (*new_parent_children).0.push(entity);
+            }
         } else {
             // The parent doesn't have a children entity, lets add it
             children_additions