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

crates/bevy_ecs/src/entity/map_entities.rs

@@ -9,7 +9,8 @@ use bevy_utils::EntityHashMap;
 ///
 /// As entity IDs are valid only for the [`World`] they're sourced from, using [`Entity`]
 /// as references in components copied from another world will be invalid. This trait
-/// allows defining custom mappings for these references via [`EntityHashMap<Entity, Entity>`]
+/// allows defining custom mappings for these references via a [`Mapper`], which is a type that knows
+/// how to perform entity mapping (usually by using an [`EntityHashMap<Entity, Entity>`]).
 ///
 /// Implementing this trait correctly is required for properly loading components
 /// with entity references from scenes.
@@ -18,7 +19,7 @@ use bevy_utils::EntityHashMap;
 ///
 /// ```
 /// use bevy_ecs::prelude::*;
-/// use bevy_ecs::entity::{EntityMapper, MapEntities};
+/// use bevy_ecs::entity::{EntityMapper, MapEntities, Mapper};
 ///
 /// #[derive(Component)]
 /// struct Spring {
@@ -27,19 +28,47 @@ 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_entities<M: Mapper>(&mut self, entity_mapper: &mut M) {
+///         self.a = entity_mapper.map(self.a);
+///         self.b = entity_mapper.map(self.b);
 ///     }
 /// }
 /// ```
 ///
 pub trait MapEntities {
     /// Updates all [`Entity`] references stored inside using `entity_mapper`.
-    ///
-    /// 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_entities<M: Mapper>(&mut self, entity_mapper: &mut M);
+}
+
+/// Any implementor of this trait knows how to map an [`Entity`] into another [`Entity`].
+///
+/// Usually this is done by using a [`EntityHashMap<Entity, Entity>`] to map the [`Entity`] references.
+/// This can be used to map [`Entity`] references from one [`World`] to another.
+pub trait Mapper {
+    /// Map an entity to another entity
+    fn map(&mut self, entity: Entity) -> Entity;
+}
+
+impl Mapper for EntityMapper<'_> {
+    /// Returns the corresponding mapped entity or reserves a new dead entity ID if it is absent.
+    fn map(&mut self, entity: Entity) -> Entity {
+        if let Some(&mapped) = self.map.get(&entity) {
+            return mapped;
+        }
+
+        // this new entity reference is specifically designed to never represent any living entity
+        let new = Entity::from_raw_and_generation(
+            self.dead_start.index(),
+            IdentifierMask::inc_masked_high_by(self.dead_start.generation, self.generations),
+        );
+
+        // Prevent generations counter from being a greater value than HIGH_MASK.
+        self.generations = (self.generations + 1) & HIGH_MASK;
+
+        self.map.insert(entity, new);
+
+        new
+    }
 }
 
 /// A wrapper for [`EntityHashMap<Entity, Entity>`], augmenting it with the ability to allocate new [`Entity`] references in a destination
