diff --git a/crates/bevy_ecs/src/reflect.rs b/crates/bevy_ecs/src/reflect.rs index 927973dea2..91781e74fc 100644 --- a/crates/bevy_ecs/src/reflect.rs +++ b/crates/bevy_ecs/src/reflect.rs @@ -413,9 +413,19 @@ impl_from_reflect_value!(Entity); #[derive(Clone)] pub struct ReflectMapEntities { map_entities: fn(&mut World, &EntityMap) -> Result<(), MapEntitiesError>, + map_specific_entities: fn(&mut World, &EntityMap, &[Entity]) -> Result<(), MapEntitiesError>, } impl ReflectMapEntities { + /// A general method for applying [`MapEntities`] behavior to all elements in an [`EntityMap`]. + /// + /// Be mindful in its usage: Works best in situations where the entities in the [`EntityMap`] are newly + /// created, before systems have a chance to add new components. If some of the entities referred to + /// by the [`EntityMap`] might already contain valid entity references, you should use [`map_entities`](Self::map_entities). + /// + /// An example of this: A scene can be loaded with `Parent` components, but then a `Parent` component can be added + /// to these entities after they have been loaded. If you reload the scene using [`map_entities`](Self::map_entities), those `Parent` + /// components with already valid entity references could be updated to point at something else entirely. pub fn map_entities( &self, world: &mut World, @@ -423,11 +433,33 @@ impl ReflectMapEntities { ) -> Result<(), MapEntitiesError> { (self.map_entities)(world, entity_map) } + + /// This is like [`map_entities`](Self::map_entities), but only applied to specific entities, not all values + /// in the [`EntityMap`]. + /// + /// This is useful mostly for when you need to be careful not to update components that already contain valid entity + /// values. See [`map_entities`](Self::map_entities) for more details. + pub fn map_specific_entities( + &self, + world: &mut World, + entity_map: &EntityMap, + entities: &[Entity], + ) -> Result<(), MapEntitiesError> { + (self.map_specific_entities)(world, entity_map, entities) + } } impl FromType for ReflectMapEntities { fn from_type() -> Self { ReflectMapEntities { + map_specific_entities: |world, entity_map, entities| { + for &entity in entities { + if let Some(mut component) = world.get_mut::(entity) { + component.map_entities(entity_map)?; + } + } + Ok(()) + }, map_entities: |world, entity_map| { for entity in entity_map.values() { if let Some(mut component) = world.get_mut::(entity) { diff --git a/crates/bevy_scene/src/dynamic_scene.rs b/crates/bevy_scene/src/dynamic_scene.rs index adac64c425..23a033c969 100644 --- a/crates/bevy_scene/src/dynamic_scene.rs +++ b/crates/bevy_scene/src/dynamic_scene.rs @@ -1,12 +1,16 @@ +use std::any::TypeId; + use crate::{DynamicSceneBuilder, Scene, SceneSpawnError}; use anyhow::Result; use bevy_app::AppTypeRegistry; use bevy_ecs::{ entity::EntityMap, + prelude::Entity, reflect::{ReflectComponent, ReflectMapEntities}, world::World, }; use bevy_reflect::{Reflect, TypeRegistryArc, TypeUuid}; +use bevy_utils::HashMap; #[cfg(feature = "serialize")] use crate::serde::SceneSerializer; @@ -86,6 +90,12 @@ impl DynamicScene { reflect_resource.apply_or_insert(world, &**resource); } + // For each component types that reference other entities, we keep track + // of which entities in the scene use that component. + // This is so we can update the scene-internal references to references + // of the actual entities in the world. + let mut scene_mappings: HashMap> = HashMap::default(); + for scene_entity in &self.entities { // Fetch the entity with the given entity id from the `entity_map` // or spawn a new entity with a transiently unique id if there is @@ -109,6 +119,15 @@ impl DynamicScene { } })?; + // If this component references entities in the scene, track it + // so we can update it to the entity in the world. + if registration.data::().is_some() { + scene_mappings + .entry(registration.type_id()) + .or_insert(Vec::new()) + .push(entity); + } + // If the entity already has the given component attached, // just apply the (possibly) new value, otherwise add the // component to the entity. @@ -116,10 +135,14 @@ impl DynamicScene { } } - for registration in type_registry.iter() { + // Updates references to entities in the scene to entities in the world + for (type_id, entities) in scene_mappings.into_iter() { + let registration = type_registry.get(type_id).expect( + "we should be getting TypeId from this TypeRegistration in the first place", + ); if let Some(map_entities_reflect) = registration.data::() { map_entities_reflect - .map_entities(world, entity_map) + .map_specific_entities(world, entity_map, &entities) .unwrap(); } } @@ -160,3 +183,89 @@ where .new_line("\n".to_string()); ron::ser::to_string_pretty(&serialize, pretty_config) } + +#[cfg(test)] +mod tests { + use bevy_app::AppTypeRegistry; + use bevy_ecs::{entity::EntityMap, system::Command, world::World}; + use bevy_hierarchy::{AddChild, Parent}; + + use crate::dynamic_scene_builder::DynamicSceneBuilder; + + #[test] + fn components_not_defined_in_scene_should_not_be_affected_by_scene_entity_map() { + // Testing that scene reloading applies EntityMap correctly to MapEntities components. + + // First, we create a simple world with a parent and a child relationship + let mut world = World::new(); + world.init_resource::(); + world + .resource_mut::() + .write() + .register::(); + let original_parent_entity = world.spawn_empty().id(); + let original_child_entity = world.spawn_empty().id(); + AddChild { + parent: original_parent_entity, + child: original_child_entity, + } + .write(&mut world); + + // We then write this relationship to a new scene, and then write that scene back to the + // world to create another parent and child relationship + let mut scene_builder = DynamicSceneBuilder::from_world(&world); + scene_builder.extract_entity(original_parent_entity); + scene_builder.extract_entity(original_child_entity); + let scene = scene_builder.build(); + let mut entity_map = EntityMap::default(); + scene.write_to_world(&mut world, &mut entity_map).unwrap(); + + let from_scene_parent_entity = entity_map.get(original_parent_entity).unwrap(); + let from_scene_child_entity = entity_map.get(original_child_entity).unwrap(); + + // We then add the parent from the scene as a child of the original child + // Hierarchy should look like: + // Original Parent <- Original Child <- Scene Parent <- Scene Child + AddChild { + parent: original_child_entity, + child: from_scene_parent_entity, + } + .write(&mut world); + + // We then reload the scene to make sure that from_scene_parent_entity's parent component + // isn't updated with the entity map, since this component isn't defined in the scene. + // With bevy_hierarchy, this can cause serious errors and malformed hierarchies. + scene.write_to_world(&mut world, &mut entity_map).unwrap(); + + assert_eq!( + original_parent_entity, + world + .get_entity(original_child_entity) + .unwrap() + .get::() + .unwrap() + .get(), + "something about reloading the scene is touching entities with the same scene Ids" + ); + assert_eq!( + original_child_entity, + world + .get_entity(from_scene_parent_entity) + .unwrap() + .get::() + .unwrap() + .get(), + "something about reloading the scene is touching components not defined in the scene but on entities defined in the scene" + ); + assert_eq!( + from_scene_parent_entity, + world + .get_entity(from_scene_child_entity) + .unwrap() + .get::() + .expect("something is wrong with this test, and the scene components don't have a parent/child relationship") + .get(), + "something is wrong with the this test or the code reloading scenes since the relationship between scene entities is broken" + ); + } +} diff --git a/crates/bevy_scene/src/scene.rs b/crates/bevy_scene/src/scene.rs index 460f7cf8b2..462c1ec41e 100644 --- a/crates/bevy_scene/src/scene.rs +++ b/crates/bevy_scene/src/scene.rs @@ -117,6 +117,7 @@ impl Scene { } } } + for registration in type_registry.iter() { if let Some(map_entities_reflect) = registration.data::() { map_entities_reflect