[Merged by Bors] - Parallelized transform propagation

crates/bevy_transform/Cargo.toml

@@ -17,5 +17,8 @@ bevy_math = { path = "../bevy_math", version = "0.9.0-dev" }
 bevy_reflect = { path = "../bevy_reflect", version = "0.9.0-dev", features = ["bevy"] }
 serde = { version = "1", features = ["derive"], optional = true }
 
+[dev_dependencies]
+bevy_tasks = { path = "../bevy_tasks", version = "0.9.0-dev" }
+
 [features]
 serialize = ["dep:serde", "bevy_math/serialize"]

crates/bevy_transform/src/components/global_transform.rs

@@ -19,8 +19,8 @@ use bevy_reflect::{std_traits::ReflectDefault, FromReflect, Reflect};
 ///
 /// [`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 labeled
+/// [`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

@@ -20,8 +20,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 labeled
+/// [`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 labeled
+/// [`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
@@ -91,11 +90,19 @@ impl Plugin for TransformPlugin {
             // add transform systems to startup so the first update is "correct"
             .add_startup_system_to_stage(
                 StartupStage::PostStartup,
-                systems::transform_propagate_system.label(TransformSystem::TransformPropagate),
+                systems::sync_simple_transforms.label(TransformSystem::TransformPropagate),
+            )
+            .add_startup_system_to_stage(
+                StartupStage::PostStartup,
+                systems::propagate_transforms.label(TransformSystem::TransformPropagate),
+            )
+            .add_system_to_stage(
+                CoreStage::PostUpdate,
+                systems::sync_simple_transforms.label(TransformSystem::TransformPropagate),
             )
             .add_system_to_stage(
                 CoreStage::PostUpdate,
-                systems::transform_propagate_system.label(TransformSystem::TransformPropagate),
+                systems::propagate_transforms.label(TransformSystem::TransformPropagate),
             );
     }
 }

crates/bevy_transform/src/systems.rs

@@ -2,74 +2,93 @@ use crate::components::{GlobalTransform, Transform};
 use bevy_ecs::prelude::{Changed, Entity, Query, With, Without};
 use bevy_hierarchy::{Children, Parent};
 
