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,6 +28,32 @@ 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};

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

crates/bevy_hierarchy/src/components/parent.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, VisitEntities},
+    entity::{Entity, VisitEntities, VisitEntitiesMut},
     traversal::Traversal,
     world::{FromWorld, World},
 };
@@ -23,11 +24,19 @@ use core::ops::Deref;
 /// [`Query`]: bevy_ecs::system::Query
 /// [`Children`]: super::children::Children
 /// [`BuildChildren::with_children`]: crate::child_builder::BuildChildren::with_children
-#[derive(Component, Debug, Eq, PartialEq, VisitEntities)]
+#[derive(Component, Debug, Eq, PartialEq, VisitEntities, VisitEntitiesMut)]
 #[cfg_attr(feature = "reflect", derive(bevy_reflect::Reflect))]
 #[cfg_attr(
     feature = "reflect",
-    reflect(Component, MapEntities, VisitEntities, PartialEq, Debug, FromWorld)
+    reflect(
+        Component,
+        MapEntities,
+        VisitEntities,
+        VisitEntitiesMut,
+        PartialEq,
+        Debug,
+        FromWorld
+    )
 )]
 pub struct Parent(pub(crate) Entity);
 
@@ -61,12 +70,6 @@ impl FromWorld for Parent {
     }
 }
 
-impl MapEntities for Parent {
-    fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M) {
-        self.0 = entity_mapper.map_entity(self.0);
-    }
-}
-
 impl Deref for Parent {
     type Target = Entity;
 

crates/bevy_render/src/mesh/mesh/skinning.rs

@@ -1,30 +1,29 @@
 use bevy_asset::{Asset, Handle};
 use bevy_ecs::{
     component::Component,
-    entity::{Entity, EntityMapper, MapEntities, VisitEntities},
+    entity::{Entity, VisitEntities, VisitEntitiesMut},
     prelude::ReflectComponent,
-    reflect::{ReflectMapEntities, ReflectVisitEntities},
+    reflect::{ReflectMapEntities, ReflectVisitEntities, ReflectVisitEntitiesMut},
 };
 use bevy_math::Mat4;
 use bevy_reflect::prelude::*;
 use core::ops::Deref;
 
-#[derive(Component, Debug, Default, Clone, Reflect, VisitEntities)]
-#[reflect(Component, MapEntities, VisitEntities, Default, Debug)]
+#[derive(Component, Debug, Default, Clone, Reflect, VisitEntities, VisitEntitiesMut)]
+#[reflect(
+    Component,
+    MapEntities,
+    VisitEntities,
+    VisitEntitiesMut,
+    Default,
+    Debug
+)]
 pub struct SkinnedMesh {
     #[visit_entities(ignore)]
     pub inverse_bindposes: Handle<SkinnedMeshInverseBindposes>,
     pub joints: Vec<Entity>,
 }
 
-impl MapEntities for SkinnedMesh {
-    fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M) {
-        for joint in &mut self.joints {
-            *joint = entity_mapper.map_entity(*joint);
-        }
-    }
-}
-
 #[derive(Asset, TypePath, Debug)]
 pub struct SkinnedMeshInverseBindposes(Box<[Mat4]>);
 

crates/bevy_scene/src/serde.rs

@@ -515,7 +515,7 @@ mod tests {
         DynamicScene, DynamicSceneBuilder,
     };
     use bevy_ecs::{
-        entity::{Entity, EntityHashMap, EntityMapper, MapEntities, VisitEntities},
+        entity::{Entity, EntityHashMap, VisitEntities, VisitEntitiesMut},
         prelude::{Component, ReflectComponent, ReflectResource, Resource, World},
         query::{With, Without},
         reflect::{AppTypeRegistry, ReflectMapEntities},
@@ -584,16 +584,10 @@ mod tests {
         foo: i32,
     }
 
-    #[derive(Clone, Component, Reflect, PartialEq, VisitEntities)]
+    #[derive(Clone, Component, Reflect, PartialEq, VisitEntities, VisitEntitiesMut)]
     #[reflect(Component, MapEntities, PartialEq)]
     struct MyEntityRef(Entity);
 
-    impl MapEntities for MyEntityRef {
-        fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M) {
-            self.0 = entity_mapper.map_entity(self.0);
-        }
-    }
-
     impl FromWorld for MyEntityRef {
         fn from_world(_world: &mut World) -> Self {
             Self(Entity::PLACEHOLDER)