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

bevy_ecs/entity/
map_entities.rs
