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());
    }
}