Iterative Transform Propagation

crates/bevy_transform/src/components/global_transform.rs

@@ -18,8 +18,8 @@ use std::ops::Mul;
 ///
 /// [`GlobalTransform`] is the position of an entity relative to the reference frame.
 ///
-/// [`GlobalTransform`] is updated from [`Transform`] in the system
-/// [`transform_propagate_system`](crate::transform_propagate_system::transform_propagate_system).
+/// [`GlobalTransform`] is updated from [`Transform`] in the systems labelled
+/// [`TransformSystem::TransformPropagate`](crate::TransformSystem::TransformPropagate).
 ///
 /// This system runs in stage [`CoreStage::PostUpdate`](crate::CoreStage::PostUpdate). If you
 /// update the[`Transform`] of an entity in this stage or after, you will notice a 1 frame lag

crates/bevy_transform/src/components/transform.rs

@@ -19,8 +19,8 @@ use std::ops::Mul;
 ///
 /// [`GlobalTransform`] is the position of an entity relative to the reference frame.
 ///
-/// [`GlobalTransform`] is updated from [`Transform`] in the system
-/// [`transform_propagate_system`](crate::transform_propagate_system::transform_propagate_system).
+/// [`GlobalTransform`] is updated from [`Transform`] in the systems labelled
+/// [`TransformSystem::TransformPropagate`](crate::TransformSystem::TransformPropagate).
 ///
 /// This system runs in stage [`CoreStage::PostUpdate`](crate::CoreStage::PostUpdate). If you
 /// update the[`Transform`] of an entity in this stage or after, you will notice a 1 frame lag

crates/bevy_transform/src/lib.rs

@@ -36,8 +36,8 @@ use prelude::{parent_update_system, Children, GlobalTransform, Parent, PreviousP
 ///
 /// [`GlobalTransform`] is the position of an entity relative to the reference frame.
 ///
-/// [`GlobalTransform`] is updated from [`Transform`] in the system
-/// [`transform_propagate_system`](crate::transform_propagate_system::transform_propagate_system).
+/// [`GlobalTransform`] is updated from [`Transform`] in the systems labelled
+/// [`TransformSystem::TransformPropagate`](crate::TransformSystem::TransformPropagate).
 ///
 /// This system runs in stage [`CoreStage::PostUpdate`](crate::CoreStage::PostUpdate). If you
 /// update the[`Transform`] of an entity in this stage or after, you will notice a 1 frame lag
@@ -112,6 +112,12 @@ impl Plugin for TransformPlugin {
                     .label(TransformSystem::TransformPropagate)
                     .after(TransformSystem::ParentUpdate),
             )
+            .add_startup_system_to_stage(
+                StartupStage::PostStartup,
+                transform_propagate_system::transform_propagate_flat_system
+                    .label(TransformSystem::TransformPropagate)
+                    .after(TransformSystem::ParentUpdate),
+            )
             .add_system_to_stage(
                 CoreStage::PostUpdate,
                 parent_update_system.label(TransformSystem::ParentUpdate),
@@ -121,6 +127,12 @@ impl Plugin for TransformPlugin {
                 transform_propagate_system::transform_propagate_system
                     .label(TransformSystem::TransformPropagate)
                     .after(TransformSystem::ParentUpdate),
+            )
+            .add_system_to_stage(
+                CoreStage::PostUpdate,
+                transform_propagate_system::transform_propagate_flat_system
+                    .label(TransformSystem::TransformPropagate)
+                    .after(TransformSystem::ParentUpdate),
             );
     }
 }

crates/bevy_transform/src/transform_propagate_system.rs

@@ -2,75 +2,97 @@ use crate::components::{Children, GlobalTransform, Parent, Transform};
 use bevy_ecs::{
     entity::Entity,
     query::{Changed, With, Without},
-    system::Query,
+    system::{Local, Query},
 };
 
