Replaced EntityMap with HashMap (#9461)

# Objective

- Fixes #9321

## Solution

- `EntityMap` has been replaced by a simple `HashMap<Entity, Entity>`.

---

## Changelog

- `EntityMap::world_scope` has been replaced with `World::world_scope`
to avoid creating a new trait. This is a public facing change to the
call semantics, but has no effect on results or behaviour.
- `EntityMap`, as a `HashMap`, now operates on `&Entity` rather than
`Entity`. This changes many standard access functions (e.g, `.get`) in a
public-facing way.

## Migration Guide

- Calls to `EntityMap::world_scope` can be directly replaced with the
following:
  `map.world_scope(&mut world)` -> `world.world_scope(&mut map)`
- Calls to legacy `EntityMap` methods such as `EntityMap::get` must
explicitly include de/reference symbols:
  `let entity = map.get(parent);` -> `let &entity = map.get(&parent);`

Author: bushrat011899

crates/bevy_ecs/src/entity/map_entities.rs

@@ -1,11 +1,11 @@
 use crate::{entity::Entity, world::World};
-use bevy_utils::{Entry, HashMap};
+use bevy_utils::HashMap;
 
 /// Operation to map all contained [`Entity`] fields in a type to new values.
 ///
 /// 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 [`EntityMap`].
+/// allows defining custom mappings for these references via [`HashMap<Entity, Entity>`]
 ///
 /// Implementing this trait correctly is required for properly loading components
 /// with entity references from scenes.
@@ -32,112 +32,29 @@ use bevy_utils::{Entry, HashMap};
 ///
 /// [`World`]: crate::world::World
 pub trait MapEntities {
-    /// Updates all [`Entity`] references stored inside using `entity_map`.
+    /// 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);
 }
 
-/// A mapping from one set of entities to another.
-///
-/// The API generally follows [`HashMap`], but each [`Entity`] is returned by value, as they are [`Copy`].
-///
-/// 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, an `EntityMap` is not capable of allocating new entity identifiers, which is needed to map references
-/// to entities that lie outside the source entity set. To do this, an `EntityMap` can be wrapped in an
-/// [`EntityMapper`] which scopes it to a particular destination [`World`] and allows new identifiers to be allocated.
-/// This functionality can be accessed through [`Self::world_scope()`].
-#[derive(Default, Debug)]
-pub struct EntityMap {
-    map: HashMap<Entity, Entity>,
-}
-
-impl EntityMap {
-    /// Inserts an entities pair into the map.
-    ///
-    /// If the map did not have `from` present, [`None`] is returned.
-    ///
-    /// If the map did have `from` present, the value is updated, and the old value is returned.
-    pub fn insert(&mut self, from: Entity, to: Entity) -> Option<Entity> {
-        self.map.insert(from, to)
-    }
-
-    /// Removes an `entity` from the map, returning the mapped value of it if the `entity` was previously in the map.
-    pub fn remove(&mut self, entity: Entity) -> Option<Entity> {
-        self.map.remove(&entity)
-    }
-
-    /// Gets the given entity's corresponding entry in the map for in-place manipulation.
-    pub fn entry(&mut self, entity: Entity) -> Entry<'_, Entity, Entity> {
-        self.map.entry(entity)
-    }
-
-    /// Returns the corresponding mapped entity.
-    pub fn get(&self, entity: Entity) -> Option<Entity> {
-        self.map.get(&entity).copied()
-    }
-
-    /// An iterator visiting all keys in arbitrary order.
-    pub fn keys(&self) -> impl Iterator<Item = Entity> + '_ {
-        self.map.keys().cloned()
-    }
-
-    /// An iterator visiting all values in arbitrary order.
-    pub fn values(&self) -> impl Iterator<Item = Entity> + '_ {
-        self.map.values().cloned()
-    }
-
-    /// Returns the number of elements in the map.
-    pub fn len(&self) -> usize {
-        self.map.len()
-    }
-
-    /// Returns true if the map contains no elements.
-    pub fn is_empty(&self) -> bool {
-        self.map.is_empty()
-    }
-
-    /// An iterator visiting all (key, value) pairs in arbitrary order.
-    pub fn iter(&self) -> impl Iterator<Item = (Entity, Entity)> + '_ {
-        self.map.iter().map(|(from, to)| (*from, *to))
-    }
-
-    /// Clears the map, removing all entity pairs. Keeps the allocated memory for reuse.
-    pub fn clear(&mut self) {
-        self.map.clear();
-    }
-
-    /// Creates an [`EntityMapper`] from this [`EntityMap`] and scoped to the provided [`World`], then calls the
-    /// provided function with it. This allows one to allocate new entity references in the provided `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
-    /// within the scope of the passed world. Its return value is then returned from `world_scope` as the generic type
-    /// parameter `R`.
-    pub fn world_scope<R>(
-        &mut self,
-        world: &mut World,
-        f: impl FnOnce(&mut World, &mut EntityMapper) -> R,
-    ) -> R {
-        let mut mapper = EntityMapper::new(self, world);
-        let result = f(world, &mut mapper);
-        mapper.finish(world);
-        result
-    }
-}
-
-/// A wrapper for [`EntityMap`], augmenting it with the ability to allocate new [`Entity`] references in a destination
+/// A wrapper for [`HashMap<Entity, Entity>`], augmenting it with the ability to allocate new [`Entity`] references in a destination
 /// world. These newly allocated references are guaranteed to never point to any living entity in that world.
 ///
 /// 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> {
