Make scene handling of entity references robust

crates/bevy_ecs/src/entity/map_entities.rs

@@ -1,4 +1,4 @@
-use crate::entity::Entity;
+use crate::{entity::Entity, world::World};
 use bevy_utils::{Entry, HashMap};
 use std::fmt;
 
@@ -33,7 +33,7 @@ impl fmt::Display for MapEntitiesError {
 ///
 /// ```rust
 /// use bevy_ecs::prelude::*;
-/// use bevy_ecs::entity::{EntityMap, MapEntities, MapEntitiesError};
+/// use bevy_ecs::entity::{EntityMapper, MapEntities, MapEntitiesError};
 ///
 /// #[derive(Component)]
 /// struct Spring {
@@ -42,7 +42,7 @@ impl fmt::Display for MapEntitiesError {
 /// }
 ///
 /// impl MapEntities for Spring {
-///     fn map_entities(&mut self, entity_map: &EntityMap) -> Result<(), MapEntitiesError> {
+///     fn map_entities(&mut self, entity_map: &mut EntityMapper) -> Result<(), MapEntitiesError> {
 ///         self.a = entity_map.get(self.a)?;
 ///         self.b = entity_map.get(self.b)?;
 ///         Ok(())
@@ -56,7 +56,7 @@ pub trait MapEntities {
     ///
     /// Implementors should look up any and all [`Entity`] values stored within and
     /// update them to the mapped values via `entity_map`.
-    fn map_entities(&mut self, entity_map: &EntityMap) -> Result<(), MapEntitiesError>;
+    fn map_entities(&mut self, entity_mapper: &mut EntityMapper) -> Result<(), MapEntitiesError>;
 }
 
 /// A mapping from one set of entities to another.
@@ -70,6 +70,77 @@ pub struct EntityMap {
     map: HashMap<Entity, Entity>,
 }
 
+/// A wrapper for [`EntityMap`], 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 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.
+    pub fn get(&self, entity: Entity) -> Result<Entity, MapEntitiesError> {
+        self.map.get(entity)
+    }
+
+    /// Returns the corresponding mapped entity or allocates a new dead entity if it is absent.
+    pub fn get_or_alloc(&mut self, entity: Entity) -> Entity {
+        if let Ok(mapped) = self.map.get(entity) {
+            return mapped;
+        }
+
+        let new = Entity {
+            generation: self.dead_start.generation + self.generations,
+            index: self.dead_start.index,
+        };
+        self.generations += 1;
+
+        self.map.insert(entity, new);
+
+        new
+    }
+
+    /// Gets a reference to the underlying [`EntityMap`].
+    pub fn get_map(&'m self) -> &'m EntityMap {
+        self.map
+    }
+
+    /// Gets a mutable reference to the underlying [`EntityMap`]
+    pub fn get_map_mut(&'m mut self) -> &'m mut EntityMap {
+        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 {
+        Self {
+            map,
+            dead_start: world.spawn_empty().id(),
+            generations: 0,
+        }
+    }
+
+    /// Reserves the allocated references to dead entities within the world. This despawns the temporary base
+    /// [`Entity`] while reserving extra generations via [`World::reserve_generations`]. Because this renders the
+    /// [`EntityMapper`] unable to safely allocate any more references, this method takes ownership of `self` in order
+    /// to render it unusable.
+    fn save(self, world: &mut World) {
+        if self.generations == 0 {
+            assert!(world.despawn(self.dead_start));
+            return;
+        }
+
+        assert!(world.reserve_generations(self.dead_start, self.generations));
+    }
+}
+
 impl EntityMap {
     /// Inserts an entities pair into the map.
     ///
@@ -122,4 +193,17 @@ impl EntityMap {
     pub fn iter(&self) -> impl Iterator<Item = (Entity, Entity)> + '_ {
         self.map.iter().map(|(from, to)| (*from, *to))
     }
+
+    /// Calls the provided closure with an [`EntityMapper`] created from this [`EntityMap`]. This allows the closure
+    /// to allocate new entity references in the provided [`World`] that will never point at a living entity.
+    pub fn with_mapper<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.save(world);
+        result
+    }
 }

crates/bevy_ecs/src/entity/mod.rs

@@ -590,6 +590,22 @@ impl Entities {
         self.meta.get_unchecked_mut(index as usize).location = location;
     }
 
+    /// Increments the generation of a despawned [`Entity`]. This may only be called on an entity
+    /// index corresponding to an entity that once existed and has since been freed.
+    pub(crate) fn reserve_generations(&mut self, index: u32, generations: u32) -> bool {
+        if (index as usize) >= self.meta.len() {
+            return false;
+        }
+
+        let meta = &mut self.meta[index as usize];
+        if meta.location.archetype_id == ArchetypeId::INVALID {
+            meta.generation += generations;
+            true
+        } else {
+            false
+        }
+    }
+
     /// Get the [`Entity`] with a given id, if it exists in this [`Entities`] collection
     /// Returns `None` if this [`Entity`] is outside of the range of currently reserved Entities
     ///

crates/bevy_ecs/src/reflect.rs

@@ -3,7 +3,7 @@
 use crate::{
     change_detection::Mut,
     component::Component,
-    entity::{Entity, EntityMap, MapEntities, MapEntitiesError},
+    entity::{Entity, EntityMap, EntityMapper, MapEntities, MapEntitiesError},
     system::Resource,
     world::{FromWorld, World},
 };
@@ -406,26 +406,27 @@ impl_from_reflect_value!(Entity);
 
 #[derive(Clone)]
 pub struct ReflectMapEntities {
-    map_entities: fn(&mut World, &EntityMap) -> Result<(), MapEntitiesError>,
+    map_entities: fn(&mut World, &mut EntityMapper) -> Result<(), MapEntitiesError>,
 }
 
 impl ReflectMapEntities {
     pub fn map_entities(
         &self,
         world: &mut World,
-        entity_map: &EntityMap,
+        entity_map: &mut EntityMap,
     ) -> Result<(), MapEntitiesError> {
-        (self.map_entities)(world, entity_map)
+        entity_map.with_mapper(world, self.map_entities)
     }
 }
 
 impl<C: Component + MapEntities> FromType<C> for ReflectMapEntities {
     fn from_type() -> Self {
         ReflectMapEntities {
-            map_entities: |world, entity_map| {
-                for entity in entity_map.values() {
+            map_entities: |world, entity_mapper| {
+                let entities = entity_mapper.get_map().values().collect::<Vec<Entity>>();
+                for entity in entities.iter().cloned() {
                     if let Some(mut component) = world.get_mut::<C>(entity) {
-                        component.map_entities(entity_map)?;
+                        component.map_entities(entity_mapper)?;
                     }
                 }
                 Ok(())

crates/bevy_ecs/src/world/mod.rs

@@ -608,6 +608,22 @@ impl World {
         }
     }
 
+    /// Despawns the given `entity`, if it exists, and reserves a number of subsequent generations.
+    ///
+    /// This function serves an extremely narrow use case of allocating a series of entity IDs that are
+    /// guaranteed to never refer to a live entity. This functionality is useful primarily for mapping references
+    /// to dead entities into a new world alongside a [`crate::entity::MapEntities`] implementation.
+    ///
+    /// See [`Self::despawn()`] for usage.
+    pub(crate) fn reserve_generations(&mut self, entity: Entity, generations: u32) -> bool {
+        if self.despawn(entity) {
+            self.entities
+                .reserve_generations(entity.index(), generations)
+        } else {
+            false
+        }
+    }
+
     /// Clears the internal component tracker state.
     ///
     /// The world maintains some internal state about changed and removed components. This state

crates/bevy_hierarchy/src/components/children.rs

@@ -1,6 +1,6 @@
 use bevy_ecs::{
     component::Component,
-    entity::{Entity, EntityMap, MapEntities, MapEntitiesError},
+    entity::{Entity, EntityMapper, MapEntities, MapEntitiesError},
     prelude::FromWorld,
     reflect::{ReflectComponent, ReflectMapEntities},
     world::World,
@@ -21,9 +21,9 @@ use std::ops::Deref;
 pub struct Children(pub(crate) SmallVec<[Entity; 8]>);
 
 impl MapEntities for Children {
-    fn map_entities(&mut self, entity_map: &EntityMap) -> Result<(), MapEntitiesError> {
+    fn map_entities(&mut self, entity_map: &mut EntityMapper) -> Result<(), MapEntitiesError> {
         for entity in &mut self.0 {
-            *entity = entity_map.get(*entity)?;
+            *entity = entity_map.get_or_alloc(*entity);
         }
 
         Ok(())

crates/bevy_hierarchy/src/components/parent.rs

@@ -1,6 +1,6 @@
 use bevy_ecs::{
     component::Component,
-    entity::{Entity, EntityMap, MapEntities, MapEntitiesError},
+    entity::{Entity, EntityMapper, MapEntities, MapEntitiesError},
     reflect::{ReflectComponent, ReflectMapEntities},
     world::{FromWorld, World},
 };
@@ -36,12 +36,8 @@ impl FromWorld for Parent {
 }
 
 impl MapEntities for Parent {
-    fn map_entities(&mut self, entity_map: &EntityMap) -> Result<(), MapEntitiesError> {
-        // Parent of an entity in the new world can be in outside world, in which case it
-        // should not be mapped.
-        if let Ok(mapped_entity) = entity_map.get(self.0) {
-            self.0 = mapped_entity;
-        }
+    fn map_entities(&mut self, entity_map: &mut EntityMapper) -> Result<(), MapEntitiesError> {
+        self.0 = entity_map.get_or_alloc(self.0);
         Ok(())
     }
 }

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

@@ -1,7 +1,7 @@
 use bevy_asset::Handle;
 use bevy_ecs::{
     component::Component,
-    entity::{Entity, EntityMap, MapEntities, MapEntitiesError},
+    entity::{Entity, EntityMapper, MapEntities, MapEntitiesError},
     prelude::ReflectComponent,
     reflect::ReflectMapEntities,
 };
@@ -17,9 +17,9 @@ pub struct SkinnedMesh {
 }
 
 impl MapEntities for SkinnedMesh {
-    fn map_entities(&mut self, entity_map: &EntityMap) -> Result<(), MapEntitiesError> {
+    fn map_entities(&mut self, entity_map: &mut EntityMapper) -> Result<(), MapEntitiesError> {
         for joint in &mut self.joints {
-            *joint = entity_map.get(*joint)?;
+            *joint = entity_map.get_or_alloc(*joint);
         }
 
         Ok(())

crates/bevy_scene/src/dynamic_scene.rs

@@ -28,8 +28,9 @@ pub struct DynamicScene {
 
 /// A reflection-powered serializable representation of an entity and its components.
 pub struct DynamicEntity {
-    /// The transiently unique identifier of a corresponding `Entity`.
-    pub entity: u32,
+    /// The identifier of the entity, unique within a scene (and the world it may have been generated from).
+    /// Components containing entitiy references will identify their referant via this identifier.
+    pub entity: u64,
     /// A vector of boxed components that belong to the given entity and
     /// implement the `Reflect` trait.
     pub components: Vec<Box<dyn Reflect>>,
@@ -69,7 +70,7 @@ impl DynamicScene {
             // or spawn a new entity with a transiently unique id if there is
             // no corresponding entry.
             let entity = *entity_map
-                .entry(bevy_ecs::entity::Entity::from_raw(scene_entity.entity))
+                .entry(bevy_ecs::entity::Entity::from_bits(scene_entity.entity))
                 .or_insert_with(|| world.spawn_empty().id());
 
             // Apply/ add each component to the given entity.

crates/bevy_scene/src/dynamic_scene_builder.rs

@@ -31,7 +31,7 @@ use std::collections::BTreeMap;
 /// let dynamic_scene = builder.build();
 /// ```
 pub struct DynamicSceneBuilder<'w> {
-    extracted_scene: BTreeMap<u32, DynamicEntity>,
+    extracted_scene: BTreeMap<u64, DynamicEntity>,
     type_registry: AppTypeRegistry,
     original_world: &'w World,
 }
@@ -113,14 +113,14 @@ impl<'w> DynamicSceneBuilder<'w> {
         let type_registry = self.type_registry.read();
 
         for entity in entities {
-            let index = entity.index();
+            let index = entity.to_bits();
 
             if self.extracted_scene.contains_key(&index) {
                 continue;
             }
 
             let mut entry = DynamicEntity {
-                entity: index,
+                entity: entity.to_bits(),
                 components: Vec::new(),
             };
 
@@ -180,7 +180,7 @@ mod tests {
         let scene = builder.build();
 
         assert_eq!(scene.entities.len(), 1);
-        assert_eq!(scene.entities[0].entity, entity.index());
+        assert_eq!(scene.entities[0].entity, entity.to_bits());
         assert_eq!(scene.entities[0].components.len(), 1);
         assert!(scene.entities[0].components[0].represents::<ComponentA>());
     }
@@ -201,7 +201,7 @@ mod tests {
         let scene = builder.build();
 
         assert_eq!(scene.entities.len(), 1);
-        assert_eq!(scene.entities[0].entity, entity.index());
+        assert_eq!(scene.entities[0].entity, entity.to_bits());
         assert_eq!(scene.entities[0].components.len(), 1);
         assert!(scene.entities[0].components[0].represents::<ComponentA>());
     }
@@ -225,7 +225,7 @@ mod tests {
         let scene = builder.build();
 
         assert_eq!(scene.entities.len(), 1);
-        assert_eq!(scene.entities[0].entity, entity.index());
+        assert_eq!(scene.entities[0].entity, entity.to_bits());
         assert_eq!(scene.entities[0].components.len(), 2);
         assert!(scene.entities[0].components[0].represents::<ComponentA>());
         assert!(scene.entities[0].components[1].represents::<ComponentB>());
@@ -252,10 +252,22 @@ mod tests {
         let mut entities = builder.build().entities.into_iter();
 
         // Assert entities are ordered
-        assert_eq!(entity_a.index(), entities.next().map(|e| e.entity).unwrap());
-        assert_eq!(entity_b.index(), entities.next().map(|e| e.entity).unwrap());
-        assert_eq!(entity_c.index(), entities.next().map(|e| e.entity).unwrap());
-        assert_eq!(entity_d.index(), entities.next().map(|e| e.entity).unwrap());
+        assert_eq!(
+            entity_a.to_bits(),
+            entities.next().map(|e| e.entity).unwrap()
+        );
+        assert_eq!(
+            entity_b.to_bits(),
+            entities.next().map(|e| e.entity).unwrap()
+        );
+        assert_eq!(
+            entity_c.to_bits(),
+            entities.next().map(|e| e.entity).unwrap()
+        );
+        assert_eq!(
+            entity_d.to_bits(),
+            entities.next().map(|e| e.entity).unwrap()
+        );
     }
 
     #[test]
@@ -282,7 +294,7 @@ mod tests {
         assert_eq!(scene.entities.len(), 2);
         let mut scene_entities = vec![scene.entities[0].entity, scene.entities[1].entity];
         scene_entities.sort();
-        assert_eq!(scene_entities, [entity_a_b.index(), entity_a.index()]);
+        assert_eq!(scene_entities, [entity_a_b.to_bits(), entity_a.to_bits()]);
     }
 
     #[test]
@@ -302,6 +314,6 @@ mod tests {
         let scene = builder.build();
 
         assert_eq!(scene.entities.len(), 1);
-        assert_eq!(scene.entities[0].entity, entity_a.index());
+        assert_eq!(scene.entities[0].entity, entity_a.to_bits());
     }
 }