Fix transform propagation for orphaned hierarchies

Co-authored-by: vyb <[email protected]>

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::{ParamSet, RemovedComponents},
 };
 use bevy_hierarchy::{Children, Parent};
 
@@ -10,16 +11,29 @@ use bevy_hierarchy::{Children, Parent};
 /// Third party plugins should use [`transform_propagate_system_set`](crate::transform_propagate_system_set)
 /// to propagate transforms correctly.
 pub fn sync_simple_transforms(
-    mut query: Query<
-        (&Transform, &mut GlobalTransform),
-        (Changed<Transform>, Without<Parent>, Without<Children>),
-    >,
+    mut query: ParamSet<(
+        Query<
+            (&Transform, &mut GlobalTransform),
+            (Changed<Transform>, Without<Parent>, Without<Children>),
+        >,
+        Query<(Ref<Transform>, &mut GlobalTransform), Without<Children>>,
+    )>,
+    orphaned: RemovedComponents<Parent>,
 ) {
     query
+        .p0()
         .par_iter_mut()
         .for_each_mut(|(transform, mut global_transform)| {
             *global_transform = GlobalTransform::from(*transform);
         });
+    // updated orphaned entities
+    let mut query = query.p1();
+    let mut iter = query.iter_many_mut(orphaned.iter());
+    while let Some((transform, mut global_transform)) = iter.fetch_next() {
+        if !transform.is_changed() {
+            *global_transform = GlobalTransform::from(*transform);
+        }
+    }
 }
 
 /// Update [`GlobalTransform`] component of entities based on entity hierarchy and
@@ -32,12 +46,15 @@ pub fn propagate_transforms(
         (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_iter_mut().for_each_mut(
         |(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);
             }
@@ -163,21 +180,74 @@ 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 offset_global_transform =
+            |offset| GlobalTransform::from(Transform::from_xyz(offset, offset, offset));
+        let offset_transform =
+            |offset| TransformBundle::from_transform(Transform::from_xyz(offset, offset, offset));
+
+        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(offset_transform(3.3)).id();
+        let parent = commands.spawn(offset_transform(4.4)).id();
+        let child = commands.spawn(offset_transform(5.5)).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(),
+            &offset_global_transform(4.4)
+        );
+
+        // Remove parent of `child`
+        let mut command_queue = CommandQueue::default();
+        let mut commands = Commands::new(&mut command_queue, &world);
+        commands.entity(child).remove_parent();
+        command_queue.apply(&mut world);
+        schedule.run(&mut world);
+
+        assert_eq!(
+            world.get::<GlobalTransform>(child).unwrap(),
+            &offset_global_transform(5.5)
+        );
+    }
+
     #[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);
@@ -218,8 +288,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);
@@ -262,8 +331,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);
@@ -342,8 +410,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);
 
@@ -389,8 +456,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);