bring back VisitEntitiesMut

Author: jrobsonchase

crates/bevy_animation/src/lib.rs

@@ -28,9 +28,9 @@ use bevy_app::{App, Plugin, PostUpdate};
 use bevy_asset::{Asset, AssetApp, Assets, Handle};
 use bevy_core::Name;
 use bevy_ecs::{
-    entity::{MapEntities, VisitEntities},
+    entity::{VisitEntities, VisitEntitiesMut},
     prelude::*,
-    reflect::{ReflectMapEntities, ReflectVisitEntities},
+    reflect::{ReflectMapEntities, ReflectVisitEntities, ReflectVisitEntitiesMut},
     world::EntityMutExcept,
 };
 use bevy_math::FloatExt;
@@ -530,8 +530,8 @@ 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, VisitEntities)]
-#[reflect(Component, MapEntities, VisitEntities)]
+#[derive(Clone, Copy, Component, Reflect, VisitEntities, VisitEntitiesMut)]
+#[reflect(Component, MapEntities, VisitEntities, VisitEntitiesMut)]
 pub struct AnimationTarget {
     /// The ID of this animation target.
     ///
@@ -1302,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

@@ -180,6 +180,85 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
     })
 }
 
+#[proc_macro_derive(VisitEntitiesMut, attributes(visit_entities))]
+pub fn derive_visit_entities_mut(input: TokenStream) -> 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(),
+            "Invalid `VisitEntitiesMut` type: at least one field",
+        )
+        .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 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::VisitEntitiesMut for #struct_name #ty_generics {
+            fn visit_entities_mut<F: FnMut(&mut Entity)>(&mut self, mut f: F) {
+                #(#field_access.visit_entities_mut(&mut f);)*
+            }
+        }
+    })
+}
+
 #[proc_macro_derive(VisitEntities, attributes(visit_entities))]
 pub fn derive_visit_entities(input: TokenStream) -> TokenStream {
     let ast = parse_macro_input!(input as DeriveInput);

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

crates/bevy_ecs/src/entity/visit_entities.rs

@@ -1,4 +1,4 @@
-pub use bevy_ecs_macros::VisitEntities;
+pub use bevy_ecs_macros::{VisitEntities, VisitEntitiesMut};
 
 use crate::entity::Entity;
 
@@ -28,14 +28,44 @@ impl VisitEntities for Entity {
     }
 }
 
+/// 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::Entities};
+    use crate::{
+        self as bevy_ecs,
+        entity::{Entities, EntityHashMap, MapEntities, SceneEntityMapper},
+        world::World,
+    };
     use bevy_utils::HashSet;
 
     use super::*;
 
