Make the MapEntities trait generic over Mappers, and add a simpler EntityMapper (#11428)

# Objective

My motivation are to resolve some of the issues I describe in this
[PR](#11415):
- not being able to easily mapping entities because the current
EntityMapper requires `&mut World` access
- not being able to create my own `EntityMapper` because some components
(`Parent` or `Children`) do not provide any public way of modifying the
inner entities

This PR makes the `MapEntities` trait accept a generic type that
implements `Mapper` to perform the mapping.
This means we don't need to use `EntityMapper` to perform our mapping,
we can use any type that implements `Mapper`. Basically this change is
very similar to what `serde` does. Instead of specifying directly how to
map entities for a given type, we have 2 distinct steps:
- the user implements `MapEntities` to define how the type will be
traversed and which `Entity`s will be mapped
  - the `Mapper` defines how the mapping is actually done
This is similar to the distinction between `Serialize` (`MapEntities`)
and `Serializer` (`Mapper`).

This allows networking library to map entities without having to use the
existing `EntityMapper` (which requires `&mut World` access and the use
of `world_scope()`)


## Migration Guide
- The existing `EntityMapper` (notably used to replicate `Scenes` across
different `World`s) has been renamed to `SceneEntityMapper`

- The `MapEntities` trait now works with a generic `EntityMapper`
instead of the specific struct `EntityMapper`.
Calls to `fn map_entities(&mut self, entity_mapper: &mut EntityMapper)`
need to be updated to
`fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M)`

- The new trait `EntityMapper` has been added to the prelude

---------

Co-authored-by: Charles Bournhonesque <[email protected]>
Co-authored-by: Alice Cecile <[email protected]>
Co-authored-by: UkoeHB <[email protected]>

Author: 4 people

crates/bevy_ecs/src/entity/map_entities.rs

@@ -9,7 +9,10 @@ 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 [`EntityMappers`](EntityMapper), which
+/// inject the entity mapping strategy between your `MapEntities` type and the current world
+/// (usually by using an [`EntityHashMap<Entity, Entity>`] between source entities and entities in the
+/// current world).
 ///
 /// Implementing this trait correctly is required for properly loading components
 /// with entity references from scenes.
@@ -18,7 +21,7 @@ use bevy_utils::EntityHashMap;
 ///
 /// ```
 /// use bevy_ecs::prelude::*;
-/// use bevy_ecs::entity::{EntityMapper, MapEntities};
+/// use bevy_ecs::entity::MapEntities;
 ///
 /// #[derive(Component)]
 /// struct Spring {
@@ -27,19 +30,52 @@ 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: EntityMapper>(&mut self, entity_mapper: &mut M) {
+///         self.a = entity_mapper.map_entity(self.a);
+///         self.b = entity_mapper.map_entity(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
+    /// Implementors should look up any and all [`Entity`] values stored within `self` and
     /// update them to the mapped values via `entity_mapper`.
-    fn map_entities(&mut self, entity_mapper: &mut EntityMapper);
+    fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M);
+}
+
+/// An implementor of this trait knows how to map an [`Entity`] into another [`Entity`].
+///
+/// Usually this is done by using an [`EntityHashMap<Entity, Entity>`] to map source entities
+/// (mapper inputs) to the current world's entities (mapper outputs).
+///
+/// More generally, this can be used to map [`Entity`] references between any two [`Worlds`](World).
+pub trait EntityMapper {
+    /// Map an entity to another entity
+    fn map_entity(&mut self, entity: Entity) -> Entity;
+}
+
+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 {
+        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
@@ -48,41 +84,30 @@ pub trait MapEntities {
 /// References are allocated by returning increasing generations starting from an internally initialized base
 /// [`Entity`]. After it is finished being used by [`MapEntities`] implementations, this entity is despawned and the
 /// requisite number of generations reserved.
-pub struct EntityMapper<'m> {
+pub struct SceneEntityMapper<'m> {
     /// A mapping from one set of entities to another.
     ///
     /// This is typically used to coordinate data transfer between sets of entities, such as between a scene and the world
     /// or over the network. This is required as [`Entity`] identifiers are opaque; you cannot and do not want to reuse
     /// identifiers directly.
     ///
     /// On its own, a [`EntityHashMap<Entity, Entity>`] is not capable of allocating new entity identifiers, which is needed to map references
-    /// to entities that lie outside the source entity set. This functionality can be accessed through [`EntityMapper::world_scope()`].
+    /// to entities that lie outside the source entity set. This functionality can be accessed through [`SceneEntityMapper::world_scope()`].
     map: &'m mut EntityHashMap<Entity, Entity>,
     /// A base [`Entity`] used to allocate new references.
     dead_start: Entity,
     /// The number of generations this mapper has allocated thus far.
     generations: u32,
 }
 
-impl<'m> EntityMapper<'m> {
-    /// Returns the corresponding mapped entity or reserves a new dead entity ID if it is absent.
+impl<'m> SceneEntityMapper<'m> {
+    #[deprecated(
+        since = "0.13.0",
+        note = "please use `EntityMapper::map_entity` instead"
+    )]
+    /// Returns the corresponding mapped entity or reserves a new dead entity ID in the current world 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(entity)
     }
 
     /// Gets a reference to the underlying [`EntityHashMap<Entity, Entity>`].
@@ -95,7 +120,7 @@ impl<'m> EntityMapper<'m> {
         self.map
     }
 
-    /// Creates a new [`EntityMapper`], spawning a temporary base [`Entity`] in the provided [`World`]
+    /// Creates a new [`SceneEntityMapper`], spawning a temporary base [`Entity`] in the provided [`World`]
     fn new(map: &'m mut EntityHashMap<Entity, Entity>, world: &mut World) -> Self {
         Self {
             map,
@@ -107,7 +132,7 @@ impl<'m> EntityMapper<'m> {
 
     /// Reserves the allocated references to dead entities within the world. This frees the temporary base
     /// [`Entity`] while reserving extra generations via [`crate::entity::Entities::reserve_generations`]. Because this
-    /// renders the [`EntityMapper`] unable to safely allocate any more references, this method takes ownership of
+    /// renders the [`SceneEntityMapper`] unable to safely allocate any more references, this method takes ownership of
     /// `self` in order to render it unusable.
     fn finish(self, world: &mut World) {
         // SAFETY: Entities data is kept in a valid state via `EntityMap::world_scope`
@@ -116,7 +141,7 @@ impl<'m> EntityMapper<'m> {
         assert!(entities.reserve_generations(self.dead_start.index(), self.generations));
     }
 
-    /// Creates an [`EntityMapper`] from a provided [`World`] and [`EntityHashMap<Entity, Entity>`], then calls the
+    /// Creates an [`SceneEntityMapper`] from a provided [`World`] and [`EntityHashMap<Entity, Entity>`], then calls the
     /// provided function with it. This allows one to allocate new entity references in this [`World`] that are
     /// guaranteed to never point at a living entity now or in the future. This functionality is useful for safely
     /// mapping entity identifiers that point at entities outside the source world. The passed function, `f`, is called
@@ -139,7 +164,7 @@ mod tests {
     use bevy_utils::EntityHashMap;
 
     use crate::{
-        entity::{Entity, EntityMapper},
+        entity::{Entity, EntityMapper, SceneEntityMapper},
         world::World,
     };
 
@@ -150,18 +175,18 @@ mod tests {
 
         let mut map = EntityHashMap::default();
         let mut world = World::new();
-        let mut mapper = EntityMapper::new(&mut map, &mut world);
+        let mut mapper = SceneEntityMapper::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_entity(mapped_ent);
 
         assert_eq!(
             dead_ref,
-            mapper.get_or_reserve(mapped_ent),
+            mapper.map_entity(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(Entity::from_raw(SECOND_IDX)).index(),
             dead_ref.index(),
             "should re-use the same index for further dead refs"
         );
@@ -178,8 +203,8 @@ mod tests {
         let mut map = EntityHashMap::default();
         let mut world = World::new();
 
-        let dead_ref = EntityMapper::world_scope(&mut map, &mut world, |_, mapper| {
-            mapper.get_or_reserve(Entity::from_raw(0))
+        let dead_ref = SceneEntityMapper::world_scope(&mut map, &mut world, |_, mapper| {
+            mapper.map_entity(Entity::from_raw(0))
         });
 
         // Next allocated entity should be a further generation on the same index

crates/bevy_ecs/src/lib.rs

@@ -36,7 +36,7 @@ pub mod prelude {
         bundle::Bundle,
         change_detection::{DetectChanges, DetectChangesMut, Mut, Ref},
         component::Component,
-        entity::Entity,
+        entity::{Entity, EntityMapper},
         event::{Event, EventReader, EventWriter, Events},
         query::{Added, AnyOf, Changed, Has, Or, QueryBuilder, QueryState, With, Without},
         removal_detection::RemovedComponents,

crates/bevy_ecs/src/reflect/map_entities.rs

@@ -1,6 +1,6 @@
 use crate::{
     component::Component,
-    entity::{Entity, EntityMapper, MapEntities},
+    entity::{Entity, MapEntities, SceneEntityMapper},
     world::World,
 };
 use bevy_reflect::FromType;
@@ -10,11 +10,11 @@ use bevy_utils::EntityHashMap;
 /// 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 [`MapEntities`] for more information.
+/// See [`SceneEntityMapper`] and [`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_all_entities: fn(&mut World, &mut SceneEntityMapper),
+    map_entities: fn(&mut World, &mut SceneEntityMapper, &[Entity]),
 }
 
 impl ReflectMapEntities {
@@ -32,7 +32,7 @@ impl ReflectMapEntities {
         world: &mut World,
         entity_map: &mut EntityHashMap<Entity, Entity>,
     ) {
-        EntityMapper::world_scope(entity_map, world, self.map_all_entities);
+        SceneEntityMapper::world_scope(entity_map, world, self.map_all_entities);
     }
 
     /// A general method for applying [`MapEntities`] behavior to elements in an [`EntityHashMap<Entity, Entity>`]. Unlike
@@ -47,7 +47,7 @@ impl ReflectMapEntities {
         entity_map: &mut EntityHashMap<Entity, Entity>,
         entities: &[Entity],
     ) {
-        EntityMapper::world_scope(entity_map, world, |world, mapper| {
+        SceneEntityMapper::world_scope(entity_map, world, |world, mapper| {
             (self.map_entities)(world, mapper, entities);
         });
     }

crates/bevy_hierarchy/src/components/children.rs

@@ -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: EntityMapper>(&mut self, entity_mapper: &mut M) {
         for entity in &mut self.0 {
-            *entity = entity_mapper.get_or_reserve(*entity);
+            *entity = entity_mapper.map_entity(*entity);
         }
     }
 }

crates/bevy_hierarchy/src/components/parent.rs

@@ -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: EntityMapper>(&mut self, entity_mapper: &mut M) {
+        self.0 = entity_mapper.map_entity(self.0);
     }
 }
 

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

@@ -17,9 +17,9 @@ pub struct SkinnedMesh {
 }
 
 impl MapEntities for SkinnedMesh {
-    fn map_entities(&mut self, entity_mapper: &mut EntityMapper) {
+    fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M) {
         for joint in &mut self.joints {
-            *joint = entity_mapper.get_or_reserve(*joint);
+            *joint = entity_mapper.map_entity(*joint);
         }
     }
 }

crates/bevy_scene/src/serde.rs

@@ -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: EntityMapper>(&mut self, entity_mapper: &mut M) {
+            self.0 = entity_mapper.map_entity(self.0);
         }
     }
 

crates/bevy_window/src/window.rs

@@ -59,10 +59,10 @@ impl WindowRef {
 }
 
 impl MapEntities for WindowRef {
-    fn map_entities(&mut self, entity_mapper: &mut EntityMapper) {
+    fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M) {
         match self {
             Self::Entity(entity) => {
-                *entity = entity_mapper.get_or_reserve(*entity);
+                *entity = entity_mapper.map_entity(*entity);
             }
             Self::Primary => {}
         };