Change ReflectMapEntities to operate on components before insertion

crates/bevy_ecs/src/entity/map_entities.rs

@@ -122,6 +122,16 @@ impl<T: EntityMapper> DynEntityMapper for T {
     }
 }
 
+impl<'a> EntityMapper for &'a mut dyn DynEntityMapper {
+    fn map_entity(&mut self, entity: Entity) -> Entity {
+        (*self).dyn_map_entity(entity)
+    }
+
+    fn mappings(&self) -> impl Iterator<Item = (Entity, Entity)> {
+        (*self).dyn_mappings().into_iter()
+    }
+}
+
 impl EntityMapper for SceneEntityMapper<'_> {
     /// Returns the corresponding mapped entity or reserves a new dead entity ID in the current world if it is absent.
     fn map_entity(&mut self, entity: Entity) -> Entity {

crates/bevy_ecs/src/reflect/map_entities.rs

@@ -1,9 +1,9 @@
 use crate::{
     component::Component,
-    entity::{Entity, EntityHashMap, MapEntities, SceneEntityMapper},
+    entity::{DynEntityMapper, Entity, EntityHashMap, MapEntities, SceneEntityMapper},
     world::World,
 };
-use bevy_reflect::FromType;
+use bevy_reflect::{FromReflect, FromType, PartialReflect};
 
 /// For a specific type of component, this maps any fields with values of type [`Entity`] to a new world.
 /// Since a given `Entity` ID is only valid for the world it came from, when performing deserialization
@@ -12,53 +12,81 @@ use bevy_reflect::FromType;
 /// See [`SceneEntityMapper`] and [`MapEntities`] for more information.
 #[derive(Clone)]
 pub struct ReflectMapEntities {
-    map_all_entities: fn(&mut World, &mut SceneEntityMapper),
-    map_entities: fn(&mut World, &mut SceneEntityMapper, &[Entity]),
+    map_entities: fn(&mut dyn PartialReflect, &mut dyn DynEntityMapper),
+    map_all_world_entities: fn(&mut World, &mut SceneEntityMapper),
+    map_world_entities: fn(&mut World, &mut SceneEntityMapper, &[Entity]),
 }
 
 impl ReflectMapEntities {
+    /// A general method for applying [`MapEntities`] behavior to a reflected component.
+    ///
+    /// Be mindful in its usage: Works best in situations where the source entities
+    /// in the [`EntityHashMap<Entity>`] have already been populated by spawning empty
+    /// entities in the destination world if needed. For example, when spawning
+    /// entities in a scene, if this is used on a component before ensuring that
+    /// all entities in the scene have been allocated, a new mapping will be created
+    /// with a "dead" entity.
+    pub fn map_entities(
+        &self,
+        component: &mut dyn PartialReflect,
+        mapper: &mut dyn DynEntityMapper,
+    ) {
+        (self.map_entities)(component, mapper);
+    }
+
     /// A general method for applying [`MapEntities`] behavior to all elements in an [`EntityHashMap<Entity>`].
     ///
     /// Be mindful in its usage: Works best in situations where the entities in the [`EntityHashMap<Entity>`] are newly
     /// created, before systems have a chance to add new components. If some of the entities referred to
-    /// by the [`EntityHashMap<Entity>`] might already contain valid entity references, you should use [`map_entities`](Self::map_entities).
+    /// by the [`EntityHashMap<Entity>`] might already contain valid entity references, you should use [`map_world_entities`](Self::map_world_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_all_entities`](Self::map_all_entities), those `Parent`
+    /// to these entities after they have been loaded. If you reload the scene using [`map_all_world_entities`](Self::map_all_world_entities), those `Parent`
     /// components with already valid entity references could be updated to point at something else entirely.
-    pub fn map_all_entities(&self, world: &mut World, entity_map: &mut EntityHashMap<Entity>) {
-        SceneEntityMapper::world_scope(entity_map, world, self.map_all_entities);
+    #[deprecated = "map_all_world_entities doesn't play well with Observers. Use map_entities instead."]
+    pub fn map_all_world_entities(
+        &self,
+        world: &mut World,
+        entity_map: &mut EntityHashMap<Entity>,
+    ) {
+        SceneEntityMapper::world_scope(entity_map, world, self.map_all_world_entities);
     }
 
     /// A general method for applying [`MapEntities`] behavior to elements in an [`EntityHashMap<Entity>`]. Unlike
-    /// [`map_all_entities`](Self::map_all_entities), this is applied to specific entities, not all values
+    /// [`map_all_world_entities`](Self::map_all_world_entities), this is applied to specific entities, not all values
     /// in the [`EntityHashMap<Entity>`].
     ///
     /// This is useful mostly for when you need to be careful not to update components that already contain valid entity
-    /// values. See [`map_all_entities`](Self::map_all_entities) for more details.
-    pub fn map_entities(
+    /// values. See [`map_all_world_entities`](Self::map_all_world_entities) for more details.
+    #[deprecated = "map_world_entities doesn't play well with Observers. Use map_entities instead."]
+    pub fn map_world_entities(
         &self,
         world: &mut World,
         entity_map: &mut EntityHashMap<Entity>,
         entities: &[Entity],
     ) {
         SceneEntityMapper::world_scope(entity_map, world, |world, mapper| {
-            (self.map_entities)(world, mapper, entities);
+            (self.map_world_entities)(world, mapper, entities);
         });
     }
 }
 
-impl<C: Component + MapEntities> FromType<C> for ReflectMapEntities {
+impl<C: Component + MapEntities + FromReflect> FromType<C> for ReflectMapEntities {
     fn from_type() -> Self {
         ReflectMapEntities {
-            map_entities: |world, entity_mapper, entities| {
+            map_entities: |component, mut entity_mapper| {
+                let mut concrete = C::from_reflect(component.as_partial_reflect()).unwrap();
+                concrete.map_entities(&mut entity_mapper);
+                component.apply(&concrete);
+            },
+            map_world_entities: |world, entity_mapper, entities| {
                 for &entity in entities {
                     if let Some(mut component) = world.get_mut::<C>(entity) {
                         component.map_entities(entity_mapper);
                     }
                 }
             },
-            map_all_entities: |world, entity_mapper| {
+            map_all_world_entities: |world, entity_mapper| {
                 let entities = entity_mapper
                     .get_map()
                     .values()

crates/bevy_scene/src/dynamic_scene.rs

@@ -1,12 +1,11 @@
 use crate::{ron, DynamicSceneBuilder, Scene, SceneSpawnError};
-use bevy_ecs::entity::EntityHashMap;
+use bevy_ecs::entity::{EntityHashMap, SceneEntityMapper};
 use bevy_ecs::{
     entity::Entity,
     reflect::{AppTypeRegistry, ReflectComponent, ReflectMapEntities},
     world::World,
 };
 use bevy_reflect::{PartialReflect, TypePath, TypeRegistry};
-use bevy_utils::TypeIdMap;
 
 #[cfg(feature = "serialize")]
 use crate::serde::SceneSerializer;
@@ -71,23 +70,26 @@ 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: TypeIdMap<Vec<Entity>> = Default::default();
-
+        // First ensure that every entity in the scene has a corresponding world
+        // entity in the entity map.
         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
             // no corresponding entry.
-            let entity = *entity_map
+            entity_map
                 .entry(scene_entity.entity)
                 .or_insert_with(|| world.spawn_empty().id());
-            let entity_mut = &mut world.entity_mut(entity);
+        }
+
+        for scene_entity in &self.entities {
+            // Fetch the entity with the given entity id from the `entity_map`.
+            let entity = *entity_map
+                .get(&scene_entity.entity)
+                .expect("should have previously spawned an empty entity");
 
             // Apply/ add each component to the given entity.
             for component in &scene_entity.components {
+                let mut component = component.clone_value();
                 let type_info = component.get_represented_type_info().ok_or_else(|| {
                     SceneSpawnError::NoRepresentedType {
                         type_path: component.reflect_type_path().to_string(),
@@ -105,36 +107,22 @@ 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_default()
-                        .push(entity);
+                // If this component references entities in the scene, update
+                // them to the entities in the world.
+                if let Some(map_entities) = registration.data::<ReflectMapEntities>() {
+                    SceneEntityMapper::world_scope(entity_map, world, |_, mapper| {
+                        map_entities.map_entities(component.as_partial_reflect_mut(), mapper);
+                    });
                 }
 
-                // 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,
+                    &mut world.entity_mut(entity),
                     component.as_partial_reflect(),
                     &type_registry,
                 );
             }
         }
 
-        // 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, &entities);
-            }
-        }
-
         // Insert resources after all entities have been added to the world.
         // This ensures the entities are available for the resources to reference during mapping.
         for resource in &self.resources {
@@ -209,11 +197,14 @@ where
 
 #[cfg(test)]
 mod tests {
+    use bevy_app::App;
     use bevy_ecs::entity::{Entity, EntityHashMap, EntityMapper, MapEntities};
+    use bevy_ecs::observer::Trigger;
     use bevy_ecs::reflect::{ReflectMapEntitiesResource, ReflectResource};
-    use bevy_ecs::system::Resource;
+    use bevy_ecs::system::{Query, Resource};
+    use bevy_ecs::world::OnAdd;
     use bevy_ecs::{reflect::AppTypeRegistry, world::Command, world::World};
-    use bevy_hierarchy::{AddChild, Parent};
+    use bevy_hierarchy::{AddChild, BuildChildren, Children, Parent};
     use bevy_reflect::Reflect;
 
     use crate::dynamic_scene_builder::DynamicSceneBuilder;
@@ -347,4 +338,52 @@ mod tests {
             "something is wrong with the this test or the code reloading scenes since the relationship between scene entities is broken"
         );
     }
+
+    #[test]
+    fn observers_should_see_mapped_entities() {
+        // Testing that Observers don't see the entities from the scene, but rather from the world itself.
+
+        let mut app = App::new();
+        app.init_resource::<AppTypeRegistry>()
+            .register_type::<Parent>()
+            .observe(|trigger: Trigger<OnAdd, Parent>, query: Query<&Parent>| {
+                let parent = query
+                    .get(trigger.entity())
+                    .expect("entity should have a parent");
+
+                assert_ne!(parent.get(), Entity::PLACEHOLDER);
+            });
+
+        let mut scene_world = World::new();
+        scene_world.insert_resource(app.world().resource::<AppTypeRegistry>().clone());
+
+        for _ in 0..5 {
+            scene_world.spawn_empty();
+        }
+
+        let scene_child = scene_world.spawn_empty().id();
+        let scene_parent = scene_world.spawn_empty().add_children(&[scene_child]).id();
+
+        app.observe(
+            move |trigger: Trigger<OnAdd, Parent>, query: Query<&Parent>| {
+                let parent = query
+                    .get(trigger.entity())
+                    .expect("entity should have a parent");
+
+                assert_ne!(parent.get(), scene_parent);
+            },
+        );
+
+        let scene = DynamicSceneBuilder::from_world(&scene_world)
+            .allow_component::<Parent>()
+            .allow_component::<Children>()
+            .extract_entities([scene_parent, scene_child].into_iter())
+            .build();
+
+        let mut entity_map = EntityHashMap::<Entity>::default();
+
+        scene
+            .write_to_world(app.world_mut(), &mut entity_map)
+            .expect("write scene to world");
+    }
 }