-    #[derive(VisitEntities)]
+    #[derive(VisitEntities, Debug, PartialEq)]
     struct Foo {
         ordered: Vec<Entity>,
         unordered: HashSet<Entity>,
@@ -45,10 +75,29 @@ mod tests {
         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 `MapEntites` 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 entities = Entities::new();
-        let foo = Foo {
+        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(),
@@ -61,21 +110,41 @@ mod tests {
             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);
     }
 }

crates/bevy_ecs/src/lib.rs

@@ -93,7 +93,7 @@ mod tests {
         world::{EntityMut, EntityRef, Mut, World},
     };
     use alloc::sync::Arc;
-    use bevy_ecs_macros::VisitEntities;
+    use bevy_ecs_macros::{VisitEntities, VisitEntitiesMut};
     use bevy_tasks::{ComputeTaskPool, TaskPool};
     use bevy_utils::HashSet;
     use core::{
@@ -2054,7 +2054,7 @@ mod tests {
     }
 
     #[allow(dead_code)]
-    #[derive(Component, VisitEntities)]
+    #[derive(Component, VisitEntities, VisitEntitiesMut)]
     struct MyEntities {
         entities: Vec<Entity>,
         another_one: Entity,
@@ -2064,6 +2064,6 @@ mod tests {
     }
 
     #[allow(dead_code)]
-    #[derive(Component, VisitEntities)]
+    #[derive(Component, VisitEntities, VisitEntitiesMut)]
     struct MyEntitiesTuple(Vec<Entity>, Entity, #[visit_entities(ignore)] usize);
 }

crates/bevy_ecs/src/reflect/mod.rs

@@ -26,7 +26,7 @@ pub use entity_commands::ReflectCommandExt;
 pub use from_world::{ReflectFromWorld, ReflectFromWorldFns};
 pub use map_entities::{ReflectMapEntities, ReflectMapEntitiesResource};
 pub use resource::{ReflectResource, ReflectResourceFns};
-pub use visit_entities::ReflectVisitEntities;
+pub use visit_entities::{ReflectVisitEntities, ReflectVisitEntitiesMut};
 
 /// A [`Resource`] storing [`TypeRegistry`] for
 /// type registrations relevant to a whole app.

crates/bevy_ecs/src/reflect/visit_entities.rs

@@ -1,4 +1,4 @@
-use crate::entity::{Entity, VisitEntities};
+use crate::entity::{Entity, VisitEntities, VisitEntitiesMut};
 use bevy_reflect::{FromReflect, FromType, PartialReflect};
 
 /// For a reflected value, apply an operation to all contained entities.
@@ -27,3 +27,36 @@ impl<C: FromReflect + VisitEntities> FromType<C> for ReflectVisitEntities {
         }
     }
 }
+
+/// For a reflected value, apply an operation to mutable references to all
+/// contained entities.
+///
+/// See [`VisitEntitiesMut`] for more details.
+#[derive(Clone)]
+pub struct ReflectVisitEntitiesMut {
+    visit_entities_mut: fn(&mut dyn PartialReflect, &mut dyn FnMut(&mut Entity)),
+}
+
+impl ReflectVisitEntitiesMut {
+    /// A general method for applying an operation to all entities in a
+    /// reflected component.
+    pub fn visit_entities(
+        &self,
+        component: &mut dyn PartialReflect,
+        f: &mut dyn FnMut(&mut Entity),
+    ) {
+        (self.visit_entities_mut)(component, f);
+    }
+}
+
+impl<C: FromReflect + VisitEntitiesMut> FromType<C> for ReflectVisitEntitiesMut {
+    fn from_type() -> Self {
+        ReflectVisitEntitiesMut {
+            visit_entities_mut: |component, f| {
+                let mut concrete = C::from_reflect(component).unwrap();
+                concrete.visit_entities_mut(f);
+                component.apply(&concrete);
+            },
+        }
+    }
+}

crates/bevy_hierarchy/src/components/children.rs

@@ -1,10 +1,11 @@
 #[cfg(feature = "reflect")]
 use bevy_ecs::reflect::{
     ReflectComponent, ReflectFromWorld, ReflectMapEntities, ReflectVisitEntities,
+    ReflectVisitEntitiesMut,
 };
 use bevy_ecs::{
     component::Component,
-    entity::{Entity, EntityMapper, MapEntities},
+    entity::{Entity, VisitEntitiesMut},
     prelude::FromWorld,
     world::World,
 };
@@ -24,22 +25,21 @@ use smallvec::SmallVec;
 /// [`Query`]: bevy_ecs::system::Query
 /// [`Parent`]: crate::components::parent::Parent
 /// [`BuildChildren::with_children`]: crate::child_builder::BuildChildren::with_children
-#[derive(Component, Debug)]
+#[derive(Component, Debug, VisitEntitiesMut)]
 #[cfg_attr(feature = "reflect", derive(bevy_reflect::Reflect))]
 #[cfg_attr(
     feature = "reflect",
-    reflect(Component, MapEntities, VisitEntities, Debug, FromWorld)
+    reflect(
+        Component,
+        MapEntities,
+        VisitEntities,
+        VisitEntitiesMut,
+        Debug,
+        FromWorld
+    )
 )]
 pub struct Children(pub(crate) SmallVec<[Entity; 8]>);
 
-impl MapEntities for Children {
-    fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M) {
-        for entity in &mut self.0 {
-            *entity = entity_mapper.map_entity(*entity);
-        }
-    }
-}
-
 // TODO: We need to impl either FromWorld or Default so Children can be registered as Reflect.
 // This is because Reflect deserialize by creating an instance and apply a patch on top.
 // However Children should only ever be set with a real user-defined entities. Its worth looking