make ReflectMapEntities match MapEntities

Author: jrobsonchase

crates/bevy_animation/src/lib.rs

@@ -19,7 +19,7 @@ use std::{
     collections::BTreeMap,
     fmt::Debug,
     hash::{Hash, Hasher},
-    iter,
+    iter, option,
 };
 
 use bevy_app::{App, Plugin, PostUpdate};
@@ -1291,9 +1291,19 @@ 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<'a> IntoIterator for &'a mut AnimationTarget {
+    type IntoIter = option::IntoIter<&'a mut Entity>;
+    type Item = <Self::IntoIter as Iterator>::Item;
+    fn into_iter(self) -> Self::IntoIter {
+        Some(&mut self.player).into_iter()
+    }
+}
+
+impl<'a> IntoIterator for &'a AnimationTarget {
+    type IntoIter = option::IntoIter<Entity>;
+    type Item = <Self::IntoIter as Iterator>::Item;
+    fn into_iter(self) -> Self::IntoIter {
+        Some(self.player).into_iter()
     }
 }
 

crates/bevy_ecs/src/entity/map_entities.rs

@@ -38,11 +38,56 @@ use super::EntityHashMap;
 /// }
 /// ```
 pub trait MapEntities {
-    /// Updates all [`Entity`] references stored inside using `entity_mapper`.
-    ///
-    /// Implementors should look up any and all [`Entity`] values stored within `self` and
-    /// update them to the mapped values via `entity_mapper`.
-    fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M);
+    fn map_entities<F: FnMut(Entity)>(&self, f: F);
+}
+
+pub trait MapEntitiesMut: MapEntities {
+    fn map_entities_mut<F: FnMut(&mut Entity)>(&mut self, f: F);
+}
+
+impl<T> MapEntities for T
+where
+    T: IterEntities,
+{
+    fn map_entities<F: FnMut(Entity)>(&self, f: F) {
+        self.iter_entities().for_each(f)
+    }
+}
+
+impl<T> MapEntitiesMut for T
+where
+    T: MapEntities + IterEntitiesMut,
+{
+    fn map_entities_mut<F: FnMut(&mut Entity)>(&mut self, f: F) {
+        self.iter_entities_mut().for_each(f)
+    }
+}
+
+pub trait IterEntities {
+    fn iter_entities(&self) -> impl Iterator<Item = Entity>;
+}
+
+impl<T> IterEntities for T
+where
+    for<'a> &'a T: IntoIterator<Item = Entity>,
+{
+    fn iter_entities(&self) -> impl Iterator<Item = Entity> {
+        self.into_iter()
+    }
+}
+
+pub trait IterEntitiesMut: IterEntities {
+    fn iter_entities_mut(&mut self) -> impl Iterator<Item = &mut Entity>;
+}
+
+impl<T> IterEntitiesMut for T
+where
+    for<'a> &'a mut T: IntoIterator<Item = &'a mut Entity>,
+    T: IterEntities,
+{
+    fn iter_entities_mut(&mut self) -> impl Iterator<Item = &mut Entity> {
+        self.into_iter()
+    }
 }
 
 /// An implementor of this trait knows how to map an [`Entity`] into another [`Entity`].

crates/bevy_ecs/src/reflect/map_entities.rs

@@ -1,9 +1,9 @@
 use crate::{
     component::Component,
-    entity::{Entity, EntityHashMap, MapEntities, SceneEntityMapper},
+    entity::{Entity, EntityHashMap, EntityMapper, MapEntities, MapEntitiesMut, SceneEntityMapper},
     world::World,
 };
-use bevy_reflect::FromType;
+use bevy_reflect::{FromReflect, FromType, PartialReflect};
 
 /// For a specific type of component, this maps any fields with values of type [`Entity`] to a new world.
 ///
@@ -15,6 +15,8 @@ use bevy_reflect::FromType;
 pub struct ReflectMapEntities {
     map_all_world_entities: fn(&mut World, &mut SceneEntityMapper),
     map_world_entities: fn(&mut World, &mut SceneEntityMapper, &[Entity]),
+    map_entities_mut: fn(&mut dyn PartialReflect, &mut dyn FnMut(&mut Entity)),
+    map_entities: fn(&dyn PartialReflect, &mut dyn FnMut(Entity)),
 }
 
 impl ReflectMapEntities {
@@ -51,15 +53,29 @@ impl ReflectMapEntities {
             (self.map_world_entities)(world, mapper, entities);
         });
     }
+
+    pub fn map_entities(&self, component: &dyn PartialReflect, f: &mut dyn FnMut(Entity)) {
+        (self.map_entities)(component, f)
+    }
+
+    pub fn map_entities_mut(
+        &self,
+        component: &mut dyn PartialReflect,
+        f: &mut dyn FnMut(&mut Entity),
+    ) {
+        (self.map_entities_mut)(component, f)
+    }
 }
 
-impl<C: Component + MapEntities> FromType<C> for ReflectMapEntities {
+impl<C: Component + FromReflect + MapEntitiesMut> FromType<C> for ReflectMapEntities {
     fn from_type() -> Self {
         ReflectMapEntities {
             map_world_entities: |world, entity_mapper, entities| {
                 for &entity in entities {
                     if let Some(mut component) = world.get_mut::<C>(entity) {
-                        component.map_entities(entity_mapper);
+                        component.map_entities_mut(|entity| {
+                            *entity = entity_mapper.map_entity(*entity);
+                        });
                     }
                 }
             },
@@ -71,10 +87,21 @@ impl<C: Component + MapEntities> FromType<C> for ReflectMapEntities {
                     .collect::<Vec<Entity>>();
                 for entity in &entities {
                     if let Some(mut component) = world.get_mut::<C>(*entity) {
-                        component.map_entities(entity_mapper);
+                        component.map_entities_mut(|entity| {
+                            *entity = entity_mapper.map_entity(*entity);
+                        });
                     }
                 }
             },
+            map_entities: |component, f| {
+                let concrete = C::from_reflect(component).unwrap();
+                concrete.map_entities(f);
+            },
+            map_entities_mut: |component, f| {
+                let mut concrete = C::from_reflect(component).unwrap();
+                concrete.map_entities_mut(f);
+                component.apply(&concrete);
+            },
         }
     }
 }
@@ -99,12 +126,14 @@ impl ReflectMapEntitiesResource {
     }
 }
 
