bevy_ecs/entity/
map_entities.rs

1use crate::{
2    entity::Entity,
3    identifier::masks::{IdentifierMask, HIGH_MASK},
4    world::World,
5};
6
7use super::EntityHashMap;
8
9/// Operation to map all contained [`Entity`] fields in a type to new values.
10///
11/// As entity IDs are valid only for the [`World`] they're sourced from, using [`Entity`]
12/// as references in components copied from another world will be invalid. This trait
13/// allows defining custom mappings for these references via [`EntityMappers`](EntityMapper), which
14/// inject the entity mapping strategy between your `MapEntities` type and the current world
15/// (usually by using an [`EntityHashMap<Entity>`] between source entities and entities in the
16/// current world).
17///
18/// Implementing this trait correctly is required for properly loading components
19/// with entity references from scenes.
20///
21/// ## Example
22///
23/// ```
24/// use bevy_ecs::prelude::*;
25/// use bevy_ecs::entity::MapEntities;
26///
27/// #[derive(Component)]
28/// struct Spring {
29///     a: Entity,
30///     b: Entity,
31/// }
32///
33/// impl MapEntities for Spring {
34///     fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M) {
35///         self.a = entity_mapper.map_entity(self.a);
36///         self.b = entity_mapper.map_entity(self.b);
37///     }
38/// }
39/// ```
40pub trait MapEntities {
41    /// Updates all [`Entity`] references stored inside using `entity_mapper`.
42    ///
43    /// Implementors should look up any and all [`Entity`] values stored within `self` and
44    /// update them to the mapped values via `entity_mapper`.
45    fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M);
46}
47
48/// An implementor of this trait knows how to map an [`Entity`] into another [`Entity`].
49///
50/// Usually this is done by using an [`EntityHashMap<Entity>`] to map source entities
51/// (mapper inputs) to the current world's entities (mapper outputs).
52///
53/// More generally, this can be used to map [`Entity`] references between any two [`Worlds`](World).
54///
55/// ## Example
56///
57/// ```
58/// # use bevy_ecs::entity::{Entity, EntityMapper};
59/// # use bevy_ecs::entity::EntityHashMap;
60/// #
61/// pub struct SimpleEntityMapper {
62///   map: EntityHashMap<Entity>,
63/// }
64///
65/// // Example implementation of EntityMapper where we map an entity to another entity if it exists
66/// // in the underlying `EntityHashMap`, otherwise we just return the original entity.
67/// impl EntityMapper for SimpleEntityMapper {
68///     fn map_entity(&mut self, entity: Entity) -> Entity {
69///         self.map.get(&entity).copied().unwrap_or(entity)
70///     }
71/// }
72/// ```
73pub trait EntityMapper {
74    /// Map an entity to another entity
75    fn map_entity(&mut self, entity: Entity) -> Entity;
76}
77
78impl EntityMapper for SceneEntityMapper<'_> {
79    /// Returns the corresponding mapped entity or reserves a new dead entity ID in the current world if it is absent.
80    fn map_entity(&mut self, entity: Entity) -> Entity {
81        if let Some(&mapped) = self.map.get(&entity) {
82            return mapped;
83        }
84
85        // this new entity reference is specifically designed to never represent any living entity
86        let new = Entity::from_raw_and_generation(
87            self.dead_start.index(),
88            IdentifierMask::inc_masked_high_by(self.dead_start.generation, self.generations),
89        );
90
91        // Prevent generations counter from being a greater value than HIGH_MASK.
92        self.generations = (self.generations + 1) & HIGH_MASK;
93
94        self.map.insert(entity, new);
95
96        new
97    }
98}
99
100/// A wrapper for [`EntityHashMap<Entity>`], augmenting it with the ability to allocate new [`Entity`] references in a destination
101/// world. These newly allocated references are guaranteed to never point to any living entity in that world.
102///
103/// References are allocated by returning increasing generations starting from an internally initialized base
104/// [`Entity`]. After it is finished being used by [`MapEntities`] implementations, this entity is despawned and the
105/// requisite number of generations reserved.
106pub struct SceneEntityMapper<'m> {
107    /// A mapping from one set of entities to another.
108    ///
109    /// This is typically used to coordinate data transfer between sets of entities, such as between a scene and the world
110    /// or over the network. This is required as [`Entity`] identifiers are opaque; you cannot and do not want to reuse
111    /// identifiers directly.
112    ///
113    /// On its own, a [`EntityHashMap<Entity>`] is not capable of allocating new entity identifiers, which is needed to map references
114    /// to entities that lie outside the source entity set. This functionality can be accessed through [`SceneEntityMapper::world_scope()`].
115    map: &'m mut EntityHashMap<Entity>,
116    /// A base [`Entity`] used to allocate new references.
117    dead_start: Entity,
118    /// The number of generations this mapper has allocated thus far.
119    generations: u32,
120}
121
122impl<'m> SceneEntityMapper<'m> {
123    /// Gets a reference to the underlying [`EntityHashMap<Entity>`].
124    pub fn get_map(&'m self) -> &'m EntityHashMap<Entity> {
125        self.map
126    }
127
128    /// Gets a mutable reference to the underlying [`EntityHashMap<Entity>`].
129    pub fn get_map_mut(&'m mut self) -> &'m mut EntityHashMap<Entity> {
130        self.map
131    }
132
133    /// Creates a new [`SceneEntityMapper`], spawning a temporary base [`Entity`] in the provided [`World`]
134    pub fn new(map: &'m mut EntityHashMap<Entity>, world: &mut World) -> Self {
135        Self {
136            map,
137            // SAFETY: Entities data is kept in a valid state via `EntityMapper::world_scope`
138            dead_start: unsafe { world.entities_mut().alloc() },
139            generations: 0,
140        }
141    }
142
143    /// Reserves the allocated references to dead entities within the world. This frees the temporary base
144    /// [`Entity`] while reserving extra generations via [`crate::entity::Entities::reserve_generations`]. Because this
145    /// renders the [`SceneEntityMapper`] unable to safely allocate any more references, this method takes ownership of
146    /// `self` in order to render it unusable.
147    pub fn finish(self, world: &mut World) {
148        // SAFETY: Entities data is kept in a valid state via `EntityMap::world_scope`
149        let entities = unsafe { world.entities_mut() };
150        assert!(entities.free(self.dead_start).is_some());
151        assert!(entities.reserve_generations(self.dead_start.index(), self.generations));
152    }
153
154    /// Creates an [`SceneEntityMapper`] from a provided [`World`] and [`EntityHashMap<Entity>`], then calls the
155    /// provided function with it. This allows one to allocate new entity references in this [`World`] that are
156    /// guaranteed to never point at a living entity now or in the future. This functionality is useful for safely
157    /// mapping entity identifiers that point at entities outside the source world. The passed function, `f`, is called
158    /// within the scope of this world. Its return value is then returned from `world_scope` as the generic type
159    /// parameter `R`.
160    pub fn world_scope<R>(
161        entity_map: &'m mut EntityHashMap<Entity>,
162        world: &mut World,
163        f: impl FnOnce(&mut World, &mut Self) -> R,
164    ) -> R {
165        let mut mapper = Self::new(entity_map, world);
166        let result = f(world, &mut mapper);
167        mapper.finish(world);
168        result
169    }
170}
171
172#[cfg(test)]
173mod tests {
174    use crate::{
175        entity::{Entity, EntityHashMap, EntityMapper, SceneEntityMapper},
176        world::World,
177    };
178
179    #[test]
180    fn entity_mapper() {
181        const FIRST_IDX: u32 = 1;
182        const SECOND_IDX: u32 = 2;
183
184        let mut map = EntityHashMap::default();
185        let mut world = World::new();
186        let mut mapper = SceneEntityMapper::new(&mut map, &mut world);
187
188        let mapped_ent = Entity::from_raw(FIRST_IDX);
189        let dead_ref = mapper.map_entity(mapped_ent);
190
191        assert_eq!(
192            dead_ref,
193            mapper.map_entity(mapped_ent),
194            "should persist the allocated mapping from the previous line"
195        );
196        assert_eq!(
197            mapper.map_entity(Entity::from_raw(SECOND_IDX)).index(),
198            dead_ref.index(),
199            "should re-use the same index for further dead refs"
200        );
201
202        mapper.finish(&mut world);
203        // Next allocated entity should be a further generation on the same index
204        let entity = world.spawn_empty().id();
205        assert_eq!(entity.index(), dead_ref.index());
206        assert!(entity.generation() > dead_ref.generation());
207    }
208
209    #[test]
210    fn world_scope_reserves_generations() {
211        let mut map = EntityHashMap::default();
212        let mut world = World::new();
213
214        let dead_ref = SceneEntityMapper::world_scope(&mut map, &mut world, |_, mapper| {
215            mapper.map_entity(Entity::from_raw(0))
216        });
217
218        // Next allocated entity should be a further generation on the same index
219        let entity = world.spawn_empty().id();
220        assert_eq!(entity.index(), dead_ref.index());
221        assert!(entity.generation() > dead_ref.generation());
222    }
223}