Add VisitEntities for generic and reflectable Entity iteration

crates/bevy_animation/src/lib.rs

@@ -28,7 +28,10 @@ use bevy_app::{App, Plugin, PostUpdate};
 use bevy_asset::{Asset, AssetApp, Assets, Handle};
 use bevy_core::Name;
 use bevy_ecs::{
-    entity::MapEntities, prelude::*, reflect::ReflectMapEntities, world::EntityMutExcept,
+    entity::{VisitEntities, VisitEntitiesMut},
+    prelude::*,
+    reflect::{ReflectMapEntities, ReflectVisitEntities, ReflectVisitEntitiesMut},
+    world::EntityMutExcept,
 };
 use bevy_math::FloatExt;
 use bevy_reflect::{
@@ -527,12 +530,13 @@ impl Hash for AnimationTargetId {
 /// Note that each entity can only be animated by one animation player at a
 /// time. However, you can change [`AnimationTarget`]'s `player` property at
 /// runtime to change which player is responsible for animating the entity.
-#[derive(Clone, Copy, Component, Reflect)]
-#[reflect(Component, MapEntities)]
+#[derive(Clone, Copy, Component, Reflect, VisitEntities, VisitEntitiesMut)]
+#[reflect(Component, MapEntities, VisitEntities, VisitEntitiesMut)]
 pub struct AnimationTarget {
     /// The ID of this animation target.
     ///
     /// Typically, this is derived from the path.
+    #[visit_entities(ignore)]
     pub id: AnimationTargetId,
 
     /// The entity containing the [`AnimationPlayer`].
@@ -1298,12 +1302,6 @@ impl From<&Name> for AnimationTargetId {
     }
 }
 
-impl MapEntities for AnimationTarget {
-    fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M) {
-        self.player = entity_mapper.map_entity(self.player);
-    }
-}
-
 impl AnimationGraphEvaluator {
     // Starts a new depth-first search.
     fn reset(&mut self, root: AnimationNodeIndex, node_count: usize) {

crates/bevy_ecs/macros/src/lib.rs

@@ -14,6 +14,7 @@ use crate::{query_data::derive_query_data_impl, query_filter::derive_query_filte
 use bevy_macro_utils::{derive_label, ensure_no_collision, get_struct_fields, BevyManifest};
 use proc_macro::TokenStream;
 use proc_macro2::Span;
+use proc_macro2::TokenStream as TokenStream2;
 use quote::{format_ident, quote};
 use syn::{
     parse_macro_input, parse_quote, punctuated::Punctuated, spanned::Spanned, token::Comma,
@@ -180,6 +181,110 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
     })
 }
 
+fn derive_visit_entities_base(
+    input: TokenStream,
+    trait_name: TokenStream2,
+    gen_methods: impl FnOnce(Vec<TokenStream2>) -> TokenStream2,
+) -> TokenStream {
+    let ast = parse_macro_input!(input as DeriveInput);
+    let ecs_path = bevy_ecs_path();
+
+    let named_fields = match get_struct_fields(&ast.data) {
+        Ok(fields) => fields,
+        Err(e) => return e.into_compile_error().into(),
+    };
+
+    let field = named_fields
+        .iter()
+        .filter_map(|field| {
+            if let Some(attr) = field
+                .attrs
+                .iter()
+                .find(|a| a.path().is_ident("visit_entities"))
+            {
+                let ignore = attr.parse_nested_meta(|meta| {
+                    if meta.path.is_ident("ignore") {
+                        Ok(())
+                    } else {
+                        Err(meta.error("Invalid visit_entities attribute. Use `ignore`"))
+                    }
+                });
+                return match ignore {
+                    Ok(()) => None,
+                    Err(e) => Some(Err(e)),
+                };
+            }
+            Some(Ok(field))
+        })
+        .map(|res| res.map(|field| field.ident.as_ref()))
+        .collect::<Result<Vec<_>, _>>();
+
+    let field = match field {
+        Ok(field) => field,
+        Err(e) => return e.into_compile_error().into(),
+    };
+
+    if field.is_empty() {
+        return syn::Error::new(
+            ast.span(),
+            format!("Invalid `{}` type: at least one field", trait_name),
+        )
+        .into_compile_error()
+        .into();
+    }
+
+    let field_access = field
+        .iter()
+        .enumerate()
+        .map(|(n, f)| {
+            if let Some(ident) = f {
+                quote! {
+                    self.#ident
+                }
+            } else {
+                let idx = Index::from(n);
+                quote! {
+                    self.#idx
+                }
+            }
+        })
+        .collect::<Vec<_>>();
+
+    let methods = gen_methods(field_access);
+
+    let generics = ast.generics;
+    let (impl_generics, ty_generics, _) = generics.split_for_impl();
+    let struct_name = &ast.ident;
+
+    TokenStream::from(quote! {
+        impl #impl_generics #ecs_path::entity:: #trait_name for #struct_name #ty_generics {
+            #methods
+        }
+    })
+}
+
+#[proc_macro_derive(VisitEntitiesMut, attributes(visit_entities))]
+pub fn derive_visit_entities_mut(input: TokenStream) -> TokenStream {
+    derive_visit_entities_base(input, quote! { VisitEntitiesMut }, |field| {
+        quote! {
+            fn visit_entities_mut<F: FnMut(&mut Entity)>(&mut self, mut f: F) {
+                #(#field.visit_entities_mut(&mut f);)*
+            }
+        }
+    })
+}
+
+#[proc_macro_derive(VisitEntities, attributes(visit_entities))]
+pub fn derive_visit_entities(input: TokenStream) -> TokenStream {
+    derive_visit_entities_base(input, quote! { VisitEntities }, |field| {
+        quote! {
+            fn visit_entities<F: FnMut(Entity)>(&self, mut f: F) {
+                #(#field.visit_entities(&mut f);)*
+            }
+        }
+    })
+}
+
 fn get_idents(fmt_string: fn(usize) -> String, count: usize) -> Vec<Ident> {
     (0..count)
         .map(|i| Ident::new(&fmt_string(i), Span::call_site()))

crates/bevy_ecs/src/entity/map_entities.rs

@@ -4,7 +4,7 @@ use crate::{
     world::World,
 };
 
-use super::EntityHashMap;
+use super::{EntityHashMap, VisitEntitiesMut};
 
 /// Operation to map all contained [`Entity`] fields in a type to new values.
 ///
@@ -45,6 +45,14 @@ pub trait MapEntities {
     fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M);
 }
 
+impl<T: VisitEntitiesMut> MapEntities for T {
+    fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M) {
+        self.visit_entities_mut(|entity| {
+            *entity = entity_mapper.map_entity(*entity);
+        });
+    }
+}
+
 /// 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
