Rework ReflectMapEntities to match MapEntities, reorder scene entity mapping

jrobsonchase · jrobsonchase · commit fe173a112436 · 2024-09-28T16:09:58.000-04:00
diff --git a/crates/bevy_ecs/src/entity/map_entities.rs b/crates/bevy_ecs/src/entity/map_entities.rs
@@ -112,6 +112,15 @@ pub trait DynEntityMapper {
     fn dyn_mappings(&self) -> Vec<(Entity, Entity)>;
 }
 
+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<T: EntityMapper> DynEntityMapper for T {
     fn dyn_map_entity(&mut self, entity: Entity) -> Entity {
         <T as EntityMapper>::map_entity(self, entity)
diff --git a/crates/bevy_ecs/src/reflect/map_entities.rs b/crates/bevy_ecs/src/reflect/map_entities.rs
@@ -1,107 +1,41 @@
-use crate::{
-    component::Component,
-    entity::{Entity, EntityHashMap, MapEntities, SceneEntityMapper},
-    world::World,
-};
-use bevy_reflect::FromType;
+use crate::entity::{DynEntityMapper, MapEntities};
+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.
+/// For a specific type of value, 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
 /// any stored IDs need to be re-allocated in the destination world.
 ///
-/// See [`SceneEntityMapper`] and [`MapEntities`] for more information.
+/// See [`EntityMapper`] and [`MapEntities`] for more information.
+///
+/// [`Entity`]: crate::entity::Entity
+/// [`EntityMapper`]: crate::entity::EntityMapper
 #[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),
 }
 
 impl ReflectMapEntities {
-    /// 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).
-    ///
-    /// 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`
-    /// 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);
-    }
-
-    /// 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
-    /// in the [`EntityHashMap<Entity>`].
+    /// A general method for remapping entities in a reflected value via a [`DynEntityMapper`].
     ///
-    /// 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.
+    /// # Panics
+    /// Panics if the the type of the reflected value doesn't match.
     pub fn map_entities(
         &self,
-        world: &mut World,
-        entity_map: &mut EntityHashMap<Entity>,
-        entities: &[Entity],
+        reflected: &mut dyn PartialReflect,
+        mapper: &mut dyn DynEntityMapper,
     ) {
-        SceneEntityMapper::world_scope(entity_map, world, |world, mapper| {
-            (self.map_entities)(world, mapper, entities);
-        });
+        (self.map_entities)(reflected, mapper);
     }
 }
 
