Iterative Transform Propagation

crates/bevy_transform/Cargo.toml

@@ -15,6 +15,7 @@ bevy_ecs = { path = "../bevy_ecs", version = "0.9.0", features = ["bevy_reflect"
 bevy_hierarchy = { path = "../bevy_hierarchy", version = "0.9.0" }
 bevy_math = { path = "../bevy_math", version = "0.9.0" }
 bevy_reflect = { path = "../bevy_reflect", version = "0.9.0", features = ["bevy"] }
+thread_local = "1.1"
 serde = { version = "1", features = ["derive"], optional = true }
 
 [dev_dependencies]

crates/bevy_transform/src/systems.rs

@@ -1,6 +1,28 @@
 use crate::components::{GlobalTransform, Transform};
-use bevy_ecs::prelude::{Changed, Entity, Query, With, Without};
+use bevy_ecs::{
+    entity::Entity,
+    query::{Changed, Or, With, Without},
+    system::{Local, Query},
+};
 use bevy_hierarchy::{Children, Parent};
+use core::cell::Cell;
+use thread_local::ThreadLocal;
+
+pub(crate) struct Pending {
+    parent: *const GlobalTransform,
+    changed: bool,
+    parent_entity: Entity,
+    child: Entity,
+}
+
+// SAFE: Access to the parent pointer is only usable in this module, the values
+// are cleared after every system execution, and the system only uses one
+// thread. There is no way to move this type across multiple threads.
+unsafe impl Send for Pending {}
+// SAFE: Access to the parent pointer is only usable in this module, the values
+// are cleared after every system execution, and the system only uses one
+// thread. There is no way to access this type across multiple threads.
+unsafe impl Sync for Pending {}
 
 /// Update [`GlobalTransform`] component of entities that aren't in the hierarchy
 pub fn sync_simple_transforms(
@@ -16,126 +38,114 @@ pub fn sync_simple_transforms(
 
 /// Update [`GlobalTransform`] component of entities based on entity hierarchy and
 /// [`Transform`] component.
-pub fn propagate_transforms(
+pub(crate) fn propagate_transforms(
     mut root_query: Query<
         (
             Entity,
-            &Children,
             &Transform,
-            Changed<Transform>,
-            Changed<Children>,
+            &Children,
+            Or<(Changed<Transform>, Changed<Children>)>,
             &mut GlobalTransform,
         ),
         Without<Parent>,
     >,
-    transform_query: Query<(&Transform, Changed<Transform>, &mut GlobalTransform), With<Parent>>,
     parent_query: Query<&Parent>,
-    children_query: Query<(&Children, Changed<Children>), (With<Parent>, With<GlobalTransform>)>,
+    transform_query: Query<
+        (
+            &Transform,
+            Or<(Changed<Transform>, Changed<Children>)>,
+            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.
+    pending_queues: Local<ThreadLocal<Cell<Vec<Pending>>>>,
 ) {
     root_query.par_for_each_mut(
         // The differing depths and sizes of hierarchy trees causes the work for each root to be
         // different. A batch size of 1 ensures that each tree gets it's own task and multiple
         // large trees are not clumped together.
         1,
-        |(entity, children, transform, mut changed, children_changed, mut global_transform)| {
+        |(root, transform, children, changed,  mut global_transform)| {
+            let pending_cell = pending_queues.get_or_default();
+            let mut pending = pending_cell.take();
+
             if changed {
                 *global_transform = GlobalTransform::from(*transform);
             }
 
-            // If our `Children` has changed, we need to recalculate everything below us
-            changed |= children_changed;
-
-            for child in children.iter() {
-                propagate_recursive(
-                    &global_transform,
-                    &transform_query,
-                    &parent_query,
-                    &children_query,
-                    entity,
-                    *child,
-                    changed,
+            pending.extend(children.iter().map(|child| Pending {
+                parent: &*global_transform as *const GlobalTransform,
+                changed,
+                parent_entity: root,
+                child: *child,
+            }));
+
+            while let Some(current) = pending.pop() {
+                let Ok(actual_parent) = parent_query.get(current.child) else {
+                    panic!("Propagated child for {:?} has no Parent component!", current.child);
+                };
+                assert_eq!(
+                    actual_parent.get(), current.parent_entity,
+                    "Malformed hierarchy. This probably means that your hierarchy has been improperly maintained, or contains a cycle"
                 );
+
+                // SAFETY: This call cannot create aliased mutable references.
+                //   - The top level iteration parallelizes on the roots of the hierarchy.
+                //   - The above assertion ensures that each child has one and only one unique parent throughout the entire
+                //     hierarchy.
+                //
+                // For example, consider the following malformed hierarchy:
+                //
+                //     A
+                //   /   \
+                //  B     C
+                //   \   /
+                //     D
+                //
+                // D has two parents, B and C. If the propagation passes through C, but the Parent component on D points to B,
+                // the above check will panic as the origin parent does match the recorded parent.
+                //
+                // Also consider the following case, where A and B are roots:
+                //
+                //  A       B
+                //   \     /
+                //    C   D
+                //     \ /
+                //      E
+                //
+                // Even if these A and B start two separate tasks running in parallel, one of them will panic before attempting
+                // to mutably access E.
+                let fetch = unsafe { transform_query.get_unchecked(current.child) };
+
+                // If our `Children` has changed, we need to recalculate everything below us
+                let Ok((transform, mut changed, children, mut global_transform)) = fetch else {
+                    continue;
+                };
+
+                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,
+                        parent_entity: current.child,
+                        child: *child,
+                    }));
+                }
             }
-        },
-    );
-}
 
-fn propagate_recursive(
-    parent: &GlobalTransform,
-    unsafe_transform_query: &Query<
-        (&Transform, Changed<Transform>, &mut GlobalTransform),
-        With<Parent>,
-    >,
-    parent_query: &Query<&Parent>,
-    children_query: &Query<(&Children, Changed<Children>), (With<Parent>, With<GlobalTransform>)>,
-    expected_parent: Entity,
-    entity: Entity,
-    mut changed: bool,
-    // We use a result here to use the `?` operator. Ideally we'd use a try block instead
-) {
-    let Ok(actual_parent) = parent_query.get(entity) else {
-        panic!("Propagated child for {:?} has no Parent component!", entity);
-    };
-    assert_eq!(
-        actual_parent.get(), expected_parent,
-        "Malformed hierarchy. This probably means that your hierarchy has been improperly maintained, or contains a cycle"
-    );
-
-    let global_matrix = {
-        let Ok((transform, transform_changed, mut global_transform)) =
-            // SAFETY: This call cannot create aliased mutable references.
-            //   - The top level iteration parallelizes on the roots of the hierarchy.
-            //   - The above assertion ensures that each child has one and only one unique parent throughout the entire
-            //     hierarchy.
-            //
-            // For example, consider the following malformed hierarchy:
-            //
-            //     A
-            //   /   \
-            //  B     C
-            //   \   /
-            //     D
-            //
-            // D has two parents, B and C. If the propagation passes through C, but the Parent component on D points to B,
-            // the above check will panic as the origin parent does match the recorded parent.
-            //
-            // Also consider the following case, where A and B are roots:
-            //
-            //  A       B
-            //   \     /
-            //    C   D
-            //     \ /
-            //      E
-            //
-            // Even if these A and B start two separate tasks running in parallel, one of them will panic before attempting
-            // to mutably access E.
-            (unsafe { unsafe_transform_query.get_unchecked(entity) }) else {
-                return;
-            };
-
-        changed |= transform_changed;
-        if changed {
-            *global_transform = parent.mul_transform(*transform);
-        }
-        *global_transform
-    };
-
-    let Ok((children, changed_children)) = children_query.get(entity) else {
-        return
-    };
-    // If our `Children` has changed, we need to recalculate everything below us
-    changed |= changed_children;
-    for child in children {
-        propagate_recursive(
-            &global_matrix,
-            unsafe_transform_query,
-            parent_query,
-            children_query,
-            entity,
-            *child,
-            changed,
-        );
-    }
+            debug_assert!(pending.is_empty());
+
+        });
 }
 
 #[cfg(test)]