Fix transform propagation for orphaned hierarchies

Author: nicopap

crates/bevy_transform/src/systems.rs

@@ -2,6 +2,7 @@ use crate::components::{GlobalTransform, Transform};
 use bevy_ecs::{
     change_detection::Ref,
     prelude::{Changed, DetectChanges, Entity, Query, With, Without},
+    system::RemovedComponents,
 };
 use bevy_hierarchy::{Children, Parent};
 
@@ -23,23 +24,28 @@ pub fn sync_simple_transforms(
 /// Update [`GlobalTransform`] component of entities based on entity hierarchy and
 /// [`Transform`] component.
 ///
+/// Note that this depends on the [`World::clear_trackers`] system to be present.
+///
 /// Third party plugins should use [`transform_propagate_system_set`](crate::transform_propagate_system_set)
 /// to propagate transforms correctly.
 pub fn propagate_transforms(
     mut root_query: Query<
         (Entity, &Children, Ref<Transform>, &mut GlobalTransform),
         Without<Parent>,
     >,
+    orphaned: RemovedComponents<Parent>,
     transform_query: Query<(Ref<Transform>, &mut GlobalTransform, Option<&Children>), With<Parent>>,
     parent_query: Query<(Entity, Ref<Parent>)>,
 ) {
+    let mut orphaned = orphaned.iter().collect::<Vec<_>>();
+    orphaned.sort_unstable();
     root_query.par_for_each_mut(
         // The differing depths and sizes of hierarchy trees causes the work for each root to be
         // different. A batch size of 1 ensures that each tree gets it's own task and multiple
         // large trees are not clumped together.
         1,
         |(entity, children, transform, mut global_transform)| {
-            let changed = transform.is_changed();
+            let changed = transform.is_changed() || orphaned.binary_search(&entity).is_ok();
             if changed {
                 *global_transform = GlobalTransform::from(*transform);
             }
@@ -165,21 +171,62 @@ mod test {
     use bevy_tasks::{ComputeTaskPool, TaskPool};
 
     use crate::components::{GlobalTransform, Transform};
-    use crate::systems::*;
-    use crate::TransformBundle;
+    use crate::{transform_propagate_system_set, TransformBundle};
     use bevy_hierarchy::{BuildChildren, BuildWorldChildren, Children, Parent};
 
     #[derive(StageLabel)]
     struct Update;
 
+    #[test]
+    fn correct_parent_removed() {
+        ComputeTaskPool::init(TaskPool::default);
+        let mut world = World::default();
+
+        let mut update_stage = SystemStage::parallel();
+        update_stage.add_system_set(transform_propagate_system_set());
+
+        let mut schedule = Schedule::default();
+        schedule.add_stage(Update, update_stage);
+
+        let mut command_queue = CommandQueue::default();
+        let mut commands = Commands::new(&mut command_queue, &world);
+        let root = commands
+            .spawn(TransformBundle::from_transform(Transform::from_xyz(
+                9.9, 9.9, 9.9,
+            )))
+            .id();
+        let parent = commands.spawn(TransformBundle::IDENTITY).id();
+        let child = commands.spawn(TransformBundle::IDENTITY).id();
+        commands.entity(parent).set_parent(root);
+        commands.entity(child).set_parent(parent);
+        command_queue.apply(&mut world);
+        schedule.run(&mut world);
+
+        assert_ne!(
+            world.get::<GlobalTransform>(parent).unwrap(),
+            &GlobalTransform::from(Transform::IDENTITY)
+        );
+
+        // Remove parent of `parent`
+        let mut command_queue = CommandQueue::default();
+        let mut commands = Commands::new(&mut command_queue, &world);
+        commands.entity(parent).remove_parent();
+        command_queue.apply(&mut world);
+        schedule.run(&mut world);
+
+        assert_eq!(
+            world.get::<GlobalTransform>(parent).unwrap(),
+            &GlobalTransform::from(Transform::IDENTITY)
+        );
+    }
+
     #[test]
     fn did_propagate() {
         ComputeTaskPool::init(TaskPool::default);
         let mut world = World::default();
 
         let mut update_stage = SystemStage::parallel();
-        update_stage.add_system(sync_simple_transforms);
-        update_stage.add_system(propagate_transforms);
+        update_stage.add_system_set(transform_propagate_system_set());
 
         let mut schedule = Schedule::default();
         schedule.add_stage(Update, update_stage);
@@ -220,8 +267,7 @@ mod test {
         let mut world = World::default();
 
         let mut update_stage = SystemStage::parallel();
-        update_stage.add_system(sync_simple_transforms);
-        update_stage.add_system(propagate_transforms);
+        update_stage.add_system_set(transform_propagate_system_set());
 
         let mut schedule = Schedule::default();
         schedule.add_stage(Update, update_stage);
@@ -264,8 +310,7 @@ mod test {
         let mut world = World::default();
 
         let mut update_stage = SystemStage::parallel();
-        update_stage.add_system(sync_simple_transforms);
-        update_stage.add_system(propagate_transforms);
+        update_stage.add_system_set(transform_propagate_system_set());
 
         let mut schedule = Schedule::default();
         schedule.add_stage(Update, update_stage);
@@ -344,8 +389,7 @@ mod test {
         let mut app = App::new();
         ComputeTaskPool::init(TaskPool::default);
 
-        app.add_system(sync_simple_transforms);
-        app.add_system(propagate_transforms);
+        app.add_system_set(transform_propagate_system_set());
 
         let translation = vec3(1.0, 0.0, 0.0);
 
@@ -391,8 +435,7 @@ mod test {
         let mut temp = World::new();
         let mut app = App::new();
 
-        app.add_system(propagate_transforms)
-            .add_system(sync_simple_transforms);
+        app.add_system_set(transform_propagate_system_set());
 
         fn setup_world(world: &mut World) -> (Entity, Entity) {
             let mut grandchild = Entity::from_raw(0);