+/// Update [`GlobalTransform`] component of entities that aren't in the hierarchy
+pub fn sync_simple_transforms(
+    mut query: Query<
+        (&Transform, &mut GlobalTransform),
+        (Changed<Transform>, Without<Parent>, Without<Children>),
+    >,
+) {
+    for (transform, mut global_transform) in 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 fn propagate_transforms(
     mut root_query: Query<
         (
-            Option<(&Children, Changed<Children>)>,
+            Entity,
+            &Children,
             &Transform,
             Changed<Transform>,
+            Changed<Children>,
             &mut GlobalTransform,
-            Entity,
         ),
         Without<Parent>,
     >,
-    mut transform_query: Query<(
-        &Transform,
-        Changed<Transform>,
-        &mut GlobalTransform,
-        &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>)>,
 ) {
-    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);
-        }
+    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)| {
+            if changed {
+                *global_transform = GlobalTransform::from(*transform);
+            }
 
-        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 {
+            changed |= children_changed;
+
+            for child in children.iter() {
                 let _ = propagate_recursive(
                     &global_transform,
-                    &mut transform_query,
+                    &transform_query,
+                    &parent_query,
                     &children_query,
-                    *child,
                     entity,
+                    *child,
                     changed,
                 );
             }
-        }
-    }
+        },
+    );
 }
 
 fn propagate_recursive(
     parent: &GlobalTransform,
-    transform_query: &mut Query<(
-        &Transform,
-        Changed<Transform>,
-        &mut GlobalTransform,
-        &Parent,
-    )>,
+    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>)>,
-    entity: Entity,
     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
 ) -> Result<(), ()> {
-    let global_matrix = {
-        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)
+    if let Ok(actual_parent) = parent_query.get(entity) {
         assert_eq!(
-            child_parent.get(), expected_parent,
+            actual_parent.get(), expected_parent,
             "Malformed hierarchy. This probably means that your hierarchy has been improperly maintained, or contains a cycle"
         );
+    } else {
+        panic!("Propagated child for {:?} has no Parent component!", entity);
+    }
+
+    // SAFE: With the check that each child has one and only one parent, each child must be globally unique within the
+    // hierarchy. Because of this, it is impossible for this query to have aliased mutable access to the same GlobalTransform.
+    // Any malformed hierarchy will cause a panic due to the above check.
+    let global_matrix = unsafe {
+        let (transform, transform_changed, mut global_transform) =
+            unsafe_transform_query.get_unchecked(entity).map_err(drop)?;
+
         changed |= transform_changed;
         if changed {
             *global_transform = parent.mul_transform(*transform);
@@ -83,10 +102,11 @@ fn propagate_recursive(
     for child in children {
         let _ = propagate_recursive(
             &global_matrix,
-            transform_query,
+            unsafe_transform_query,
+            parent_query,
             children_query,
-            *child,
             entity,
+            *child,
             changed,
         );
     }
@@ -99,9 +119,10 @@ mod test {
     use bevy_ecs::prelude::*;
     use bevy_ecs::system::CommandQueue;
     use bevy_math::vec3;
+    use bevy_tasks::{ComputeTaskPool, TaskPool};
 
     use crate::components::{GlobalTransform, Transform};
-    use crate::systems::transform_propagate_system;
+    use crate::systems::*;
     use crate::TransformBundle;
     use bevy_hierarchy::{BuildChildren, BuildWorldChildren, Children, Parent};
 
@@ -110,10 +131,12 @@ mod test {
 
     #[test]
     fn did_propagate() {
+        ComputeTaskPool::init(TaskPool::default);
         let mut world = World::default();
 
         let mut update_stage = SystemStage::parallel();
-        update_stage.add_system(transform_propagate_system);
+        update_stage.add_system(sync_simple_transforms);
+        update_stage.add_system(propagate_transforms);
 
         let mut schedule = Schedule::default();
         schedule.add_stage(Update, update_stage);
@@ -152,8 +175,10 @@ mod test {
     #[test]
     fn did_propagate_command_buffer() {
         let mut world = World::default();
+
         let mut update_stage = SystemStage::parallel();
-        update_stage.add_system(transform_propagate_system);
+        update_stage.add_system(sync_simple_transforms);
+        update_stage.add_system(propagate_transforms);
 
         let mut schedule = Schedule::default();
         schedule.add_stage(Update, update_stage);
@@ -192,10 +217,12 @@ mod test {
 
     #[test]
     fn correct_children() {
+        ComputeTaskPool::init(TaskPool::default);
         let mut world = World::default();
 
         let mut update_stage = SystemStage::parallel();
-        update_stage.add_system(transform_propagate_system);
+        update_stage.add_system(sync_simple_transforms);
+        update_stage.add_system(propagate_transforms);
 
         let mut schedule = Schedule::default();
         schedule.add_stage(Update, update_stage);
@@ -272,8 +299,10 @@ mod test {
     #[test]
     fn correct_transforms_when_no_children() {
         let mut app = App::new();
+        ComputeTaskPool::init(TaskPool::default);
 
-        app.add_system(transform_propagate_system);
+        app.add_system(sync_simple_transforms);
+        app.add_system(propagate_transforms);
 
         let translation = vec3(1.0, 0.0, 0.0);
 
@@ -313,6 +342,7 @@ mod test {
     #[test]
     #[should_panic]
     fn panic_when_hierarchy_cycle() {
+        ComputeTaskPool::init(TaskPool::default);
         // We cannot directly edit Parent and Children, so we use a temp world to break
         // the hierarchy's invariants.
         let mut temp = World::new();
@@ -321,7 +351,8 @@ mod test {
         // FIXME: Parallel executors seem to have some odd interaction with the other
         // tests in this crate. Using single_threaded until a root cause can be found.
         app.add_stage("single", SystemStage::single_threaded())
-            .add_system_to_stage("single", transform_propagate_system);
+            .add_system_to_stage("single", sync_simple_transforms)
+            .add_system_to_stage("single", propagate_transforms);
 
         fn setup_world(world: &mut World) -> (Entity, Entity) {
             let mut grandchild = Entity::from_raw(0);