@@ -65,24 +94,10 @@ pub struct EntityMapper<'m> {
 }
 
 impl<'m> EntityMapper<'m> {
+    #[deprecated(since = "0.13.0", note = "please use `EntityMapper::map` instead")]
     /// Returns the corresponding mapped entity or reserves a new dead entity ID if it is absent.
     pub fn get_or_reserve(&mut self, entity: Entity) -> Entity {
-        if let Some(&mapped) = self.map.get(&entity) {
-            return mapped;
-        }
-
-        // this new entity reference is specifically designed to never represent any living entity
-        let new = Entity::from_raw_and_generation(
-            self.dead_start.index(),
-            IdentifierMask::inc_masked_high_by(self.dead_start.generation, self.generations),
-        );
-
-        // Prevent generations counter from being a greater value than HIGH_MASK.
-        self.generations = (self.generations + 1) & HIGH_MASK;
-
-        self.map.insert(entity, new);
-
-        new
+        self.map(entity)
     }
 
     /// Gets a reference to the underlying [`EntityHashMap<Entity, Entity>`].
@@ -153,15 +168,15 @@ mod tests {
         let mut mapper = EntityMapper::new(&mut map, &mut world);
 
         let mapped_ent = Entity::from_raw(FIRST_IDX);
-        let dead_ref = mapper.get_or_reserve(mapped_ent);
+        let dead_ref = mapper.map(mapped_ent);
 
         assert_eq!(
             dead_ref,
-            mapper.get_or_reserve(mapped_ent),
+            mapper.map(mapped_ent),
             "should persist the allocated mapping from the previous line"
         );
         assert_eq!(
-            mapper.get_or_reserve(Entity::from_raw(SECOND_IDX)).index(),
+            mapper.map(Entity::from_raw(SECOND_IDX)).index(),
             dead_ref.index(),
             "should re-use the same index for further dead refs"
         );
@@ -179,7 +194,7 @@ mod tests {
         let mut world = World::new();
 
         let dead_ref = EntityMapper::world_scope(&mut map, &mut world, |_, mapper| {
-            mapper.get_or_reserve(Entity::from_raw(0))
+            mapper.map(Entity::from_raw(0))
         });
 
         // Next allocated entity should be a further generation on the same index

crates/bevy_hierarchy/src/components/children.rs

@@ -2,7 +2,7 @@
 use bevy_ecs::reflect::{ReflectComponent, ReflectMapEntities};
 use bevy_ecs::{
     component::Component,
-    entity::{Entity, EntityMapper, MapEntities},
+    entity::{Entity, MapEntities, Mapper},
     prelude::FromWorld,
     world::World,
 };
@@ -29,9 +29,9 @@ 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_entities<M: Mapper>(&mut self, entity_mapper: &mut M) {
         for entity in &mut self.0 {
-            *entity = entity_mapper.get_or_reserve(*entity);
+            *entity = entity_mapper.map(*entity);
         }
     }
 }

crates/bevy_hierarchy/src/components/parent.rs

@@ -2,7 +2,7 @@
 use bevy_ecs::reflect::{ReflectComponent, ReflectMapEntities};
 use bevy_ecs::{
     component::Component,
-    entity::{Entity, EntityMapper, MapEntities},
+    entity::{Entity, MapEntities, Mapper},
     world::{FromWorld, World},
 };
 use std::ops::Deref;
@@ -56,8 +56,8 @@ 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_entities<M: Mapper>(&mut self, entity_mapper: &mut M) {
+        self.0 = entity_mapper.map(self.0);
     }
 }
 

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

@@ -1,7 +1,7 @@
 use bevy_asset::{Asset, Handle};
 use bevy_ecs::{
     component::Component,
-    entity::{Entity, EntityMapper, MapEntities},
+    entity::{Entity, MapEntities, Mapper},
     prelude::ReflectComponent,
     reflect::ReflectMapEntities,
 };
@@ -17,9 +17,9 @@ pub struct SkinnedMesh {
 }
 
 impl MapEntities for SkinnedMesh {
-    fn map_entities(&mut self, entity_mapper: &mut EntityMapper) {
+    fn map_entities<M: Mapper>(&mut self, entity_mapper: &mut M) {
         for joint in &mut self.joints {
-            *joint = entity_mapper.get_or_reserve(*joint);
+            *joint = entity_mapper.map(*joint);
         }
     }
 }

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, MapEntities, Mapper};
     use bevy_ecs::prelude::{Component, ReflectComponent, ReflectResource, Resource, World};
     use bevy_ecs::query::{With, Without};
     use bevy_ecs::reflect::{AppTypeRegistry, ReflectMapEntities};
@@ -550,8 +550,8 @@ 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_entities<M: Mapper>(&mut self, entity_mapper: &mut M) {
+            self.0 = entity_mapper.map(self.0);
         }
     }
 

crates/bevy_window/src/window.rs

@@ -1,5 +1,5 @@
 use bevy_ecs::{
-    entity::{Entity, EntityMapper, MapEntities},
+    entity::{Entity, MapEntities, Mapper},
     prelude::{Component, ReflectComponent},
 };
 use bevy_math::{DVec2, IVec2, Vec2};
@@ -59,10 +59,10 @@ impl WindowRef {
 }
 
 impl MapEntities for WindowRef {
-    fn map_entities(&mut self, entity_mapper: &mut EntityMapper) {
+    fn map_entities<M: Mapper>(&mut self, entity_mapper: &mut M) {
         match self {
             Self::Entity(entity) => {
-                *entity = entity_mapper.get_or_reserve(*entity);
+                *entity = entity_mapper.map(*entity);
             }
             Self::Primary => {}
         };