Make Transform propagation correct in the presence of updated children (#4608)

Supercedes https://github.com/bevyengine/bevy/pull/3340, and absorbs the test from there.

# Objective

- Fixes #3329

## Solution

- If the `Children` component has changed, we currently do not have a way to know how it has changed.
- Therefore, we must update the hierarchy downwards  from that point to be correct.

Co-authored-by: Daniel McNab <36049421+DJMcNab@users.noreply.github.com>

crates/bevy_transform/src/systems.rs

@@ -1,10 +1,5 @@
 use crate::components::{GlobalTransform, Transform};
-use bevy_ecs::{
-    entity::Entity,
-    prelude::Changed,
-    query::{With, Without},
-    system::Query,
-};
+use bevy_ecs::prelude::{Changed, Entity, Query, With, Without};
 use bevy_hierarchy::{Children, Parent};
 
 /// Update [`GlobalTransform`] component of entities based on entity hierarchy and
@@ -12,7 +7,7 @@ use bevy_hierarchy::{Children, Parent};
 pub fn transform_propagate_system(
     mut root_query: Query<
         (
-            Option<&Children>,
+            Option<(&Children, Changed<Children>)>,
             &Transform,
             Changed<Transform>,
             &mut GlobalTransform,
@@ -23,18 +18,19 @@ pub fn transform_propagate_system(
         (&Transform, Changed<Transform>, &mut GlobalTransform),
         With<Parent>,
     >,
-    children_query: Query<Option<&Children>, (With<Parent>, With<GlobalTransform>)>,
+    children_query: Query<(&Children, Changed<Children>), (With<Parent>, With<GlobalTransform>)>,
 ) {
     for (children, transform, transform_changed, mut global_transform) in root_query.iter_mut() {
-        let mut changed = false;
+        let mut changed = transform_changed;
         if transform_changed {
             *global_transform = GlobalTransform::from(*transform);
-            changed = true;
         }
 
-        if let Some(children) = children {
+        if let Some((children, changed_children)) = children {
+            // If our `Children` has changed, we need to recalculate everything below us
+            changed |= changed_children;
             for child in children.iter() {
-                propagate_recursive(
+                let _ = propagate_recursive(
                     &global_transform,
                     &mut transform_query,
                     &children_query,
@@ -52,44 +48,43 @@ fn propagate_recursive(
         (&Transform, Changed<Transform>, &mut GlobalTransform),
         With<Parent>,
     >,
-    children_query: &Query<Option<&Children>, (With<Parent>, With<GlobalTransform>)>,
+    children_query: &Query<(&Children, Changed<Children>), (With<Parent>, With<GlobalTransform>)>,
     entity: Entity,
     mut changed: bool,
-) {
+    // We use a result here to use the `?` operator. Ideally we'd use a try block instead
+) -> Result<(), ()> {
     let global_matrix = {
-        if let Ok((transform, transform_changed, mut global_transform)) =
-            transform_query.get_mut(entity)
-        {
-            changed |= transform_changed;
-            if changed {
-                *global_transform = parent.mul_transform(*transform);
-            }
-            *global_transform
-        } else {
-            return;
+        let (transform, transform_changed, mut global_transform) =
+            transform_query.get_mut(entity).map_err(drop)?;
+
+        changed |= transform_changed;
+        if changed {
+            *global_transform = parent.mul_transform(*transform);
         }
+        *global_transform
     };
 
-    if let Ok(Some(children)) = children_query.get(entity) {
-        for child in children.iter() {
-            propagate_recursive(
-                &global_matrix,
-                transform_query,
-                children_query,
-                *child,
-                changed,
-            );
-        }
+    let (children, changed_children) = children_query.get(entity).map_err(drop)?;
+    // If our `Children` has changed, we need to recalculate everything below us
+    changed |= changed_children;
+    for child in children.iter() {
+        let _ = propagate_recursive(
+            &global_matrix,
+            transform_query,
+            children_query,
+            *child,
+            changed,
+        );
     }
+    Ok(())
 }
 
 #[cfg(test)]
 mod test {
-    use bevy_ecs::{
-        schedule::{Schedule, Stage, SystemStage},
-        system::{CommandQueue, Commands},
-        world::World,
-    };
+    use bevy_app::prelude::*;
+    use bevy_ecs::prelude::*;
+    use bevy_ecs::system::CommandQueue;
+    use bevy_math::vec3;
 
     use crate::components::{GlobalTransform, Transform};
     use crate::systems::transform_propagate_system;
@@ -271,4 +266,60 @@ mod test {
             vec![children[1]]
         );
     }
+
+    #[test]
+    fn correct_transforms_when_no_children() {
+        let mut app = App::new();
+
+        app.add_system(parent_update_system);
+        app.add_system(transform_propagate_system);
+
+        let translation = vec3(1.0, 0.0, 0.0);
+
+        let parent = app
+            .world
+            .spawn()
+            .insert(Transform::from_translation(translation))
+            .insert(GlobalTransform::default())
+            .id();
+
+        let child = app
+            .world
+            .spawn()
+            .insert_bundle((
+                Transform::identity(),
+                GlobalTransform::default(),
+                Parent(parent),
+            ))
+            .id();
+
+        let grandchild = app
+            .world
+            .spawn()
+            .insert_bundle((
+                Transform::identity(),
+                GlobalTransform::default(),
+                Parent(child),
+            ))
+            .id();
+
+        app.update();
+
+        // check the `Children` structure is spawned
+        assert_eq!(&**app.world.get::<Children>(parent).unwrap(), &[child]);
+        assert_eq!(&**app.world.get::<Children>(child).unwrap(), &[grandchild]);
+        // Note that at this point, the `GlobalTransform`s will not have updated yet, due to `Commands` delay
+        app.update();
+
+        let mut state = app.world.query::<&GlobalTransform>();
+        for global in state.iter(&app.world) {
+            assert_eq!(
+                global,
+                &GlobalTransform {
+                    translation,
+                    ..Default::default()
+                },
+            );
+        }
+    }
 }