1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223
use crate::{
entity::Entity,
identifier::masks::{IdentifierMask, HIGH_MASK},
world::World,
};
use super::EntityHashMap;
/// 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 [`EntityMappers`](EntityMapper), which
/// inject the entity mapping strategy between your `MapEntities` type and the current world
/// (usually by using an [`EntityHashMap<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.
///
/// ## Example
///
/// ```
/// use bevy_ecs::prelude::*;
/// use bevy_ecs::entity::MapEntities;
///
/// #[derive(Component)]
/// struct Spring {
/// a: Entity,
/// b: Entity,
/// }
///
/// impl MapEntities for Spring {
/// 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 `self` and
/// update them to the mapped values via `entity_mapper`.
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>`] 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).
///
/// ## Example
///
/// ```
/// # use bevy_ecs::entity::{Entity, EntityMapper};
/// # use bevy_ecs::entity::EntityHashMap;
/// #
/// pub struct SimpleEntityMapper {
/// map: EntityHashMap<Entity>,
/// }
///
/// // Example implementation of EntityMapper where we map an entity to another entity if it exists
/// // in the underlying `EntityHashMap`, otherwise we just return the original entity.
/// impl EntityMapper for SimpleEntityMapper {
/// fn map_entity(&mut self, entity: Entity) -> Entity {
/// self.map.get(&entity).copied().unwrap_or(entity)
/// }
/// }
/// ```
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>`], 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 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>`] 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 [`SceneEntityMapper::world_scope()`].
map: &'m mut EntityHashMap<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> SceneEntityMapper<'m> {
/// Gets a reference to the underlying [`EntityHashMap<Entity>`].
pub fn get_map(&'m self) -> &'m EntityHashMap<Entity> {
self.map
}
/// Gets a mutable reference to the underlying [`EntityHashMap<Entity>`].
pub fn get_map_mut(&'m mut self) -> &'m mut EntityHashMap<Entity> {
self.map
}
/// Creates a new [`SceneEntityMapper`], spawning a temporary base [`Entity`] in the provided [`World`]
pub fn new(map: &'m mut EntityHashMap<Entity>, world: &mut World) -> Self {
Self {
map,
// SAFETY: Entities data is kept in a valid state via `EntityMapper::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 [`SceneEntityMapper`] unable to safely allocate any more references, this method takes ownership of
/// `self` in order to render it unusable.
pub 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));
}
/// Creates an [`SceneEntityMapper`] from a provided [`World`] and [`EntityHashMap<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 EntityHashMap<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 crate::{
entity::{Entity, EntityHashMap, EntityMapper, SceneEntityMapper},
world::World,
};
#[test]
fn entity_mapper() {
const FIRST_IDX: u32 = 1;
const SECOND_IDX: u32 = 2;
let mut map = EntityHashMap::default();
let mut world = World::new();
let mut mapper = SceneEntityMapper::new(&mut map, &mut world);
let mapped_ent = Entity::from_raw(FIRST_IDX);
let dead_ref = mapper.map_entity(mapped_ent);
assert_eq!(
dead_ref,
mapper.map_entity(mapped_ent),
"should persist the allocated mapping from the previous line"
);
assert_eq!(
mapper.map_entity(Entity::from_raw(SECOND_IDX)).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 = EntityHashMap::default();
let mut world = World::new();
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
let entity = world.spawn_empty().id();
assert_eq!(entity.index(), dead_ref.index());
assert!(entity.generation() > dead_ref.generation());
}
}