bevy_ecs/entity/
map_entities.rs

1pub use bevy_ecs_macros::MapEntities;
2use indexmap::{IndexMap, IndexSet};
3
4use crate::{
5    entity::{hash_map::EntityHashMap, Entity},
6    world::World,
7};
8
9use alloc::{
10    collections::{BTreeMap, BTreeSet, VecDeque},
11    vec::Vec,
12};
13use bevy_platform::collections::{HashMap, HashSet};
14use core::{
15    hash::{BuildHasher, Hash},
16    mem,
17};
18use smallvec::SmallVec;
19
20use super::EntityIndexSet;
21
22/// Operation to map all contained [`Entity`] fields in a type to new values.
23///
24/// As entity IDs are valid only for the [`World`] they're sourced from, using [`Entity`]
25/// as references in components copied from another world will be invalid. This trait
26/// allows defining custom mappings for these references via [`EntityMappers`](EntityMapper), which
27/// inject the entity mapping strategy between your `MapEntities` type and the current world
28/// (usually by using an [`EntityHashMap<Entity>`] between source entities and entities in the
29/// current world).
30///
31/// Components use [`Component::map_entities`](crate::component::Component::map_entities) to map
32/// entities in the context of scenes and entity cloning, which generally uses [`MapEntities`] internally
33/// to map each field (see those docs for usage).
34///
35/// [`HashSet<Entity>`]: bevy_platform::collections::HashSet
36///
37/// ## Example
38///
39/// ```
40/// use bevy_ecs::prelude::*;
41/// use bevy_ecs::entity::MapEntities;
42///
43/// #[derive(Component)]
44/// struct Spring {
45///     a: Entity,
46///     b: Entity,
47/// }
48///
49/// impl MapEntities for Spring {
50///     fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M) {
51///         self.a = entity_mapper.get_mapped(self.a);
52///         self.b = entity_mapper.get_mapped(self.b);
53///     }
54/// }
55/// ```
56pub trait MapEntities {
57    /// Updates all [`Entity`] references stored inside using `entity_mapper`.
58    ///
59    /// Implementers should look up any and all [`Entity`] values stored within `self` and
60    /// update them to the mapped values via `entity_mapper`.
61    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E);
62}
63
64impl MapEntities for Entity {
65    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E) {
66        *self = entity_mapper.get_mapped(*self);
67    }
68}
69
70impl<T: MapEntities> MapEntities for Option<T> {
71    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E) {
72        if let Some(entities) = self {
73            entities.map_entities(entity_mapper);
74        }
75    }
76}
77
78impl<K: MapEntities + Eq + Hash, V: MapEntities, S: BuildHasher + Default> MapEntities
79    for HashMap<K, V, S>
80{
81    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E) {
82        *self = self
83            .drain()
84            .map(|(mut key_entities, mut value_entities)| {
85                key_entities.map_entities(entity_mapper);
86                value_entities.map_entities(entity_mapper);
87                (key_entities, value_entities)
88            })
89            .collect();
90    }
91}
92
93impl<T: MapEntities + Eq + Hash, S: BuildHasher + Default> MapEntities for HashSet<T, S> {
94    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E) {
95        *self = self
96            .drain()
97            .map(|mut entities| {
98                entities.map_entities(entity_mapper);
99                entities
100            })
101            .collect();
102    }
103}
104
105impl<K: MapEntities + Eq + Hash, V: MapEntities, S: BuildHasher + Default> MapEntities
106    for IndexMap<K, V, S>
107{
108    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E) {
109        *self = self
110            .drain(..)
111            .map(|(mut key_entities, mut value_entities)| {
112                key_entities.map_entities(entity_mapper);
113                value_entities.map_entities(entity_mapper);
114                (key_entities, value_entities)
115            })
116            .collect();
117    }
118}
119
120impl<T: MapEntities + Eq + Hash, S: BuildHasher + Default> MapEntities for IndexSet<T, S> {
121    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E) {
122        *self = self
123            .drain(..)
124            .map(|mut entities| {
125                entities.map_entities(entity_mapper);
126                entities
127            })
128            .collect();
129    }
130}
131
132impl MapEntities for EntityIndexSet {
133    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E) {
134        *self = self
135            .drain(..)
136            .map(|e| entity_mapper.get_mapped(e))
137            .collect();
138    }
139}
140
141impl<K: MapEntities + Ord, V: MapEntities> MapEntities for BTreeMap<K, V> {
142    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E) {
143        *self = mem::take(self)
144            .into_iter()
145            .map(|(mut key_entities, mut value_entities)| {
146                key_entities.map_entities(entity_mapper);
147                value_entities.map_entities(entity_mapper);
148                (key_entities, value_entities)
149            })
150            .collect();
151    }
152}
153
154impl<T: MapEntities + Ord> MapEntities for BTreeSet<T> {
155    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E) {
156        *self = mem::take(self)
157            .into_iter()
158            .map(|mut entities| {
159                entities.map_entities(entity_mapper);
160                entities
161            })
162            .collect();
163    }
164}
165
166impl<T: MapEntities, const N: usize> MapEntities for [T; N] {
167    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E) {
168        for entities in self.iter_mut() {
169            entities.map_entities(entity_mapper);
170        }
171    }
172}
173
174impl<T: MapEntities> MapEntities for Vec<T> {
175    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E) {
176        for entities in self.iter_mut() {
177            entities.map_entities(entity_mapper);
178        }
179    }
180}
181
182impl<T: MapEntities> MapEntities for VecDeque<T> {
183    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E) {
184        for entities in self.iter_mut() {
185            entities.map_entities(entity_mapper);
186        }
187    }
188}
189
190impl<T: MapEntities, A: smallvec::Array<Item = T>> MapEntities for SmallVec<A> {
191    fn map_entities<E: EntityMapper>(&mut self, entity_mapper: &mut E) {
192        for entities in self.iter_mut() {
193            entities.map_entities(entity_mapper);
194        }
195    }
196}
197
198/// An implementor of this trait knows how to map an [`Entity`] into another [`Entity`].
199///
200/// Usually this is done by using an [`EntityHashMap<Entity>`] to map source entities
201/// (mapper inputs) to the current world's entities (mapper outputs).
202///
203/// More generally, this can be used to map [`Entity`] references between any two [`Worlds`](World).
204///
205/// This is used by [`MapEntities`] implementers.
206///
207/// ## Example
208///
209/// ```
210/// # use bevy_ecs::entity::{Entity, EntityMapper};
211/// # use bevy_ecs::entity::EntityHashMap;
212/// #
213/// pub struct SimpleEntityMapper {
214///   map: EntityHashMap<Entity>,
215/// }
216///
217/// // Example implementation of EntityMapper where we map an entity to another entity if it exists
218/// // in the underlying `EntityHashMap`, otherwise we just return the original entity.
219/// impl EntityMapper for SimpleEntityMapper {
220///     fn get_mapped(&mut self, entity: Entity) -> Entity {
221///         self.map.get(&entity).copied().unwrap_or(entity)
222///     }
223///
224///     fn set_mapped(&mut self, source: Entity, target: Entity) {
225///         self.map.insert(source, target);
226///     }
227/// }
228/// ```
229pub trait EntityMapper {
230    /// Returns the "target" entity that maps to the given `source`.
231    fn get_mapped(&mut self, source: Entity) -> Entity;
232
233    /// Maps the `target` entity to the given `source`. For some implementations this might not actually determine the result
234    /// of [`EntityMapper::get_mapped`].
235    fn set_mapped(&mut self, source: Entity, target: Entity);
236}
237
238impl EntityMapper for () {
239    #[inline]
240    fn get_mapped(&mut self, source: Entity) -> Entity {
241        source
242    }
243
244    #[inline]
245    fn set_mapped(&mut self, _source: Entity, _target: Entity) {}
246}
247
248impl EntityMapper for (Entity, Entity) {
249    #[inline]
250    fn get_mapped(&mut self, source: Entity) -> Entity {
251        if source == self.0 {
252            self.1
253        } else {
254            source
255        }
256    }
257
258    fn set_mapped(&mut self, _source: Entity, _target: Entity) {}
259}
260
261impl EntityMapper for &mut dyn EntityMapper {
262    fn get_mapped(&mut self, source: Entity) -> Entity {
263        (*self).get_mapped(source)
264    }
265
266    fn set_mapped(&mut self, source: Entity, target: Entity) {
267        (*self).set_mapped(source, target);
268    }
269}
270
271impl EntityMapper for SceneEntityMapper<'_> {
272    /// Returns the corresponding mapped entity or reserves a new dead entity ID in the current world if it is absent.
273    fn get_mapped(&mut self, source: Entity) -> Entity {
274        if let Some(&mapped) = self.map.get(&source) {
275            return mapped;
276        }
277
278        // this new entity reference is specifically designed to never represent any living entity
279        let new = Entity::from_index_and_generation(
280            self.dead_start.index(),
281            self.dead_start.generation.after_versions(self.generations),
282        );
283        self.generations = self.generations.wrapping_add(1);
284
285        self.map.insert(source, new);
286
287        new
288    }
289
290    fn set_mapped(&mut self, source: Entity, target: Entity) {
291        self.map.insert(source, target);
292    }
293}
294
295impl EntityMapper for EntityHashMap<Entity> {
296    /// Returns the corresponding mapped entity or returns `entity` if there is no mapped entity
297    fn get_mapped(&mut self, source: Entity) -> Entity {
298        self.get(&source).cloned().unwrap_or(source)
299    }
300
301    fn set_mapped(&mut self, source: Entity, target: Entity) {
302        self.insert(source, target);
303    }
304}
305
306/// A wrapper for [`EntityHashMap<Entity>`], augmenting it with the ability to allocate new [`Entity`] references in a destination
307/// world. These newly allocated references are guaranteed to never point to any living entity in that world.
308///
309/// References are allocated by returning increasing generations starting from an internally initialized base
310/// [`Entity`]. After it is finished being used, this entity is despawned and the requisite number of generations reserved.
311pub struct SceneEntityMapper<'m> {
312    /// A mapping from one set of entities to another.
313    ///
314    /// This is typically used to coordinate data transfer between sets of entities, such as between a scene and the world
315    /// or over the network. This is required as [`Entity`] identifiers are opaque; you cannot and do not want to reuse
316    /// identifiers directly.
317    ///
318    /// On its own, a [`EntityHashMap<Entity>`] is not capable of allocating new entity identifiers, which is needed to map references
319    /// to entities that lie outside the source entity set. This functionality can be accessed through [`SceneEntityMapper::world_scope()`].
320    map: &'m mut EntityHashMap<Entity>,
321    /// A base [`Entity`] used to allocate new references.
322    dead_start: Entity,
323    /// The number of generations this mapper has allocated thus far.
324    generations: u32,
325}
326
327impl<'m> SceneEntityMapper<'m> {
328    /// Gets a reference to the underlying [`EntityHashMap<Entity>`].
329    pub fn get_map(&'m self) -> &'m EntityHashMap<Entity> {
330        self.map
331    }
332
333    /// Gets a mutable reference to the underlying [`EntityHashMap<Entity>`].
334    pub fn get_map_mut(&'m mut self) -> &'m mut EntityHashMap<Entity> {
335        self.map
336    }
337
338    /// Creates a new [`SceneEntityMapper`], spawning a temporary base [`Entity`] in the provided [`World`]
339    pub fn new(map: &'m mut EntityHashMap<Entity>, world: &World) -> Self {
340        Self {
341            map,
342            dead_start: world.allocator.alloc(),
343            generations: 0,
344        }
345    }
346
347    /// Reserves the allocated references to dead entities within the world. This frees the temporary base
348    /// [`Entity`] while reserving extra generations. Because this makes the [`SceneEntityMapper`] unable to
349    /// safely allocate any more references, this method takes ownership of `self` in order to render it unusable.
350    pub fn finish(self, world: &mut World) {
351        // SAFETY: We never constructed the entity and never released it for something else to construct.
352        let reuse_row = unsafe {
353            world
354                .entities
355                .mark_free(self.dead_start.index(), self.generations)
356        };
357        world.allocator.free(reuse_row);
358    }
359
360    /// Creates an [`SceneEntityMapper`] from a provided [`World`] and [`EntityHashMap<Entity>`], then calls the
361    /// provided function with it. This allows one to allocate new entity references in this [`World`] that are
362    /// guaranteed to never point at a living entity now or in the future. This functionality is useful for safely
363    /// mapping entity identifiers that point at entities outside the source world. The passed function, `f`, is called
364    /// within the scope of this world. Its return value is then returned from `world_scope` as the generic type
365    /// parameter `R`.
366    pub fn world_scope<R>(
367        entity_map: &'m mut EntityHashMap<Entity>,
368        world: &mut World,
369        f: impl FnOnce(&mut World, &mut Self) -> R,
370    ) -> R {
371        let mut mapper = Self::new(entity_map, world);
372        let result = f(world, &mut mapper);
373        mapper.finish(world);
374        result
375    }
376}
377
378#[cfg(test)]
379mod tests {
380
381    use crate::{
382        entity::{Entity, EntityHashMap, EntityMapper, SceneEntityMapper},
383        world::World,
384    };
385
386    #[test]
387    fn entity_mapper() {
388        let mut map = EntityHashMap::default();
389        let mut world = World::new();
390        let mut mapper = SceneEntityMapper::new(&mut map, &world);
391
392        let mapped_ent = Entity::from_raw_u32(1).unwrap();
393        let dead_ref = mapper.get_mapped(mapped_ent);
394
395        assert_eq!(
396            dead_ref,
397            mapper.get_mapped(mapped_ent),
398            "should persist the allocated mapping from the previous line"
399        );
400        assert_eq!(
401            mapper.get_mapped(Entity::from_raw_u32(2).unwrap()).index(),
402            dead_ref.index(),
403            "should re-use the same index for further dead refs"
404        );
405
406        mapper.finish(&mut world);
407        // Next allocated entity should be a further generation on the same index
408        let entity = world.spawn_empty().id();
409        assert_eq!(entity.index(), dead_ref.index());
410        assert!(entity
411            .generation()
412            .cmp_approx(&dead_ref.generation())
413            .is_gt());
414    }
415
416    #[test]
417    fn world_scope_reserves_generations() {
418        let mut map = EntityHashMap::default();
419        let mut world = World::new();
420
421        let dead_ref = SceneEntityMapper::world_scope(&mut map, &mut world, |_, mapper| {
422            mapper.get_mapped(Entity::from_raw_u32(0).unwrap())
423        });
424
425        // Next allocated entity should be a further generation on the same index
426        let entity = world.spawn_empty().id();
427        assert_eq!(entity.index(), dead_ref.index());
428        assert!(entity
429            .generation()
430            .cmp_approx(&dead_ref.generation())
431            .is_gt());
432    }
433}
bevy_ecs/entity/map_entities.rs

bevy_ecs/entity/
map_entities.rs