@@ -122,6 +130,16 @@ impl<T: EntityMapper> DynEntityMapper for T {
     }
 }
 
+impl<'a> EntityMapper for &'a mut dyn DynEntityMapper {
+    fn map_entity(&mut self, entity: Entity) -> Entity {
+        (*self).dyn_map_entity(entity)
+    }
+
+    fn mappings(&self) -> impl Iterator<Item = (Entity, Entity)> {
+        (*self).dyn_mappings().into_iter()
+    }
+}
+
 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 {
@@ -152,8 +170,7 @@ impl EntityMapper for SceneEntityMapper<'_> {
 /// 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.
+/// [`Entity`]. After it is finished being used, this entity is despawned and the requisite number of generations reserved.
 pub struct SceneEntityMapper<'m> {
     /// A mapping from one set of entities to another.
     ///

crates/bevy_ecs/src/entity/mod.rs

@@ -36,11 +36,13 @@
 //! [`EntityWorldMut::insert`]: crate::world::EntityWorldMut::insert
 //! [`EntityWorldMut::remove`]: crate::world::EntityWorldMut::remove
 mod map_entities;
+mod visit_entities;
 #[cfg(feature = "bevy_reflect")]
 use bevy_reflect::Reflect;
 #[cfg(all(feature = "bevy_reflect", feature = "serialize"))]
 use bevy_reflect::{ReflectDeserialize, ReflectSerialize};
 pub use map_entities::*;
+pub use visit_entities::*;
 
 mod hash;
 pub use hash::*;

crates/bevy_ecs/src/entity/visit_entities.rs

@@ -0,0 +1,150 @@
+pub use bevy_ecs_macros::{VisitEntities, VisitEntitiesMut};
+
+use crate::entity::Entity;
+
+/// Apply an operation to all entities in a container.
+///
+/// This is implemented by default for types that implement [`IntoIterator`].
+///
+/// It may be useful to implement directly for types that can't produce an
+/// iterator for lifetime reasons, such as those involving internal mutexes.
+pub trait VisitEntities {
+    /// Apply an operation to all contained entities.
+    fn visit_entities<F: FnMut(Entity)>(&self, f: F);
+}
+
+impl<T> VisitEntities for T
+where
+    for<'a> &'a T: IntoIterator<Item = &'a Entity>,
+{
+    fn visit_entities<F: FnMut(Entity)>(&self, f: F) {
+        self.into_iter().copied().for_each(f);
+    }
+}
+
+impl VisitEntities for Entity {
+    fn visit_entities<F: FnMut(Entity)>(&self, mut f: F) {
+        f(*self);
+    }
+}
+
+/// Apply an operation to mutable references to all entities in a container.
+///
+/// This is implemented by default for types that implement [`IntoIterator`].
+///
+/// It may be useful to implement directly for types that can't produce an
+/// iterator for lifetime reasons, such as those involving internal mutexes.
+pub trait VisitEntitiesMut: VisitEntities {
+    /// Apply an operation to mutable references to all contained entities.
+    fn visit_entities_mut<F: FnMut(&mut Entity)>(&mut self, f: F);
+}
+
+impl<T: VisitEntities> VisitEntitiesMut for T
+where
+    for<'a> &'a mut T: IntoIterator<Item = &'a mut Entity>,
+{
+    fn visit_entities_mut<F: FnMut(&mut Entity)>(&mut self, f: F) {
+        self.into_iter().for_each(f);
+    }
+}
+
+impl VisitEntitiesMut for Entity {
+    fn visit_entities_mut<F: FnMut(&mut Entity)>(&mut self, mut f: F) {
+        f(self);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        self as bevy_ecs,
+        entity::{EntityHashMap, MapEntities, SceneEntityMapper},
+        world::World,
+    };
+    use bevy_utils::HashSet;
+
+    use super::*;
+
+    #[derive(VisitEntities, Debug, PartialEq)]
+    struct Foo {
+        ordered: Vec<Entity>,
+        unordered: HashSet<Entity>,
+        single: Entity,
+        #[allow(dead_code)]
+        #[visit_entities(ignore)]
+        not_an_entity: String,
+    }
+
+    // Need a manual impl since VisitEntitiesMut isn't implemented for `HashSet`.
+    // We don't expect users to actually do this - it's only for test purposes
+    // to prove out the automatic `MapEntities` impl we get with `VisitEntitiesMut`.
+    impl VisitEntitiesMut for Foo {
+        fn visit_entities_mut<F: FnMut(&mut Entity)>(&mut self, mut f: F) {
+            self.ordered.visit_entities_mut(&mut f);
+            self.unordered = self
+                .unordered
+                .drain()
+                .map(|mut entity| {
+                    f(&mut entity);
+                    entity
+                })
+                .collect();
+            f(&mut self.single);
+        }
+    }
+
+    #[test]
+    fn visit_entities() {
+        let mut world = World::new();
+        let entities = world.entities();
+        let mut foo = Foo {
+            ordered: vec![entities.reserve_entity(), entities.reserve_entity()],
+            unordered: [
+                entities.reserve_entity(),
+                entities.reserve_entity(),
+                entities.reserve_entity(),
+            ]
+            .into_iter()
+            .collect(),
+            single: entities.reserve_entity(),
+            not_an_entity: "Bar".into(),
+        };
+
+        let mut entity_map = EntityHashMap::<Entity>::default();
+        let mut remapped = Foo {
+            ordered: vec![],
+            unordered: HashSet::new(),
+            single: Entity::PLACEHOLDER,
+            not_an_entity: foo.not_an_entity.clone(),
+        };
+
+        // Note: this assumes that the VisitEntities derive is field-ordered,
+        //       which isn't explicitly stated/guaranteed.
+        //       If that changes, this test will fail, but that might be OK if
+        //       we're intentionally breaking that assumption.
+        let mut i = 0;
+        foo.visit_entities(|entity| {
+            let new_entity = entities.reserve_entity();
+            if i < foo.ordered.len() {
+                assert_eq!(entity, foo.ordered[i]);
+                remapped.ordered.push(new_entity);
+            } else if i < foo.ordered.len() + foo.unordered.len() {
+                assert!(foo.unordered.contains(&entity));
+                remapped.unordered.insert(new_entity);
+            } else {
+                assert_eq!(entity, foo.single);
+                remapped.single = new_entity;
+            }
+
+            entity_map.insert(entity, new_entity);
+
+            i += 1;
+        });
+
+        SceneEntityMapper::world_scope(&mut entity_map, &mut world, |_, mapper| {
+            foo.map_entities(mapper);
+        });
+
+        assert_eq!(foo, remapped);
+    }
+}