Make scene handling of entity references robust

crates/bevy_ecs/Cargo.toml

@@ -26,7 +26,7 @@ thread_local = "1.1.4"
 fixedbitset = "0.4.2"
 rustc-hash = "1.1"
 downcast-rs = "1.2"
-serde = { version = "1", features = ["derive"] }
+serde = "1"
 
 [dev-dependencies]
 rand = "0.8"

crates/bevy_ecs/src/entity/map_entities.rs

@@ -1,24 +1,5 @@
-use crate::entity::Entity;
+use crate::{entity::Entity, world::World};
 use bevy_utils::{Entry, HashMap};
-use std::fmt;
-
-/// The errors that might be returned while using [`MapEntities::map_entities`].
-#[derive(Debug)]
-pub enum MapEntitiesError {
-    EntityNotFound(Entity),
-}
-
-impl std::error::Error for MapEntitiesError {}
-
-impl fmt::Display for MapEntitiesError {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        match self {
-            MapEntitiesError::EntityNotFound(_) => {
-                write!(f, "the given entity does not exist in the map")
-            }
-        }
-    }
-}
 
 /// Operation to map all contained [`Entity`] fields in a type to new values.
 ///
@@ -33,7 +14,7 @@ impl fmt::Display for MapEntitiesError {
 ///
 /// ```rust
 /// use bevy_ecs::prelude::*;
-/// use bevy_ecs::entity::{EntityMap, MapEntities, MapEntitiesError};
+/// use bevy_ecs::entity::{EntityMapper, MapEntities};
 ///
 /// #[derive(Component)]
 /// struct Spring {
@@ -42,10 +23,9 @@ impl fmt::Display for MapEntitiesError {
 /// }
 ///
 /// impl MapEntities for Spring {
-///     fn map_entities(&mut self, entity_map: &EntityMap) -> Result<(), MapEntitiesError> {
-///         self.a = entity_map.get(self.a)?;
-///         self.b = entity_map.get(self.b)?;
-///         Ok(())
+///     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);
 ///     }
 /// }
 /// ```
@@ -55,16 +35,22 @@ pub trait MapEntities {
     /// Updates all [`Entity`] references stored inside using `entity_map`.
     ///
     /// 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>;
+    /// 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.
+/// 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>,
@@ -91,11 +77,8 @@ impl EntityMap {
     }
 
     /// Returns the corresponding mapped entity.
-    pub fn get(&self, entity: Entity) -> Result<Entity, MapEntitiesError> {
-        self.map
-            .get(&entity)
-            .cloned()
-            .ok_or(MapEntitiesError::EntityNotFound(entity))
+    pub fn get(&self, entity: Entity) -> Option<Entity> {
+        self.map.get(&entity).copied()
     }
 
     /// An iterator visiting all keys in arbitrary order.
@@ -122,4 +105,138 @@ impl EntityMap {
     pub fn iter(&self) -> impl Iterator<Item = (Entity, Entity)> + '_ {
         self.map.iter().map(|(from, to)| (*from, *to))
     }
+
+    /// 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
+/// 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 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 {
+            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,
+            // SAFETY: Entities data is kept in a valid state via `EntityMap::world_scope`
+            dead_start: unsafe { world.entities_mut().alloc() },
+            generations: 0,
+        }
+    }
+
+    /// 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
+    /// `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`
+        let entities = unsafe { world.entities_mut() };
+        assert!(entities.free(self.dead_start).is_some());
+        assert!(entities.reserve_generations(self.dead_start.index, self.generations));
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{EntityMap, EntityMapper};
+    use crate::{entity::Entity, world::World};
+
+    #[test]
+    fn entity_mapper() {
+        const FIRST_IDX: u32 = 1;
+        const SECOND_IDX: u32 = 2;
+
+        let mut map = EntityMap::default();
+        let mut world = World::new();
+        let mut mapper = EntityMapper::new(&mut map, &mut world);
+
+        let mapped_ent = Entity::new(FIRST_IDX, 0);
+        let dead_ref = mapper.get_or_reserve(mapped_ent);
+
+        assert_eq!(
+            dead_ref,
+            mapper.get_or_reserve(mapped_ent),
+            "should persist the allocated mapping from the previous line"
+        );
+        assert_eq!(
+            mapper.get_or_reserve(Entity::new(SECOND_IDX, 0)).index(),
+            dead_ref.index(),
+            "should re-use the same index for further dead refs"
+        );
+
+        mapper.finish(&mut world);
+        // Next allocated entity should be a further generation on the same index
+        let entity = world.spawn_empty().id();
+        assert_eq!(entity.index(), dead_ref.index());
+        assert!(entity.generation() > dead_ref.generation());
+    }
+
+    #[test]
+    fn world_scope_reserves_generations() {
+        let mut map = EntityMap::default();
+        let mut world = World::new();
+
+        let dead_ref = map.world_scope(&mut world, |_, mapper| {
+            mapper.get_or_reserve(Entity::new(0, 0))
+        });
+
+        // Next allocated entity should be a further generation on the same index
+        let entity = world.spawn_empty().id();
+        assert_eq!(entity.index(), dead_ref.index());
+        assert!(entity.generation() > dead_ref.generation());
+    }
 }

crates/bevy_ecs/src/entity/mod.rs

@@ -115,7 +115,7 @@ type IdCursor = isize;
 /// [`EntityCommands`]: crate::system::EntityCommands
 /// [`Query::get`]: crate::system::Query::get
 /// [`World`]: crate::world::World
-#[derive(Clone, Copy, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
+#[derive(Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
 pub struct Entity {
     generation: u32,
     index: u32,
@@ -227,6 +227,25 @@ impl Entity {
     }
 }
 
+impl Serialize for Entity {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        serializer.serialize_u64(self.to_bits())
+    }
+}
+
+impl<'de> Deserialize<'de> for Entity {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let id: u64 = serde::de::Deserialize::deserialize(deserializer)?;
+        Ok(Entity::from_bits(id))
+    }
+}
+
 impl fmt::Debug for Entity {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "{}v{}", self.index, self.generation)
