Make the MapEntities trait generic over Mappers, and add a simpler EntityMapper

crates/bevy_ecs/src/entity/map_entities.rs

@@ -18,7 +18,7 @@ use bevy_utils::EntityHashMap;
 ///
 /// ```
 /// use bevy_ecs::prelude::*;
-/// use bevy_ecs::entity::{EntityMapper, MapEntities};
+/// use bevy_ecs::entity::{EntityMapper, MapEntities, SimpleEntityMapper};
 ///
 /// #[derive(Component)]
 /// struct Spring {
@@ -27,9 +27,13 @@ use bevy_utils::EntityHashMap;
 /// }
 ///
 /// impl MapEntities for Spring {
-///     fn map_entities(&mut self, entity_mapper: &mut EntityMapper) {
-///         self.a = entity_mapper.get_or_reserve(self.a);
-///         self.b = entity_mapper.get_or_reserve(self.b);
+///     fn map_or_gen_entities(&mut self, entity_mapper: &mut EntityMapper) {
+///         self.a.map_or_gen_entities(entity_mapper);
+///         self.b.map_or_gen_entities(entity_mapper);
+///     }
+///     fn map_entities(&mut self, entity_mapper: &mut SimpleEntityMapper) {
+///         self.a.map_entities(entity_mapper);
+///         self.b.map_entities(entity_mapper)
 ///     }
 /// }
 /// ```
