bevy_ecs/entity/
visit_entities.rs

1pub use bevy_ecs_macros::{VisitEntities, VisitEntitiesMut};
2
3use crate::entity::Entity;
4
5/// Apply an operation to all entities in a container.
6///
7/// This is implemented by default for types that implement [`IntoIterator`].
8///
9/// It may be useful to implement directly for types that can't produce an
10/// iterator for lifetime reasons, such as those involving internal mutexes.
11pub trait VisitEntities {
12    /// Apply an operation to all contained entities.
13    fn visit_entities<F: FnMut(Entity)>(&self, f: F);
14}
15
16impl<T> VisitEntities for T
17where
18    for<'a> &'a T: IntoIterator<Item = &'a Entity>,
19{
20    fn visit_entities<F: FnMut(Entity)>(&self, f: F) {
21        self.into_iter().copied().for_each(f);
22    }
23}
24
25impl VisitEntities for Entity {
26    fn visit_entities<F: FnMut(Entity)>(&self, mut f: F) {
27        f(*self);
28    }
29}
30
31/// Apply an operation to mutable references to all entities in a container.
32///
33/// This is implemented by default for types that implement [`IntoIterator`].
34///
35/// It may be useful to implement directly for types that can't produce an
36/// iterator for lifetime reasons, such as those involving internal mutexes.
37pub trait VisitEntitiesMut: VisitEntities {
38    /// Apply an operation to mutable references to all contained entities.
39    fn visit_entities_mut<F: FnMut(&mut Entity)>(&mut self, f: F);
40}
41
42impl<T: VisitEntities> VisitEntitiesMut for T
43where
44    for<'a> &'a mut T: IntoIterator<Item = &'a mut Entity>,
45{
46    fn visit_entities_mut<F: FnMut(&mut Entity)>(&mut self, f: F) {
47        self.into_iter().for_each(f);
48    }
49}
50
51impl VisitEntitiesMut for Entity {
52    fn visit_entities_mut<F: FnMut(&mut Entity)>(&mut self, mut f: F) {
53        f(self);
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use crate::{
60        self as bevy_ecs,
61        entity::{EntityHashMap, MapEntities, SceneEntityMapper},
62        world::World,
63    };
64    use bevy_utils::HashSet;
65
66    use super::*;
67
68    #[derive(VisitEntities, Debug, PartialEq)]
69    struct Foo {
70        ordered: Vec<Entity>,
71        unordered: HashSet<Entity>,
72        single: Entity,
73        #[allow(dead_code)]
74        #[visit_entities(ignore)]
75        not_an_entity: String,
76    }
77
78    // Need a manual impl since VisitEntitiesMut isn't implemented for `HashSet`.
79    // We don't expect users to actually do this - it's only for test purposes
80    // to prove out the automatic `MapEntities` impl we get with `VisitEntitiesMut`.
81    impl VisitEntitiesMut for Foo {
82        fn visit_entities_mut<F: FnMut(&mut Entity)>(&mut self, mut f: F) {
83            self.ordered.visit_entities_mut(&mut f);
84            self.unordered = self
85                .unordered
86                .drain()
87                .map(|mut entity| {
88                    f(&mut entity);
89                    entity
90                })
91                .collect();
92            f(&mut self.single);
93        }
94    }
95
96    #[test]
97    fn visit_entities() {
98        let mut world = World::new();
99        let entities = world.entities();
100        let mut foo = Foo {
101            ordered: vec![entities.reserve_entity(), entities.reserve_entity()],
102            unordered: [
103                entities.reserve_entity(),
104                entities.reserve_entity(),
105                entities.reserve_entity(),
106            ]
107            .into_iter()
108            .collect(),
109            single: entities.reserve_entity(),
110            not_an_entity: "Bar".into(),
111        };
112
113        let mut entity_map = EntityHashMap::<Entity>::default();
114        let mut remapped = Foo {
115            ordered: vec![],
116            unordered: HashSet::new(),
117            single: Entity::PLACEHOLDER,
118            not_an_entity: foo.not_an_entity.clone(),
119        };
120
121        // Note: this assumes that the VisitEntities derive is field-ordered,
122        //       which isn't explicitly stated/guaranteed.
123        //       If that changes, this test will fail, but that might be OK if
124        //       we're intentionally breaking that assumption.
125        let mut i = 0;
126        foo.visit_entities(|entity| {
127            let new_entity = entities.reserve_entity();
128            if i < foo.ordered.len() {
129                assert_eq!(entity, foo.ordered[i]);
130                remapped.ordered.push(new_entity);
131            } else if i < foo.ordered.len() + foo.unordered.len() {
132                assert!(foo.unordered.contains(&entity));
133                remapped.unordered.insert(new_entity);
134            } else {
135                assert_eq!(entity, foo.single);
136                remapped.single = new_entity;
137            }
138
139            entity_map.insert(entity, new_entity);
140
141            i += 1;
142        });
143
144        SceneEntityMapper::world_scope(&mut entity_map, &mut world, |_, mapper| {
145            foo.map_entities(mapper);
146        });
147
148        assert_eq!(foo, remapped);
149    }
150}
bevy_ecs/entity/visit_entities.rs

bevy_ecs/entity/
visit_entities.rs
