Bugfix: Scene reload fix (nonbreaking)

crates/bevy_ecs/src/reflect.rs

@@ -413,21 +413,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;
@@ -86,6 +90,12 @@ impl DynamicScene {
             reflect_resource.apply_or_insert(world, &**resource);
         }
 
+        // 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
@@ -109,17 +119,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();
             }
         }
@@ -160,3 +183,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

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