Bugfix: Scene reload fix (nonbreaking) (#7951)

Fix a bug with scene reload.

(This is a copy of #7570 but without the breaking API change, in order
to allow the bugfix to be introduced in 0.10.1)

When a scene was reloaded, it was corrupting components that weren't
native to the scene itself. In particular, when a DynamicScene was
created on Entity (A), all components in the scene without parents are
automatically added as children of Entity (A). But if that scene was
reloaded and the same ID of Entity (A) was a scene ID as well*, that
parent component was corrupted, causing the hierarchy to become
malformed and bevy to panic.

*For example, if Entity (A)'s ID was 3, and the scene contained an
entity with ID 3

This issue could affect any components that:
* Implemented `MapEntities`, basically components that contained
references to other entities
* Were added to entities from a scene file but weren't defined in the
scene file

- Fixes #7529

The solution was to keep track of entities+components that had
`MapEntities` functionality during scene load, and only apply the entity
update behavior to them. They were tracked with a HashMap from the
component's TypeID to a vector of entity ID's. Then the
`ReflectMapEntities` struct was updated to hold a function that took a
list of entities to be applied to, instead of naively applying itself to
all values in the EntityMap.

(See this PR comment
#7570 (comment) for
a story-based explanation of this bug and solution)

- Components that implement `MapEntities` added to scene entities after
load are not corrupted during scene reload.

crates/bevy_ecs/src/reflect.rs

@@ -405,21 +405,53 @@ impl_from_reflect_value!(Entity);
 #[derive(Clone)]
 pub struct ReflectMapEntities {
     map_entities: fn(&mut World, &EntityMap) -> Result<(), MapEntitiesError>,
+    map_specific_entities: fn(&mut World, &EntityMap, &[Entity]) -> Result<(), MapEntitiesError>,
 }
 
 impl ReflectMapEntities {
+    /// A general method for applying [`MapEntities`] behavior to all elements in an [`EntityMap`].
+    ///
+    /// Be mindful in its usage: Works best in situations where the entities in the [`EntityMap`] are newly
+    /// created, before systems have a chance to add new components. If some of the entities referred to
+    /// by the [`EntityMap`] might already contain valid entity references, you should use [`map_entities`](Self::map_entities).
+    ///
+    /// An example of this: A scene can be loaded with `Parent` components, but then a `Parent` component can be added
+    /// to these entities after they have been loaded. If you reload the scene using [`map_entities`](Self::map_entities), those `Parent`
+    /// components with already valid entity references could be updated to point at something else entirely.
     pub fn map_entities(
         &self,
         world: &mut World,
         entity_map: &EntityMap,
     ) -> Result<(), MapEntitiesError> {
         (self.map_entities)(world, entity_map)
     }
+
+    /// This is like [`map_entities`](Self::map_entities), but only applied to specific entities, not all values
+    /// in the [`EntityMap`].
+    ///
+    /// This is useful mostly for when you need to be careful not to update components that already contain valid entity
+    /// values. See [`map_entities`](Self::map_entities) for more details.
+    pub fn map_specific_entities(
+        &self,
+        world: &mut World,
+        entity_map: &EntityMap,
+        entities: &[Entity],
+    ) -> Result<(), MapEntitiesError> {
+        (self.map_specific_entities)(world, entity_map, entities)
+    }
 }
 
 impl<C: Component + MapEntities> FromType<C> for ReflectMapEntities {
     fn from_type() -> Self {
         ReflectMapEntities {
+            map_specific_entities: |world, entity_map, entities| {
+                for &entity in entities {
+                    if let Some(mut component) = world.get_mut::<C>(entity) {
+                        component.map_entities(entity_map)?;
+                    }
+                }
+                Ok(())
+            },
             map_entities: |world, entity_map| {
                 for entity in entity_map.values() {
                     if let Some(mut component) = world.get_mut::<C>(entity) {

crates/bevy_scene/src/dynamic_scene.rs

@@ -1,12 +1,16 @@
+use std::any::TypeId;
+
 use crate::{DynamicSceneBuilder, Scene, SceneSpawnError};
 use anyhow::Result;
 use bevy_app::AppTypeRegistry;
 use bevy_ecs::{
     entity::EntityMap,
+    prelude::Entity,
     reflect::{ReflectComponent, ReflectMapEntities},
     world::World,
 };
 use bevy_reflect::{Reflect, TypeRegistryArc, TypeUuid};
+use bevy_utils::HashMap;
 
 #[cfg(feature = "serialize")]
 use crate::serde::SceneSerializer;
@@ -64,6 +68,12 @@ impl DynamicScene {
     ) -> Result<(), SceneSpawnError> {
         let type_registry = type_registry.read();
 
+        // For each component types that reference other entities, we keep track
+        // of which entities in the scene use that component.
+        // This is so we can update the scene-internal references to references
+        // of the actual entities in the world.
+        let mut scene_mappings: HashMap<TypeId, Vec<Entity>> = HashMap::default();
+
         for scene_entity in &self.entities {
             // Fetch the entity with the given entity id from the `entity_map`
             // or spawn a new entity with a transiently unique id if there is
@@ -87,17 +97,30 @@ impl DynamicScene {
                         }
                     })?;
 
+                // If this component references entities in the scene, track it
+                // so we can update it to the entity in the world.
+                if registration.data::<ReflectMapEntities>().is_some() {
+                    scene_mappings
+                        .entry(registration.type_id())
+                        .or_insert(Vec::new())
+                        .push(entity);
+                }
+
                 // If the entity already has the given component attached,
                 // just apply the (possibly) new value, otherwise add the
                 // component to the entity.
                 reflect_component.apply_or_insert(entity_mut, &**component);
             }
         }
 
-        for registration in type_registry.iter() {
+        // Updates references to entities in the scene to entities in the world
+        for (type_id, entities) in scene_mappings.into_iter() {
+            let registration = type_registry.get(type_id).expect(
+                "we should be getting TypeId from this TypeRegistration in the first place",
+            );
             if let Some(map_entities_reflect) = registration.data::<ReflectMapEntities>() {
                 map_entities_reflect
-                    .map_entities(world, entity_map)
+                    .map_specific_entities(world, entity_map, &entities)
                     .unwrap();
             }
         }
@@ -138,3 +161,89 @@ where
         .new_line("\n".to_string());
     ron::ser::to_string_pretty(&serialize, pretty_config)
 }
+
+#[cfg(test)]
+mod tests {
+    use bevy_app::AppTypeRegistry;
+    use bevy_ecs::{entity::EntityMap, system::Command, world::World};
+    use bevy_hierarchy::{AddChild, Parent};
+
+    use crate::dynamic_scene_builder::DynamicSceneBuilder;
+
+    #[test]
+    fn components_not_defined_in_scene_should_not_be_affected_by_scene_entity_map() {
+        // Testing that scene reloading applies EntityMap correctly to MapEntities components.
+
+        // First, we create a simple world with a parent and a child relationship
+        let mut world = World::new();
+        world.init_resource::<AppTypeRegistry>();
+        world
+            .resource_mut::<AppTypeRegistry>()
+            .write()
+            .register::<Parent>();
+        let original_parent_entity = world.spawn_empty().id();
+        let original_child_entity = world.spawn_empty().id();
+        AddChild {
+            parent: original_parent_entity,
+            child: original_child_entity,
+        }
+        .write(&mut world);
+
+        // We then write this relationship to a new scene, and then write that scene back to the
+        // world to create another parent and child relationship
+        let mut scene_builder = DynamicSceneBuilder::from_world(&world);
+        scene_builder.extract_entity(original_parent_entity);
+        scene_builder.extract_entity(original_child_entity);
+        let scene = scene_builder.build();
+        let mut entity_map = EntityMap::default();
+        scene.write_to_world(&mut world, &mut entity_map).unwrap();
+
+        let from_scene_parent_entity = entity_map.get(original_parent_entity).unwrap();
+        let from_scene_child_entity = entity_map.get(original_child_entity).unwrap();
+
+        // We then add the parent from the scene as a child of the original child
+        // Hierarchy should look like:
+        // Original Parent <- Original Child <- Scene Parent <- Scene Child
+        AddChild {
+            parent: original_child_entity,
+            child: from_scene_parent_entity,
+        }
+        .write(&mut world);
+
+        // We then reload the scene to make sure that from_scene_parent_entity's parent component
+        // isn't updated with the entity map, since this component isn't defined in the scene.
+        // With bevy_hierarchy, this can cause serious errors and malformed hierarchies.
+        scene.write_to_world(&mut world, &mut entity_map).unwrap();
+
+        assert_eq!(
+            original_parent_entity,
+            world
+                .get_entity(original_child_entity)
+                .unwrap()
+                .get::<Parent>()
+                .unwrap()
+                .get(),
+            "something about reloading the scene is touching entities with the same scene Ids"
+        );
+        assert_eq!(
+            original_child_entity,
+            world
+                .get_entity(from_scene_parent_entity)
+                .unwrap()
+                .get::<Parent>()
+                .unwrap()
+                .get(),
+            "something about reloading the scene is touching components not defined in the scene but on entities defined in the scene"
+        );
+        assert_eq!(
+            from_scene_parent_entity,
+            world
+                .get_entity(from_scene_child_entity)
+                .unwrap()
+                .get::<Parent>()
+                .expect("something is wrong with this test, and the scene components don't have a parent/child relationship")
+                .get(),
+            "something is wrong with the this test or the code reloading scenes since the relationship between scene entities is broken"
+        );
+    }
+}

crates/bevy_scene/src/scene.rs

@@ -90,6 +90,7 @@ impl Scene {
                 }
             }
         }
+
         for registration in type_registry.iter() {
             if let Some(map_entities_reflect) = registration.data::<ReflectMapEntities>() {
                 map_entities_reflect