Change ReflectMapEntities to operate on components before insertion (#15422)

Previous PR #14549 was closed in
error and couldn't be reopened since I had updated the branch
:crying_cat_face:

# Objective

Fixes #14465

## Solution

`ReflectMapEntities` now works similarly to `MapEntities` in that it
works on the reflected value itself rather than the component in the
world after insertion. This makes it so that observers see the remapped
entities on insertion rather than the entity IDs from the scene.

`ReflectMapEntities` now works for both components and resources, so we
only need the one.

## Testing

* New unit test for `Observer`s + `DynamicScene`s
* New unit test for `Observer`s + `Scene`s
* Open to suggestions for other tests!

---

## Migration Guide

- Consumers of `ReflectMapEntities` will need to call `map_entities` on
values prior to inserting them into the world.
- Implementors of `MapEntities` will need to remove the `mappings`
method, which is no longer needed for `ReflectMapEntities` and has been
removed from the trait.

---------

Co-authored-by: Alice Cecile <[email protected]>
Co-authored-by: Hennadii Chernyshchyk <[email protected]>

Author: 3 people

crates/bevy_ecs/src/entity/map_entities.rs

@@ -15,9 +15,16 @@ use super::{EntityHashMap, VisitEntitiesMut};
 /// (usually by using an [`EntityHashMap<Entity>`] between source entities and entities in the
 /// current world).
 ///
+/// This trait is similar to [`VisitEntitiesMut`]. They differ in that [`VisitEntitiesMut`] operates
+/// on `&mut Entity` and allows for in-place modification, while this trait makes no assumption that
+/// such in-place modification is occurring, which is impossible for types such as [`HashSet<Entity>`]
+/// and [`EntityHashMap`] which must be rebuilt when their contained [`Entity`]s are remapped.
+///
 /// Implementing this trait correctly is required for properly loading components
 /// with entity references from scenes.
 ///
+/// [`HashSet<Entity>`]: bevy_utils::HashSet
+///
 /// ## Example
 ///
 /// ```
@@ -60,9 +67,6 @@ impl<T: VisitEntitiesMut> MapEntities for T {
 ///
 /// More generally, this can be used to map [`Entity`] references between any two [`Worlds`](World).
 ///
-/// Note that this trait is _not_ [object safe](https://doc.rust-lang.org/reference/items/traits.html#object-safety).
-/// Please see [`DynEntityMapper`] for an object safe alternative.
-///
 /// ## Example
 ///
 /// ```
@@ -79,64 +83,16 @@ impl<T: VisitEntitiesMut> MapEntities for T {
 ///     fn map_entity(&mut self, entity: Entity) -> Entity {
 ///         self.map.get(&entity).copied().unwrap_or(entity)
 ///     }
-///
-///     fn mappings(&self) -> impl Iterator<Item = (Entity, Entity)> {
-///         self.map.iter().map(|(&source, &target)| (source, target))
-///     }
 /// }
 /// ```
 pub trait EntityMapper {
     /// Map an entity to another entity
     fn map_entity(&mut self, entity: Entity) -> Entity;
-
-    /// Iterate over all entity to entity mappings.
-    ///
-    /// # Examples
-    ///
-    /// ```rust
-    /// # use bevy_ecs::entity::{Entity, EntityMapper};
-    /// # fn example(mapper: impl EntityMapper) {
-    /// for (source, target) in mapper.mappings() {
-    ///     println!("Will map from {source} to {target}");
-    /// }
-    /// # }
-    /// ```
-    fn mappings(&self) -> impl Iterator<Item = (Entity, Entity)>;
-}
-
-/// An [object safe](https://doc.rust-lang.org/reference/items/traits.html#object-safety) version
-/// of [`EntityMapper`]. This trait is automatically implemented for type that implements `EntityMapper`.
-pub trait DynEntityMapper {
-    /// Map an entity to another entity.
-    ///
-    /// This is an [object safe](https://doc.rust-lang.org/reference/items/traits.html#object-safety)
-    /// alternative to [`EntityMapper::map_entity`].
-    fn dyn_map_entity(&mut self, entity: Entity) -> Entity;
-
-    /// Iterate over all entity to entity mappings.
-    ///
-    /// This is an [object safe](https://doc.rust-lang.org/reference/items/traits.html#object-safety)
-    /// alternative to [`EntityMapper::mappings`].
-    fn dyn_mappings(&self) -> Vec<(Entity, Entity)>;
-}
-
-impl<T: EntityMapper> DynEntityMapper for T {
-    fn dyn_map_entity(&mut self, entity: Entity) -> Entity {
-        <T as EntityMapper>::map_entity(self, entity)
-    }
-
-    fn dyn_mappings(&self) -> Vec<(Entity, Entity)> {
-        <T as EntityMapper>::mappings(self).collect()
-    }
 }
 
-impl<'a> EntityMapper for &'a mut dyn DynEntityMapper {
+impl EntityMapper for &mut dyn EntityMapper {
     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()
+        (*self).map_entity(entity)
     }
 }
 
@@ -160,10 +116,6 @@ impl EntityMapper for SceneEntityMapper<'_> {
 
         new
     }
-
-    fn mappings(&self) -> impl Iterator<Item = (Entity, Entity)> {
-        self.map.iter().map(|(&source, &target)| (source, target))
-    }
 }
 
 /// A wrapper for [`EntityHashMap<Entity>`], augmenting it with the ability to allocate new [`Entity`] references in a destination
@@ -242,10 +194,9 @@ impl<'m> SceneEntityMapper<'m> {
 #[cfg(test)]
 mod tests {
     use crate::{
-        entity::{DynEntityMapper, Entity, EntityHashMap, EntityMapper, SceneEntityMapper},
+        entity::{Entity, EntityHashMap, EntityMapper, SceneEntityMapper},
         world::World,
     };
-    use bevy_utils::assert_object_safe;
 
     #[test]
     fn entity_mapper() {
@@ -292,26 +243,6 @@ mod tests {
         assert!(entity.generation() > dead_ref.generation());
     }
 
-    #[test]
-    fn entity_mapper_iteration() {
-        let mut old_world = World::new();
-        let mut new_world = World::new();
-
-        let mut map = EntityHashMap::default();
-        let mut mapper = SceneEntityMapper::new(&mut map, &mut new_world);
-
-        assert_eq!(mapper.mappings().collect::<Vec<_>>(), vec![]);
-
-        let old_entity = old_world.spawn_empty().id();
-
-        let new_entity = mapper.map_entity(old_entity);
-
-        assert_eq!(
-            mapper.mappings().collect::<Vec<_>>(),
-            vec![(old_entity, new_entity)]
-        );
-    }
-
     #[test]
     fn entity_mapper_no_panic() {
         let mut world = World::new();
@@ -328,9 +259,4 @@ mod tests {
         // The SceneEntityMapper should leave `Entities` in a flushed state.
         assert!(!world.entities.needs_flush());
     }
-
-    #[test]
-    fn dyn_entity_mapper_object_safe() {
-        assert_object_safe::<dyn DynEntityMapper>();
-    }
 }

crates/bevy_ecs/src/reflect/map_entities.rs

@@ -1,107 +1,37 @@
-use crate::{
-    component::Component,
-    entity::{Entity, EntityHashMap, MapEntities, SceneEntityMapper},
-    world::World,
-};
-use bevy_reflect::FromType;
+use crate::entity::{EntityMapper, 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 EntityMapper),
 }
 
 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 an [`EntityMapper`].
     ///
-    /// 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(
-        &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);
-        });
+    /// # Panics
+    /// Panics if the the type of the reflected value doesn't match.
+    pub fn map_entities(&self, reflected: &mut dyn PartialReflect, mapper: &mut dyn EntityMapper) {
+        (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("reflected type should match");
+                concrete.map_entities(&mut mapper);
+                reflected.apply(&concrete);
             },
         }
     }

crates/bevy_ecs/src/reflect/mod.rs

@@ -24,7 +24,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};
 pub use visit_entities::{ReflectVisitEntities, ReflectVisitEntitiesMut};