-impl<R: crate::system::Resource + MapEntities> FromType<R> for ReflectMapEntitiesResource {
+impl<R: crate::system::Resource + MapEntitiesMut> FromType<R> for ReflectMapEntitiesResource {
     fn from_type() -> Self {
         ReflectMapEntitiesResource {
             map_entities: |world, entity_mapper| {
                 if let Some(mut resource) = world.get_resource_mut::<R>() {
-                    resource.map_entities(entity_mapper);
+                    resource.map_entities_mut(|entity| {
+                        *entity = entity_mapper.map_entity(*entity);
+                    });
                 }
             },
         }

crates/bevy_hierarchy/src/components/children.rs

@@ -2,7 +2,7 @@
 use bevy_ecs::reflect::{ReflectComponent, ReflectFromWorld, ReflectMapEntities};
 use bevy_ecs::{
     component::Component,
-    entity::{Entity, EntityMapper, MapEntities},
+    entity::{Entity, MapEntitiesMut},
     prelude::FromWorld,
     world::World,
 };
@@ -28,14 +28,6 @@ use std::ops::Deref;
 #[cfg_attr(feature = "reflect", reflect(Component, MapEntities, 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
@@ -152,13 +144,36 @@ impl Deref for Children {
     }
 }
 
+impl<'a> IntoIterator for &'a mut Children {
+    type Item = <Self::IntoIter as Iterator>::Item;
+
+    type IntoIter = slice::IterMut<'a, Entity>;
+
+    #[inline(always)]
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.iter_mut()
+    }
+}
+
 impl<'a> IntoIterator for &'a Children {
     type Item = <Self::IntoIter as Iterator>::Item;
 
-    type IntoIter = slice::Iter<'a, Entity>;
+    type IntoIter = std::iter::Copied<slice::Iter<'a, Entity>>;
 
     #[inline(always)]
     fn into_iter(self) -> Self::IntoIter {
-        self.0.iter()
+        self.0.iter().copied()
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    fn assert_impls_map_entities_mut<M: MapEntitiesMut>() {}
+
+    #[test]
+    fn children_impls_map_entities_mut() {
+        assert_impls_map_entities_mut::<Children>();
     }
 }

crates/bevy_hierarchy/src/components/parent.rs

@@ -2,11 +2,11 @@
 use bevy_ecs::reflect::{ReflectComponent, ReflectFromWorld, ReflectMapEntities};
 use bevy_ecs::{
     component::Component,
-    entity::{Entity, EntityMapper, MapEntities},
+    entity::Entity,
     traversal::Traversal,
     world::{FromWorld, World},
 };
-use std::ops::Deref;
+use std::{ops::Deref, option};
 
 /// Holds a reference to the parent entity of this entity.
 /// This component should only be present on entities that actually have a parent entity.
@@ -59,9 +59,21 @@ 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<'a> IntoIterator for &'a Parent {
+    type Item = Entity;
+    type IntoIter = option::IntoIter<Entity>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        Some(self.0).into_iter()
+    }
+}
+
+impl<'a> IntoIterator for &'a mut Parent {
+    type Item = &'a mut Entity;
+    type IntoIter = option::IntoIter<&'a mut Entity>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        Some(&mut self.0).into_iter()
     }
 }
 

crates/bevy_hierarchy/src/query_extension.rs

@@ -92,12 +92,7 @@ where
     pub fn new(children_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
         DescendantIter {
             children_query,
-            vecdeque: children_query
-                .get(entity)
-                .into_iter()
-                .flatten()
-                .copied()
-                .collect(),
+            vecdeque: children_query.get(entity).into_iter().flatten().collect(),
         }
     }
 }

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

@@ -16,11 +16,21 @@ pub struct SkinnedMesh {
     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);
-        }
+impl<'a> IntoIterator for &'a mut SkinnedMesh {
+    type IntoIter = std::slice::IterMut<'a, Entity>;
+    type Item = <Self::IntoIter as Iterator>::Item;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.joints.iter_mut()
+    }
+}
+
+impl<'a> IntoIterator for &'a SkinnedMesh {
+    type IntoIter = std::iter::Copied<std::slice::Iter<'a, Entity>>;
+    type Item = <Self::IntoIter as Iterator>::Item;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.joints.iter().copied()
     }
 }
 

crates/bevy_render/src/view/visibility/mod.rs

@@ -370,7 +370,7 @@ fn visibility_propagate_system(
             inherited_visibility.0 = is_visible;
 
             // Recursively update the visibility of each child.
-            for &child in children.into_iter().flatten() {
+            for child in children.into_iter().flatten() {
                 let _ =
                     propagate_recursive(is_visible, child, &mut visibility_query, &children_query);
             }
@@ -401,7 +401,7 @@ fn propagate_recursive(
         inherited_visibility.0 = is_visible;
 
         // Recursively update the visibility of each child.
-        for &child in children_query.get(entity).ok().into_iter().flatten() {
+        for child in children_query.get(entity).ok().into_iter().flatten() {
             let _ = propagate_recursive(is_visible, child, visibility_query, children_query);
         }
     }