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    /// Implementors 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`] implementors.
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_row_and_generation(
280            self.dead_start.row(),
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: &mut World) -> Self {
340        // We're going to be calling methods on `Entities` that require advance
341        // flushing, such as `alloc` and `free`.
342        world.flush_entities();
343        Self {
344            map,
345            // SAFETY: Entities data is kept in a valid state via `EntityMapper::world_scope`
346            dead_start: unsafe { world.entities_mut().alloc() },
347            generations: 0,
348        }
349    }
350
351    /// Reserves the allocated references to dead entities within the world. This frees the temporary base
352    /// [`Entity`] while reserving extra generations. Because this makes the [`SceneEntityMapper`] unable to
353    /// safely allocate any more references, this method takes ownership of `self` in order to render it unusable.
354    pub fn finish(self, world: &mut World) {
355        // SAFETY: Entities data is kept in a valid state via `EntityMap::world_scope`
356        let entities = unsafe { world.entities_mut() };
357        assert!(entities.free(self.dead_start).is_some());
358        assert!(entities.reserve_generations(self.dead_start.index(), self.generations));
359    }
360
361    /// Creates an [`SceneEntityMapper`] from a provided [`World`] and [`EntityHashMap<Entity>`], then calls the
362    /// provided function with it. This allows one to allocate new entity references in this [`World`] that are
363    /// guaranteed to never point at a living entity now or in the future. This functionality is useful for safely
364    /// mapping entity identifiers that point at entities outside the source world. The passed function, `f`, is called
365    /// within the scope of this world. Its return value is then returned from `world_scope` as the generic type
366    /// parameter `R`.
367    pub fn world_scope<R>(
368        entity_map: &'m mut EntityHashMap<Entity>,
369        world: &mut World,
370        f: impl FnOnce(&mut World, &mut Self) -> R,
371    ) -> R {
372        let mut mapper = Self::new(entity_map, world);
373        let result = f(world, &mut mapper);
374        mapper.finish(world);
375        result
376    }
377}
378
379#[cfg(test)]
380mod tests {
381
382    use crate::{
383        entity::{Entity, EntityHashMap, EntityMapper, SceneEntityMapper},
384        world::World,
385    };
386
387    #[test]
388    fn entity_mapper() {
389        let mut map = EntityHashMap::default();
390        let mut world = World::new();
391        let mut mapper = SceneEntityMapper::new(&mut map, &mut world);
392
393        let mapped_ent = Entity::from_raw_u32(1).unwrap();
394        let dead_ref = mapper.get_mapped(mapped_ent);
395
396        assert_eq!(
397            dead_ref,
398            mapper.get_mapped(mapped_ent),
399            "should persist the allocated mapping from the previous line"
400        );
401        assert_eq!(
402            mapper.get_mapped(Entity::from_raw_u32(2).unwrap()).index(),
403            dead_ref.index(),
404            "should re-use the same index for further dead refs"
405        );
406
407        mapper.finish(&mut world);
408        // Next allocated entity should be a further generation on the same index
409        let entity = world.spawn_empty().id();
410        assert_eq!(entity.index(), dead_ref.index());
411        assert!(entity
412            .generation()
413            .cmp_approx(&dead_ref.generation())
414            .is_gt());
415    }
416
417    #[test]
418    fn world_scope_reserves_generations() {
419        let mut map = EntityHashMap::default();
420        let mut world = World::new();
421
422        let dead_ref = SceneEntityMapper::world_scope(&mut map, &mut world, |_, mapper| {
423            mapper.get_mapped(Entity::from_raw_u32(0).unwrap())
424        });
425
426        // Next allocated entity should be a further generation on the same index
427        let entity = world.spawn_empty().id();
428        assert_eq!(entity.index(), dead_ref.index());
429        assert!(entity
430            .generation()
431            .cmp_approx(&dead_ref.generation())
432            .is_gt());
433    }
434
435    #[test]
436    fn entity_mapper_no_panic() {
437        let mut world = World::new();
438        // "Dirty" the `Entities`, requiring a flush afterward.
439        world.entities.reserve_entity();
440        assert!(world.entities.needs_flush());
441
442        // Create and exercise a SceneEntityMapper - should not panic because it flushes
443        // `Entities` first.
444        SceneEntityMapper::world_scope(&mut Default::default(), &mut world, |_, m| {
445            m.get_mapped(Entity::PLACEHOLDER);
446        });
447
448        // The SceneEntityMapper should leave `Entities` in a flushed state.
449        assert!(!world.entities.needs_flush());
450    }
451}
bevy_ecs/entity/map_entities.rs

bevy_ecs/entity/
map_entities.rs