-    /// The wrapped [`EntityMap`].
-    map: &'m mut EntityMap,
+    /// 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 [`HashMap<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()`].
+    map: &'m mut HashMap<Entity, Entity>,
     /// A base [`Entity`] used to allocate new references.
     dead_start: Entity,
     /// The number of generations this mapper has allocated thus far.
@@ -147,7 +64,7 @@ pub struct EntityMapper<'m> {
 impl<'m> EntityMapper<'m> {
     /// 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) {
+        if let Some(&mapped) = self.map.get(&entity) {
             return mapped;
         }
 
@@ -163,21 +80,21 @@ impl<'m> EntityMapper<'m> {
         new
     }
 
-    /// Gets a reference to the underlying [`EntityMap`].
-    pub fn get_map(&'m self) -> &'m EntityMap {
+    /// Gets a reference to the underlying [`HashMap<Entity, Entity>`].
+    pub fn get_map(&'m self) -> &'m HashMap<Entity, Entity> {
         self.map
     }
 
-    /// Gets a mutable reference to the underlying [`EntityMap`]
-    pub fn get_map_mut(&'m mut self) -> &'m mut EntityMap {
+    /// Gets a mutable reference to the underlying [`HashMap<Entity, Entity>`].
+    pub fn get_map_mut(&'m mut self) -> &'m mut HashMap<Entity, Entity> {
         self.map
     }
 
     /// Creates a new [`EntityMapper`], spawning a temporary base [`Entity`] in the provided [`World`]
-    fn new(map: &'m mut EntityMap, world: &mut World) -> Self {
+    fn new(map: &'m mut HashMap<Entity, Entity>, world: &mut World) -> Self {
         Self {
             map,
-            // SAFETY: Entities data is kept in a valid state via `EntityMap::world_scope`
+            // SAFETY: Entities data is kept in a valid state via `EntityMapper::world_scope`
             dead_start: unsafe { world.entities_mut().alloc() },
             generations: 0,
         }
@@ -193,19 +110,40 @@ impl<'m> EntityMapper<'m> {
         assert!(entities.free(self.dead_start).is_some());
         assert!(entities.reserve_generations(self.dead_start.index, self.generations));
     }
+
+    /// Creates an [`EntityMapper`] from a provided [`World`] and [`HashMap<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
+    /// within the scope of this world. Its return value is then returned from `world_scope` as the generic type
+    /// parameter `R`.
+    pub fn world_scope<R>(
+        entity_map: &'m mut HashMap<Entity, Entity>,
+        world: &mut World,
+        f: impl FnOnce(&mut World, &mut Self) -> R,
+    ) -> R {
+        let mut mapper = Self::new(entity_map, world);
+        let result = f(world, &mut mapper);
+        mapper.finish(world);
+        result
+    }
 }
 
 #[cfg(test)]
 mod tests {
-    use super::{EntityMap, EntityMapper};
-    use crate::{entity::Entity, world::World};
+    use bevy_utils::HashMap;
+
+    use crate::{
+        entity::{Entity, EntityMapper},
+        world::World,
+    };
 
     #[test]
     fn entity_mapper() {
         const FIRST_IDX: u32 = 1;
         const SECOND_IDX: u32 = 2;
 
-        let mut map = EntityMap::default();
+        let mut map = HashMap::default();
         let mut world = World::new();
         let mut mapper = EntityMapper::new(&mut map, &mut world);
 
@@ -232,10 +170,10 @@ mod tests {
 
     #[test]
     fn world_scope_reserves_generations() {
-        let mut map = EntityMap::default();
+        let mut map = HashMap::default();
         let mut world = World::new();
 
-        let dead_ref = map.world_scope(&mut world, |_, mapper| {
+        let dead_ref = EntityMapper::world_scope(&mut map, &mut world, |_, mapper| {
             mapper.get_or_reserve(Entity::new(0, 0))
         });
 

crates/bevy_ecs/src/reflect/map_entities.rs

@@ -1,9 +1,10 @@
 use crate::{
     component::Component,
-    entity::{Entity, EntityMap, EntityMapper, MapEntities},
+    entity::{Entity, EntityMapper, MapEntities},
     world::World,
 };
 use bevy_reflect::FromType;
+use bevy_utils::HashMap;
 
 /// For a specific type of component, 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
@@ -17,27 +18,32 @@ pub struct ReflectMapEntities {
 }
 
 impl ReflectMapEntities {
-    /// A general method for applying [`MapEntities`] behavior to all elements in an [`EntityMap`].
+    /// A general method for applying [`MapEntities`] behavior to all elements in an [`HashMap<Entity, Entity>`].
     ///
-    /// Be mindful in its usage: Works best in situations where the entities in the [`EntityMap`] are newly
+    /// Be mindful in its usage: Works best in situations where the entities in the [`HashMap<Entity, Entity>`] are newly
     /// created, before systems have a chance to add new components. If some of the entities referred to
-    /// by the [`EntityMap`] might already contain valid entity references, you should use [`map_entities`](Self::map_entities).
+    /// by the [`HashMap<Entity, 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 EntityMap) {
-        entity_map.world_scope(world, self.map_all_entities);
+    pub fn map_all_entities(&self, world: &mut World, entity_map: &mut HashMap<Entity, Entity>) {
+        EntityMapper::world_scope(entity_map, world, self.map_all_entities);
     }
 
-    /// A general method for applying [`MapEntities`] behavior to elements in an [`EntityMap`]. Unlike
+    /// A general method for applying [`MapEntities`] behavior to elements in an [`HashMap<Entity, Entity>`]. Unlike
     /// [`map_all_entities`](Self::map_all_entities), this is applied to specific entities, not all values
-    /// in the [`EntityMap`].
+    /// in the [`HashMap<Entity, Entity>`].
     ///
     /// 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 EntityMap, entities: &[Entity]) {
-        entity_map.world_scope(world, |world, mapper| {
+    pub fn map_entities(
+        &self,
+        world: &mut World,
+        entity_map: &mut HashMap<Entity, Entity>,
+        entities: &[Entity],
+    ) {
+        EntityMapper::world_scope(entity_map, world, |world, mapper| {
             (self.map_entities)(world, mapper, entities);
         });
     }
@@ -54,7 +60,11 @@ impl<C: Component + MapEntities> FromType<C> for ReflectMapEntities {
                 }
             },
             map_all_entities: |world, entity_mapper| {
-                let entities = entity_mapper.get_map().values().collect::<Vec<Entity>>();
+                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);

crates/bevy_scene/src/dynamic_scene.rs

@@ -3,7 +3,7 @@ use std::any::TypeId;
 use crate::{DynamicSceneBuilder, Scene, SceneSpawnError};
 use anyhow::Result;
 use bevy_ecs::{
-    entity::{Entity, EntityMap},
+    entity::Entity,
     reflect::{AppTypeRegistry, ReflectComponent, ReflectMapEntities},
     world::World,
 };
@@ -67,7 +67,7 @@ impl DynamicScene {
     pub fn write_to_world_with(
         &self,
         world: &mut World,
-        entity_map: &mut EntityMap,
+        entity_map: &mut HashMap<Entity, Entity>,
         type_registry: &AppTypeRegistry,
     ) -> Result<(), SceneSpawnError> {
         let type_registry = type_registry.read();
@@ -155,7 +155,7 @@ impl DynamicScene {
     pub fn write_to_world(
         &self,
         world: &mut World,
-        entity_map: &mut EntityMap,
+        entity_map: &mut HashMap<Entity, Entity>,
     ) -> Result<(), SceneSpawnError> {
         let registry = world.resource::<AppTypeRegistry>().clone();
         self.write_to_world_with(world, entity_map, &registry)
@@ -183,8 +183,9 @@ where
 
 #[cfg(test)]
 mod tests {
-    use bevy_ecs::{entity::EntityMap, reflect::AppTypeRegistry, system::Command, world::World};
+    use bevy_ecs::{reflect::AppTypeRegistry, system::Command, world::World};
     use bevy_hierarchy::{AddChild, Parent};
+    use bevy_utils::HashMap;
 
     use crate::dynamic_scene_builder::DynamicSceneBuilder;
 
@@ -213,11 +214,11 @@ mod tests {
         scene_builder.extract_entity(original_parent_entity);
         scene_builder.extract_entity(original_child_entity);
         let scene = scene_builder.build();
-        let mut entity_map = EntityMap::default();
+        let mut entity_map = HashMap::default();
         scene.write_to_world(&mut world, &mut entity_map).unwrap();
 
-        let from_scene_parent_entity = entity_map.get(original_parent_entity).unwrap();
-        let from_scene_child_entity = entity_map.get(original_child_entity).unwrap();
+        let &from_scene_parent_entity = entity_map.get(&original_parent_entity).unwrap();
+        let &from_scene_child_entity = entity_map.get(&original_child_entity).unwrap();
 
         // We then add the parent from the scene as a child of the original child
         // Hierarchy should look like:

crates/bevy_scene/src/scene.rs

@@ -1,9 +1,9 @@
 use bevy_ecs::{
-    entity::EntityMap,
     reflect::{AppTypeRegistry, ReflectComponent, ReflectMapEntities, ReflectResource},
     world::World,
 };
 use bevy_reflect::{TypePath, TypeUuid};
+use bevy_utils::HashMap;
 
 use crate::{DynamicScene, InstanceInfo, SceneSpawnError};
 
@@ -30,7 +30,7 @@ impl Scene {
         type_registry: &AppTypeRegistry,
     ) -> Result<Scene, SceneSpawnError> {
         let mut world = World::new();
-        let mut entity_map = EntityMap::default();
+        let mut entity_map = HashMap::default();
         dynamic_scene.write_to_world_with(&mut world, &mut entity_map, type_registry)?;
 
         Ok(Self { world })
@@ -56,7 +56,7 @@ impl Scene {
         type_registry: &AppTypeRegistry,
     ) -> Result<InstanceInfo, SceneSpawnError> {
         let mut instance_info = InstanceInfo {
-            entity_map: EntityMap::default(),
+            entity_map: HashMap::default(),
         };
 
         let type_registry = type_registry.read();