+/// Used for [`transform_propagate_system`]. Ignore otherwise.
+pub(crate) struct Pending {
+    parent: *const GlobalTransform,
+    changed: bool,
+    child: Entity,
+}
+
+// SAFE: Values are cleared after every frame and cannot otherwise be
+// constructed without transmute.
+unsafe impl Send for Pending {}
+// SAFE: Values are cleared after every frame and cannot otherwise be
+// constructed without transmute.
+unsafe impl Sync for Pending {}
+
+/// Update [`GlobalTransform`] component of entities based on entity hierarchy and
+/// [`Transform`] component.
+pub(crate) fn transform_propagate_flat_system(
+    mut root_query: Query<
+        (&Transform, &mut GlobalTransform),
+        (Without<Parent>, Without<Children>, Changed<Transform>),
+    >,
+) {
+    for (transform, mut global_transform) in root_query.iter_mut() {
+        *global_transform = GlobalTransform::from(*transform);
+    }
+}
+
 /// Update [`GlobalTransform`] component of entities based on entity hierarchy and
 /// [`Transform`] component.
-pub fn transform_propagate_system(
+pub(crate) fn transform_propagate_system(
     mut root_query: Query<
-        (Entity, Option<&Children>, &Transform, &mut GlobalTransform),
+        (
+            &Transform,
+            Changed<Transform>,
+            &Children,
+            &mut GlobalTransform,
+        ),
         Without<Parent>,
     >,
-    mut transform_query: Query<(&Transform, &mut GlobalTransform), With<Parent>>,
-    changed_transform_query: Query<Entity, Changed<Transform>>,
-    children_query: Query<Option<&Children>, (With<Parent>, With<GlobalTransform>)>,
+    mut transform_query: Query<
+        (
+            &Transform,
+            Changed<Transform>,
+            Option<&Children>,
+            &mut GlobalTransform,
+        ),
+        With<Parent>,
+    >,
+    mut pending: Local<Vec<Pending>>,
 ) {
-    for (entity, children, transform, mut global_transform) in root_query.iter_mut() {
+    for (transform, is_changed, children, mut global_transform) in root_query.iter_mut() {
         let mut changed = false;
-        if changed_transform_query.get(entity).is_ok() {
+        if is_changed {
             *global_transform = GlobalTransform::from(*transform);
             changed = true;
         }
 
-        if let Some(children) = children {
-            for child in children.0.iter() {
-                propagate_recursive(
-                    &global_transform,
-                    &changed_transform_query,
-                    &mut transform_query,
-                    &children_query,
-                    *child,
-                    changed,
-                );
+        pending.extend(children.0.iter().map(|child| Pending {
+            parent: &*global_transform as *const GlobalTransform,
+            changed,
+            child: *child,
+        }));
+
+        while let Some(mut current) = pending.pop() {
+            if let Ok((transform, changed, children, mut global_transform)) =
+                transform_query.get_mut(current.child)
+            {
+                current.changed |= changed;
+                // SAFE: The pointers here are generated only during this one traversal
+                // from one given run of the system.
+                let global_matrix = unsafe {
+                    let parent = current.parent.as_ref().unwrap();
+                    if current.changed {
+                        *global_transform = parent.mul_transform(*transform);
+                    }
+                    &*global_transform as *const GlobalTransform
+                };
+                if let Some(children) = children {
+                    pending.extend(children.0.iter().map(|child| Pending {
+                        parent: global_matrix,
+                        changed: current.changed,
+                        child: *child,
+                    }));
+                }
             }
         }
     }
-}
-
-fn propagate_recursive(
-    parent: &GlobalTransform,
-    changed_transform_query: &Query<Entity, Changed<Transform>>,
-    transform_query: &mut Query<(&Transform, &mut GlobalTransform), With<Parent>>,
-    children_query: &Query<Option<&Children>, (With<Parent>, With<GlobalTransform>)>,
-    entity: Entity,
-    mut changed: bool,
-) {
-    changed |= changed_transform_query.get(entity).is_ok();
-
-    let global_matrix = {
-        if let Ok((transform, mut global_transform)) = transform_query.get_mut(entity) {
-            if changed {
-                *global_transform = parent.mul_transform(*transform);
-            }
-            *global_transform
-        } else {
-            return;
-        }
-    };
-
-    if let Ok(Some(children)) = children_query.get(entity) {
-        for child in children.0.iter() {
-            propagate_recursive(
-                &global_matrix,
-                changed_transform_query,
-                transform_query,
-                children_query,
-                *child,
-                changed,
-            );
-        }
-    }
+    debug_assert!(pending.is_empty());
 }
 
 #[cfg(test)]