@@ -39,7 +43,33 @@ pub trait MapEntities {
     ///
     /// Implementors should look up any and all [`Entity`] values stored within and
     /// update them to the mapped values via `entity_mapper`.
-    fn map_entities(&mut self, entity_mapper: &mut EntityMapper);
+    fn map_or_gen_entities(&mut self, entity_mapper: &mut EntityMapper);
+
+    /// Updates all [`Entity`] references stored inside using `entity_mapper`.
+    ///
+    /// Only updates the references for which there is a mapping.
+    fn map_entities(&mut self, entity_mapper: &mut SimpleEntityMapper);
+}
+
+/// Similar to `EntityMapper`, but does not allocate new [`Entity`] references in case we couldn't map the entity.
+pub struct SimpleEntityMapper<'m> {
+    map: &'m EntityHashMap<Entity, Entity>,
+}
+
+impl<'m> SimpleEntityMapper<'m> {
+    pub fn new(map: &'m EntityHashMap<Entity, Entity>) -> Self {
+        Self { map }
+    }
+
+    /// Returns the corresponding mapped entity or None if it is absent.
+    pub fn get(&mut self, entity: Entity) -> Option<Entity> {
+        self.map.get(&entity).copied()
+    }
+
+    /// Gets a reference to the underlying [`EntityHashMap<Entity, Entity>`].
+    pub fn get_map(&'m self) -> &'m EntityHashMap<Entity, Entity> {
+        self.map
+    }
 }
 
 /// A wrapper for [`EntityHashMap<Entity, Entity>`], augmenting it with the ability to allocate new [`Entity`] references in a destination

crates/bevy_ecs/src/entity/mod.rs

@@ -136,6 +136,18 @@ pub struct Entity {
     index: u32,
 }
 
+impl MapEntities for Entity {
+    fn map_or_gen_entities(&mut self, entity_mapper: &mut EntityMapper) {
+        *self = entity_mapper.get_or_reserve(*self);
+    }
+
+    fn map_entities(&mut self, entity_mapper: &mut SimpleEntityMapper) {
+        if let Some(mapped) = entity_mapper.get(*self) {
+            *self = mapped;
+        }
+    }
+}
+
 // By not short-circuiting in comparisons, we get better codegen.
 // See <https://github.com/rust-lang/rust/issues/117800>
 impl PartialEq for Entity {

crates/bevy_ecs/src/reflect/map_entities.rs

@@ -13,8 +13,8 @@ use bevy_utils::EntityHashMap;
 /// See [`MapEntities`] for more information.
 #[derive(Clone)]
 pub struct ReflectMapEntities {
-    map_all_entities: fn(&mut World, &mut EntityMapper),
-    map_entities: fn(&mut World, &mut EntityMapper, &[Entity]),
+    map_or_gen_all_entities: fn(&mut World, &mut EntityMapper),
+    map_or_gen_entities: fn(&mut World, &mut EntityMapper, &[Entity]),
 }
 
 impl ReflectMapEntities {
@@ -27,12 +27,12 @@ impl ReflectMapEntities {
     /// 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(
+    pub fn map_or_gen_all_entities(
         &self,
         world: &mut World,
         entity_map: &mut EntityHashMap<Entity, Entity>,
     ) {
-        EntityMapper::world_scope(entity_map, world, self.map_all_entities);
+        EntityMapper::world_scope(entity_map, world, self.map_or_gen_all_entities);
     }
 
     /// A general method for applying [`MapEntities`] behavior to elements in an [`EntityHashMap<Entity, Entity>`]. Unlike
@@ -41,37 +41,37 @@ impl ReflectMapEntities {
     ///
     /// 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(
+    pub fn map_or_gen_entities(
         &self,
         world: &mut World,
         entity_map: &mut EntityHashMap<Entity, Entity>,
         entities: &[Entity],
     ) {
         EntityMapper::world_scope(entity_map, world, |world, mapper| {
-            (self.map_entities)(world, mapper, entities);
+            (self.map_or_gen_entities)(world, mapper, entities);
         });
     }
 }
 
 impl<C: Component + MapEntities> FromType<C> for ReflectMapEntities {
     fn from_type() -> Self {
         ReflectMapEntities {
-            map_entities: |world, entity_mapper, entities| {
+            map_or_gen_entities: |world, entity_mapper, entities| {
                 for &entity in entities {
                     if let Some(mut component) = world.get_mut::<C>(entity) {
-                        component.map_entities(entity_mapper);
+                        component.map_or_gen_entities(entity_mapper);
                     }
                 }
             },
-            map_all_entities: |world, entity_mapper| {
+            map_or_gen_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);
+                        component.map_or_gen_entities(entity_mapper);
                     }
                 }
             },

crates/bevy_hierarchy/src/components/children.rs

@@ -9,6 +9,7 @@ use bevy_ecs::{
 use bevy_utils::smallvec::SmallVec;
 use core::slice;
 use std::ops::Deref;
+use bevy_ecs::entity::SimpleEntityMapper;
 
 /// Contains references to the child entities of this entity.
 ///
@@ -29,9 +30,15 @@ use std::ops::Deref;
 pub struct Children(pub(crate) SmallVec<[Entity; 8]>);
 
 impl MapEntities for Children {
-    fn map_entities(&mut self, entity_mapper: &mut EntityMapper) {
+    fn map_or_gen_entities(&mut self, entity_mapper: &mut EntityMapper) {
         for entity in &mut self.0 {
-            *entity = entity_mapper.get_or_reserve(*entity);
+            entity.map_or_gen_entities(entity_mapper);
+        }
+    }
+
+    fn map_entities(&mut self, entity_mapper: &mut SimpleEntityMapper) {
+        for entity in &mut self.0 {
+            entity.map_entities(entity_mapper);
         }
     }
 }

crates/bevy_hierarchy/src/components/parent.rs

@@ -6,6 +6,7 @@ use bevy_ecs::{
     world::{FromWorld, World},
 };
 use std::ops::Deref;
+use bevy_ecs::entity::SimpleEntityMapper;
 
 /// Holds a reference to the parent entity of this entity.
 /// This component should only be present on entities that actually have a parent entity.
@@ -56,8 +57,12 @@ impl FromWorld for Parent {
 }
 
 impl MapEntities for Parent {
-    fn map_entities(&mut self, entity_mapper: &mut EntityMapper) {
-        self.0 = entity_mapper.get_or_reserve(self.0);
+    fn map_or_gen_entities(&mut self, entity_mapper: &mut EntityMapper) {
+        self.0.map_or_gen_entities(entity_mapper);
+    }
+
+    fn map_entities(&mut self, entity_mapper: &mut SimpleEntityMapper) {
+        self.0.map_entities(entity_mapper);
     }
 }
 

crates/bevy_render/src/mesh/mesh/skinning.rs

@@ -8,6 +8,7 @@ use bevy_ecs::{
 use bevy_math::Mat4;
 use bevy_reflect::{Reflect, TypePath};
 use std::ops::Deref;
+use bevy_ecs::entity::SimpleEntityMapper;
 
 #[derive(Component, Debug, Default, Clone, Reflect)]
 #[reflect(Component, MapEntities)]
@@ -17,9 +18,15 @@ pub struct SkinnedMesh {
 }
 
 impl MapEntities for SkinnedMesh {
-    fn map_entities(&mut self, entity_mapper: &mut EntityMapper) {
+    fn map_or_gen_entities(&mut self, entity_mapper: &mut EntityMapper) {
         for joint in &mut self.joints {
-            *joint = entity_mapper.get_or_reserve(*joint);
+            *joint.map_or_gen_entities(entity_mapper);
+        }
+    }
+
+    fn map_entities(&mut self, entity_mapper: &mut SimpleEntityMapper) {
+        for joint in &mut self.joints {
+            *joint.map_entities(entity_mapper);
         }
     }
 }

crates/bevy_scene/src/dynamic_scene.rs

@@ -149,7 +149,7 @@ impl DynamicScene {
                 "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);
+                map_entities_reflect.map_or_gen_entities(world, entity_map, &entities);
             }
         }
 

crates/bevy_scene/src/scene.rs

@@ -130,7 +130,7 @@ impl Scene {
 
         for registration in type_registry.iter() {
             if let Some(map_entities_reflect) = registration.data::<ReflectMapEntities>() {
-                map_entities_reflect.map_all_entities(world, &mut instance_info.entity_map);
+                map_entities_reflect.map_or_gen_all_entities(world, &mut instance_info.entity_map);
             }
         }
 

crates/bevy_scene/src/serde.rs

@@ -499,7 +499,7 @@ mod tests {
     use crate::ron;
     use crate::serde::{SceneDeserializer, SceneSerializer};
     use crate::{DynamicScene, DynamicSceneBuilder};
-    use bevy_ecs::entity::{Entity, EntityMapper, MapEntities};
+    use bevy_ecs::entity::{Entity, EntityMapper, MapEntities, SimpleEntityMapper};
     use bevy_ecs::prelude::{Component, ReflectComponent, ReflectResource, Resource, World};
     use bevy_ecs::query::{With, Without};
     use bevy_ecs::reflect::{AppTypeRegistry, ReflectMapEntities};
@@ -550,8 +550,12 @@ mod tests {
     struct MyEntityRef(Entity);
 
     impl MapEntities for MyEntityRef {
-        fn map_entities(&mut self, entity_mapper: &mut EntityMapper) {
-            self.0 = entity_mapper.get_or_reserve(self.0);
+        fn map_or_gen_entities(&mut self, entity_mapper: &mut EntityMapper) {
+            self.0.map_or_gen_entities(entity_mapper);
+        }
+
+        fn map_entities(&mut self, entity_mapper: &mut SimpleEntityMapper) {
+            self.0.map_entities(entity_mapper);
         }
     }
 

crates/bevy_window/src/window.rs

@@ -2,6 +2,7 @@ use bevy_ecs::{
     entity::{Entity, EntityMapper, MapEntities},
     prelude::{Component, ReflectComponent},
 };
+use bevy_ecs::entity::SimpleEntityMapper;
 use bevy_math::{DVec2, IVec2, Vec2};
 use bevy_reflect::{std_traits::ReflectDefault, Reflect};
 
@@ -59,10 +60,19 @@ impl WindowRef {
 }
 
 impl MapEntities for WindowRef {
-    fn map_entities(&mut self, entity_mapper: &mut EntityMapper) {
+    fn map_or_gen_entities(&mut self, entity_mapper: &mut EntityMapper) {
         match self {
             Self::Entity(entity) => {
-                *entity = entity_mapper.get_or_reserve(*entity);
+                *entity.map_entities(entity_mapper);
+            }
+            Self::Primary => {}
+        };
+    }
+
+    fn map_entities(&mut self, entity_mapper: &mut SimpleEntityMapper) {
+        match self {
+            Self::Entity(entity) => {
+                *entity.map_or_gen_entities(entity_mapper);
             }
             Self::Primary => {}
         };