From 84c783b1002d6af6ad1da83d6d15d0dee3bce023 Mon Sep 17 00:00:00 2001
From: Daniel McNab <36049421+DJMcNab@users.noreply.github.com>
Date: Mon, 16 May 2022 21:25:34 +0000
Subject: [PATCH] Ensure that the parent is always the expected entity (#4717)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

# Objective

- Transform propogation could stack overflow when there was a cycle.
- I think https://github.com/bevyengine/bevy/pull/4203 would use all available memory.

## Solution

- Make sure that the child entity's `Parent`s are their parents.

This is also required for when parallelising, although as noted in the comment, the naïve solution would be UB.
(The best way to fix this would probably be an `&mut UnsafeCell<T>` `WorldQuery`, or wrapper type with the same effect)
---
 crates/bevy_transform/src/systems.rs | 68 +++++++++++++++++++++++-----
 1 file changed, 57 insertions(+), 11 deletions(-)

diff --git a/crates/bevy_transform/src/systems.rs b/crates/bevy_transform/src/systems.rs
index da063f8241597..e7bd9789d3df7 100644
--- a/crates/bevy_transform/src/systems.rs
+++ b/crates/bevy_transform/src/systems.rs
@@ -11,16 +11,21 @@ pub fn transform_propagate_system(
             &Transform,
             Changed<Transform>,
             &mut GlobalTransform,
+            Entity,
         ),
         Without<Parent>,
     >,
-    mut transform_query: Query<
-        (&Transform, Changed<Transform>, &mut GlobalTransform),
-        With<Parent>,
-    >,
+    mut transform_query: Query<(
+        &Transform,
+        Changed<Transform>,
+        &mut GlobalTransform,
+        &Parent,
+    )>,
     children_query: Query<(&Children, Changed<Children>), (With<Parent>, With<GlobalTransform>)>,
 ) {
-    for (children, transform, transform_changed, mut global_transform) in root_query.iter_mut() {
+    for (children, transform, transform_changed, mut global_transform, entity) in
+        root_query.iter_mut()
+    {
         let mut changed = transform_changed;
         if transform_changed {
             *global_transform = GlobalTransform::from(*transform);
@@ -35,6 +40,7 @@ pub fn transform_propagate_system(
                     &mut transform_query,
                     &children_query,
                     *child,
+                    entity,
                     changed,
                 );
             }
@@ -44,19 +50,26 @@ pub fn transform_propagate_system(
 
 fn propagate_recursive(
     parent: &GlobalTransform,
-    transform_query: &mut Query<
-        (&Transform, Changed<Transform>, &mut GlobalTransform),
-        With<Parent>,
-    >,
+    transform_query: &mut Query<(
+        &Transform,
+        Changed<Transform>,
+        &mut GlobalTransform,
+        &Parent,
+    )>,
     children_query: &Query<(&Children, Changed<Children>), (With<Parent>, With<GlobalTransform>)>,
     entity: Entity,
+    expected_parent: 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 = {
-        let (transform, transform_changed, mut global_transform) =
+        let (transform, transform_changed, mut global_transform, child_parent) =
             transform_query.get_mut(entity).map_err(drop)?;
-
+        // Note that for parallelising, this check cannot occur here, since there is an `&mut GlobalTransform` (in global_transform)
+        assert_eq!(
+            child_parent.0, expected_parent,
+            "Malformed hierarchy. This probably means that your hierarchy has been improperly maintained, or contains a cycle"
+        );
         changed |= transform_changed;
         if changed {
             *global_transform = parent.mul_transform(*transform);
@@ -73,6 +86,7 @@ fn propagate_recursive(
             transform_query,
             children_query,
             *child,
+            entity,
             changed,
         );
     }
@@ -322,4 +336,36 @@ mod test {
             );
         }
     }
+    #[test]
+    #[should_panic]
+    fn panic_when_hierarchy_cycle() {
+        let mut app = App::new();
+
+        app.add_system(parent_update_system);
+        app.add_system(transform_propagate_system);
+
+        let child = app
+            .world
+            .spawn()
+            .insert_bundle((Transform::identity(), GlobalTransform::default()))
+            .id();
+
+        let grandchild = app
+            .world
+            .spawn()
+            .insert_bundle((
+                Transform::identity(),
+                GlobalTransform::default(),
+                Parent(child),
+            ))
+            .id();
+        app.world.spawn().insert_bundle((
+            Transform::default(),
+            GlobalTransform::default(),
+            Children::with(&[child]),
+        ));
+        app.world.entity_mut(child).insert(Parent(grandchild));
+
+        app.update();
+    }
 }