-impl<C: Component + MapEntities> FromType<C> for ReflectMapEntities {
+impl<C: FromReflect + MapEntities> FromType<C> for ReflectMapEntities {
     fn from_type() -> Self {
         ReflectMapEntities {
-            map_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| {
-                let entities = entity_mapper
-                    .get_map()
-                    .values()
-                    .copied()
-                    .collect::<Vec<Entity>>();
-                for entity in &entities {
-                    if let Some(mut component) = world.get_mut::<C>(*entity) {
-                        component.map_entities(entity_mapper);
-                    }
-                }
-            },
-        }
-    }
-}
-
-/// For a specific type of resource, 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
-/// any stored IDs need to be re-allocated in the destination world.
-///
-/// See [`SceneEntityMapper`] and [`MapEntities`] for more information.
-#[derive(Clone)]
-pub struct ReflectMapEntitiesResource {
-    map_entities: fn(&mut World, &mut SceneEntityMapper),
-}
-
-impl ReflectMapEntitiesResource {
-    /// A method for applying [`MapEntities`] behavior to elements in an [`EntityHashMap<Entity>`].
-    pub fn map_entities(&self, world: &mut World, entity_map: &mut EntityHashMap<Entity>) {
-        SceneEntityMapper::world_scope(entity_map, world, |world, mapper| {
-            (self.map_entities)(world, mapper);
-        });
-    }
-}
-
-impl<R: crate::system::Resource + MapEntities> FromType<R> for ReflectMapEntitiesResource {
-    fn from_type() -> Self {
-        ReflectMapEntitiesResource {
-            map_entities: |world, entity_mapper| {
-                if let Some(mut resource) = world.get_resource_mut::<R>() {
-                    resource.map_entities(entity_mapper);
-                }
+            map_entities: |reflected, mut mapper| {
+                let mut concrete = C::from_reflect(reflected).expect("mismatched reflected type");
+                concrete.map_entities(&mut mapper);
+                reflected.apply(&concrete);
             },
         }
     }
diff --git a/crates/bevy_ecs/src/reflect/mod.rs b/crates/bevy_ecs/src/reflect/mod.rs
@@ -23,7 +23,7 @@ pub use bundle::{ReflectBundle, ReflectBundleFns};
 pub use component::{ReflectComponent, ReflectComponentFns};
 pub use entity_commands::ReflectCommandExt;
 pub use from_world::{ReflectFromWorld, ReflectFromWorldFns};
-pub use map_entities::{ReflectMapEntities, ReflectMapEntitiesResource};
+pub use map_entities::ReflectMapEntities;
 pub use resource::{ReflectResource, ReflectResourceFns};
 
 /// A [`Resource`] storing [`TypeRegistry`] for
diff --git a/crates/bevy_scene/src/dynamic_scene.rs b/crates/bevy_scene/src/dynamic_scene.rs
@@ -1,16 +1,15 @@
 use crate::{ron, DynamicSceneBuilder, Scene, SceneSpawnError};
+use bevy_asset::Asset;
+use bevy_ecs::reflect::ReflectResource;
 use bevy_ecs::{
-    entity::{Entity, EntityHashMap},
+    entity::{Entity, EntityHashMap, SceneEntityMapper},
     reflect::{AppTypeRegistry, ReflectComponent, ReflectMapEntities},
     world::World,
 };
 use bevy_reflect::{PartialReflect, TypePath, TypeRegistry};
-use bevy_utils::TypeIdMap;
 
 #[cfg(feature = "serialize")]
 use crate::serde::SceneSerializer;
-use bevy_asset::Asset;
-use bevy_ecs::reflect::{ReflectMapEntitiesResource, ReflectResource};
 #[cfg(feature = "serialize")]
 use serde::Serialize;
 
@@ -70,23 +69,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(),
@@ -104,39 +106,26 @@ 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 {
+            let mut resource = resource.clone_value();
             let type_info = resource.get_represented_type_info().ok_or_else(|| {
                 SceneSpawnError::NoRepresentedType {
                     type_path: resource.reflect_type_path().to_string(),
@@ -153,14 +142,17 @@ impl DynamicScene {
                 }
             })?;
 
+            // 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(resource.as_partial_reflect_mut(), mapper);
+                });
+            }
+
             // If the world already contains an instance of the given resource
             // just apply the (possibly) new value, otherwise insert the resource
-            reflect_resource.apply_or_insert(world, &**resource, &type_registry);
-
-            // Map entities in the resource if it implements [`MapEntities`].
-            if let Some(map_entities_reflect) = registration.data::<ReflectMapEntitiesResource>() {
-                map_entities_reflect.map_entities(world, entity_map);
-            }
+            reflect_resource.apply_or_insert(world, resource.as_partial_reflect(), &type_registry);
         }
 
         Ok(())
@@ -211,10 +203,7 @@ mod tests {
     use bevy_ecs::{
         component::Component,
         entity::{Entity, EntityHashMap, EntityMapper, MapEntities},
-        reflect::{
-            AppTypeRegistry, ReflectComponent, ReflectMapEntities, ReflectMapEntitiesResource,
-            ReflectResource,
-        },
+        reflect::{AppTypeRegistry, ReflectComponent, ReflectMapEntities, ReflectResource},
         system::Resource,
         world::{Command, World},
     };
@@ -225,7 +214,7 @@ mod tests {
     use crate::dynamic_scene_builder::DynamicSceneBuilder;
 
     #[derive(Resource, Reflect, Debug)]
-    #[reflect(Resource, MapEntitiesResource)]
+    #[reflect(Resource, MapEntities)]
     struct TestResource {
         entity_a: Entity,
         entity_b: Entity,
diff --git a/crates/bevy_scene/src/scene.rs b/crates/bevy_scene/src/scene.rs
@@ -1,11 +1,11 @@
 use crate::{DynamicScene, SceneSpawnError};
 use bevy_asset::Asset;
 use bevy_ecs::{
-    entity::{Entity, EntityHashMap},
+    entity::{Entity, EntityHashMap, SceneEntityMapper},
     reflect::{AppTypeRegistry, ReflectComponent, ReflectMapEntities, ReflectResource},
     world::World,
 };
-use bevy_reflect::TypePath;
+use bevy_reflect::{PartialReflect, TypePath};
 
 /// To spawn a scene, you can use either:
 /// * [`SceneSpawner::spawn`](crate::SceneSpawner::spawn)
@@ -90,47 +90,64 @@ impl Scene {
             reflect_resource.copy(&self.world, world, &type_registry);
         }
 
+        // Ensure that all scene entities have been allocated in the destination
+        // world before handling components that may contain references that need mapping.
         for archetype in self.world.archetypes().iter() {
             for scene_entity in archetype.entities() {
-                let entity = entity_map
+                entity_map
                     .entry(scene_entity.id())
                     .or_insert_with(|| world.spawn_empty().id());
+            }
+        }
+
+        for archetype in self.world.archetypes().iter() {
+            for scene_entity in archetype.entities() {
+                let entity = *entity_map
+                    .get(&scene_entity.id())
+                    .expect("should have previously spawned an entity");
+
                 for component_id in archetype.components() {
                     let component_info = self
                         .world
                         .components()
                         .get_info(component_id)
                         .expect("component_ids in archetypes should have ComponentInfo");
 
-                    let reflect_component = type_registry
+                    let registration = type_registry
                         .get(component_info.type_id().unwrap())
                         .ok_or_else(|| SceneSpawnError::UnregisteredType {
                             std_type_name: component_info.name().to_string(),
-                        })
-                        .and_then(|registration| {
-                            registration.data::<ReflectComponent>().ok_or_else(|| {
-                                SceneSpawnError::UnregisteredComponent {
-                                    type_path: registration.type_info().type_path().to_string(),
-                                }
-                            })
                         })?;
-                    reflect_component.copy(
-                        &self.world,
-                        world,
-                        scene_entity.id(),
-                        *entity,
+                    let reflect_component =
+                        registration.data::<ReflectComponent>().ok_or_else(|| {
+                            SceneSpawnError::UnregisteredComponent {
+                                type_path: registration.type_info().type_path().to_string(),
+                            }
+                        })?;
+
+                    let Some(mut component) = reflect_component
+                        .reflect(self.world.entity(scene_entity.id()))
+                        .map(PartialReflect::clone_value)
+                    else {
+                        continue;
+                    };
+
+                    // 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);
+                        });
+                    }
+                    reflect_component.apply_or_insert(
+                        &mut world.entity_mut(entity),
+                        component.as_partial_reflect(),
                         &type_registry,
                     );
                 }
             }
         }
 
-        for registration in type_registry.iter() {
-            if let Some(map_entities_reflect) = registration.data::<ReflectMapEntities>() {
-                map_entities_reflect.map_all_entities(world, entity_map);
-            }
-        }
-
         Ok(())
     }
 }