@@ -596,6 +615,25 @@ impl Entities {
         self.meta.get_unchecked_mut(index as usize).location = location;
     }
 
+    /// Increments the `generation` of a freed [`Entity`]. The next entity ID allocated with this
+    /// `index` will count `generation` starting from the prior `generation` + the specified
+    /// value + 1.
+    ///
+    /// Does nothing if no entity with this `index` has been allocated yet.
+    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
     ///
@@ -839,4 +877,29 @@ mod tests {
         const C4: u32 = Entity::from_bits(0x00dd_00ff_0000_0000).generation();
         assert_eq!(0x00dd_00ff, C4);
     }
+
+    #[test]
+    fn reserve_generations() {
+        let mut entities = Entities::new();
+        let entity = entities.alloc();
+        entities.free(entity);
+
+        assert!(entities.reserve_generations(entity.index, 1));
+    }
+
+    #[test]
+    fn reserve_generations_and_alloc() {
+        const GENERATIONS: u32 = 10;
+
+        let mut entities = Entities::new();
+        let entity = entities.alloc();
+        entities.free(entity);
+
+        assert!(entities.reserve_generations(entity.index, GENERATIONS));
+
+        // The very next entitiy allocated should be a further generation on the same index
+        let next_entity = entities.alloc();
+        assert_eq!(next_entity.index(), entity.index());
+        assert!(next_entity.generation > entity.generation + GENERATIONS);
+    }
 }

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},
     system::Resource,
     world::{
         unsafe_world_cell::{UnsafeEntityCell, UnsafeWorldCell},
@@ -404,29 +404,25 @@ 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),
 }
 
 impl ReflectMapEntities {
-    pub fn map_entities(
-        &self,
-        world: &mut World,
-        entity_map: &EntityMap,
-    ) -> Result<(), MapEntitiesError> {
-        (self.map_entities)(world, entity_map)
+    pub fn map_entities(&self, world: &mut World, entity_map: &mut EntityMap) {
+        entity_map.world_scope(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() {
-                    if let Some(mut component) = world.get_mut::<C>(entity) {
-                        component.map_entities(entity_map)?;
+            map_entities: |world, entity_mapper| {
+                let entities = entity_mapper.get_map().values().collect::<Vec<Entity>>();
+                for entity in &entities {
+                    if let Some(mut component) = world.get_mut::<C>(*entity) {
+                        component.map_entities(entity_mapper);
                     }
                 }
-                Ok(())
             },
         }
     }