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).
+/// [`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).
+/// [`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

@@ -4,7 +4,6 @@
 /// The basic components of the transform crate
 pub mod components;
 mod systems;
-pub use crate::systems::transform_propagate_system;
 
 #[doc(hidden)]
 pub mod prelude {
@@ -32,8 +31,8 @@ use prelude::{GlobalTransform, Transform};
 ///
 /// [`GlobalTransform`] is the position of an entity relative to the reference frame.
 ///
-/// [`GlobalTransform`] is updated from [`Transform`] in the 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
@@ -99,11 +98,23 @@ impl Plugin for TransformPlugin {
                     .label(TransformSystem::TransformPropagate)
                     .after(HierarchySystem::ParentUpdate),
             )
+            .add_startup_system_to_stage(
+                StartupStage::PostStartup,
+                systems::transform_propagate_flat_system
+                    .label(TransformSystem::TransformPropagate)
+                    .after(HierarchySystem::ParentUpdate),
+            )
             .add_system_to_stage(
                 CoreStage::PostUpdate,
                 systems::transform_propagate_system
                     .label(TransformSystem::TransformPropagate)
                     .after(HierarchySystem::ParentUpdate),
+            )
+            .add_system_to_stage(
+                CoreStage::PostUpdate,
+                systems::transform_propagate_flat_system
+                    .label(TransformSystem::TransformPropagate)
+                    .after(HierarchySystem::ParentUpdate),
             );
     }
 }

crates/bevy_transform/src/systems.rs

@@ -2,76 +2,96 @@ use crate::components::{GlobalTransform, Transform};
 use bevy_ecs::{
     entity::Entity,
     query::{Changed, With, Without},
-    system::Query,
+    system::{Local, Query},
 };
 use bevy_hierarchy::{Children, Parent};
 
+/// 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 fn transform_propagate_system(
+pub(crate) fn transform_propagate_flat_system(
     mut root_query: Query<
-        (Entity, Option<&Children>, &Transform, &mut GlobalTransform),
-        Without<Parent>,
+        (&Transform, &mut GlobalTransform),
+        (Without<Parent>, Without<Children>, Changed<Transform>),
     >,
-    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>)>,
 ) {
-    for (entity, children, transform, mut global_transform) in root_query.iter_mut() {
-        let mut changed = false;
-        if changed_transform_query.get(entity).is_ok() {
-            *global_transform = GlobalTransform::from(*transform);
-            changed = true;
-        }
-
-        if let Some(children) = children {
-            for child in children.iter() {
-                propagate_recursive(
-                    &global_transform,
-                    &changed_transform_query,
-                    &mut transform_query,
-                    &children_query,
-                    *child,
-                    changed,
-                );
-            }
-        }
-    }
+    root_query.for_each_mut(|(transform, mut global_transform)| {
+        *global_transform = GlobalTransform::from(*transform);
+    });
 }
 
-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,
+/// Update [`GlobalTransform`] component of entities based on entity hierarchy and
+/// [`Transform`] component.
+pub(crate) fn transform_propagate_system(
+    mut root_query: Query<
+        (
+            &Transform,
+            Changed<Transform>,
+            &Children,
+            &mut GlobalTransform,
+        ),
+        Without<Parent>,
+    >,
+    mut transform_query: Query<
+        (
+            &Transform,
+            Changed<Transform>,
+            Option<&Children>,
+            &mut GlobalTransform,
+        ),
+        With<Parent>,
+    >,
+    // Stack space for the depth-first search of a given hierarchy. Used as a Local to
+    // avoid reallocating the stack space used here.
+    mut pending: Local<Vec<Pending>>,
 ) {
-    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;
+    root_query.for_each_mut(|(transform, changed, children, mut global_transform)| {
+        if changed {
+            *global_transform = GlobalTransform::from(*transform);
         }
-    };
 
-    if let Ok(Some(children)) = children_query.get(entity) {
-        for child in children.iter() {
-            propagate_recursive(
-                &global_matrix,
-                changed_transform_query,
-                transform_query,
-                children_query,
-                *child,
-                changed,
-            );
+        pending.extend(children.iter().map(|child| Pending {
+            parent: &*global_transform as *const GlobalTransform,
+            changed,
+            child: *child,
+        }));
+
+        while let Some(current) = pending.pop() {
+            if let Ok((transform, mut changed, children, mut global_transform)) =
+                transform_query.get_mut(current.child)
+            {
+                changed |= current.changed;
+                if changed {
+                    // SAFE: The pointers here are generated only during this one traversal
+                    // from one given run of the system.
+                    unsafe {
+                        *global_transform = current.parent.read().mul_transform(*transform);
+                    };
+                }
+                if let Some(children) = children {
+                    pending.extend(children.iter().map(|child| Pending {
+                        parent: &*global_transform as *const GlobalTransform,
+                        changed: current.changed,
+                        child: *child,
+                    }));
+                }
+            }
         }
-    }
+    });
+    debug_assert!(pending.is_empty());
 }
 
 #[cfg(test)]