add Transform2d to propagation

NthTensor · NthTensor · commit 358f2b8de8ce · 2025-05-26T19:36:07.000-04:00
diff --git a/crates/bevy_transform/src/components/mod.rs b/crates/bevy_transform/src/components/mod.rs
@@ -1,7 +1,9 @@
 mod global_transform;
 mod transform_2d;
 mod transform_3d;
+mod tree_changed;
 
 pub use global_transform::*;
 pub use transform_2d::*;
 pub use transform_3d::*;
+pub(crate) use tree_changed::TransformTreeChanged;
diff --git a/crates/bevy_transform/src/components/transform_2d.rs b/crates/bevy_transform/src/components/transform_2d.rs
@@ -1,10 +1,30 @@
+use super::{GlobalTransform, TransformTreeChanged};
+use bevy_math::{Affine2, Affine3A, Dir2, Isometry2d, Mat3, Quat, Rot2, Vec2, Vec3};
 use core::ops::Mul;
 
-use bevy_math::{Affine2, Affine3A, Dir2, Isometry2d, Mat3, Quat, Rot2, Vec2, Vec3};
+#[cfg(feature = "bevy-support")]
+use bevy_ecs::component::Component;
 
-use super::GlobalTransform;
+#[cfg(feature = "bevy_reflect")]
+use {bevy_ecs::reflect::ReflectComponent, bevy_reflect::prelude::*};
 
 /// TODO
+#[derive(Debug, PartialEq, Clone, Copy)]
+#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
+#[cfg_attr(
+    feature = "bevy-support",
+    derive(Component),
+    require(GlobalTransform, TransformTreeChanged)
+)]
+#[cfg_attr(
+    feature = "bevy_reflect",
+    derive(Reflect),
+    reflect(Component, Default, PartialEq, Debug, Clone)
+)]
+#[cfg_attr(
+    all(feature = "bevy_reflect", feature = "serialize"),
+    reflect(Serialize, Deserialize)
+)]
 pub struct Transform2d {
     /// TODO
     pub translation: Vec2,
diff --git a/crates/bevy_transform/src/components/transform_3d.rs b/crates/bevy_transform/src/components/transform_3d.rs
@@ -1,4 +1,4 @@
-use super::GlobalTransform;
+use super::{GlobalTransform, TransformTreeChanged};
 use bevy_math::{Affine3A, Dir3, Isometry3d, Mat3, Mat4, Quat, Vec3};
 use core::ops::Mul;
 
@@ -650,20 +650,3 @@ impl Mul<Vec3> for Transform3d {
         self.transform_point(value)
     }
 }
-
-/// An optimization for transform propagation. This ZST marker component uses change detection to
-/// mark all entities of the hierarchy as "dirty" if any of their descendants have a changed
-/// `Transform3d`. If this component is *not* marked `is_changed()`, propagation will halt.
-#[derive(Clone, Copy, Default, PartialEq, Debug)]
-#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
-#[cfg_attr(feature = "bevy-support", derive(Component))]
-#[cfg_attr(
-    feature = "bevy_reflect",
-    derive(Reflect),
-    reflect(Component, Default, PartialEq, Debug)
-)]
-#[cfg_attr(
-    all(feature = "bevy_reflect", feature = "serialize"),
-    reflect(Serialize, Deserialize)
-)]
-pub struct TransformTreeChanged;
diff --git a/crates/bevy_transform/src/components/tree_changed.rs b/crates/bevy_transform/src/components/tree_changed.rs
@@ -0,0 +1,22 @@
+#[cfg(feature = "bevy-support")]
+use bevy_ecs::component::Component;
+
+#[cfg(feature = "bevy_reflect")]
+use {bevy_ecs::reflect::ReflectComponent, bevy_reflect::prelude::*};
+
+/// An optimization for transform propagation. This ZST marker component uses change detection to
+/// mark all entities of the hierarchy as "dirty" if any of their descendants have a changed
+/// `Transform3d`. If this component is *not* marked `is_changed()`, propagation will halt.
+#[derive(Clone, Copy, Default, PartialEq, Debug)]
+#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
+#[cfg_attr(feature = "bevy-support", derive(Component))]
+#[cfg_attr(
+    feature = "bevy_reflect",
+    derive(Reflect),
+    reflect(Component, Default, PartialEq, Debug)
+)]
+#[cfg_attr(
+    all(feature = "bevy_reflect", feature = "serialize"),
+    reflect(Serialize, Deserialize)
+)]
+pub struct TransformTreeChanged;
diff --git a/crates/bevy_transform/src/plugins.rs b/crates/bevy_transform/src/plugins.rs
@@ -2,6 +2,8 @@ use crate::systems::{mark_dirty_trees, propagate_parent_transforms, sync_simple_
 use bevy_app::{App, Plugin, PostStartup, PostUpdate};
 use bevy_ecs::schedule::{IntoScheduleConfigs, SystemSet};
 
+use crate::components::{GlobalTransform, Transform2d, Transform3d, TransformTreeChanged};
+
 /// Set enum for the systems relating to transform propagation
 #[derive(Debug, Hash, PartialEq, Eq, Clone, SystemSet)]
 pub enum TransformSystems {
@@ -20,31 +22,40 @@ pub struct TransformPlugin;
 impl Plugin for TransformPlugin {
     fn build(&self, app: &mut App) {
         #[cfg(feature = "bevy_reflect")]
-        app.register_type::<crate::components::Transform3d>()
-            .register_type::<crate::components::TransformTreeChanged>()
-            .register_type::<crate::components::GlobalTransform>();
+        app.register_type::<Transform3d>()
+            .register_type::<TransformTreeChanged>()
+            .register_type::<GlobalTransform>();
 
         app
             // add transform systems to startup so the first update is "correct"
+            .add_systems(
+                PostStartup,
+                (mark_dirty_trees, propagate_parent_transforms)
+                    .chain()
+                    .in_set(TransformSystems::Propagate),
+            )
             .add_systems(
                 PostStartup,
                 (
-                    mark_dirty_trees,
-                    propagate_parent_transforms,
-                    sync_simple_transforms,
+                    sync_simple_transforms::<Transform3d>,
+                    sync_simple_transforms::<Transform2d>,
                 )
+                    .after(propagate_parent_transforms)
+                    .in_set(TransformSystems::Propagate),
+            )
+            .add_systems(
+                PostUpdate,
+                (mark_dirty_trees, propagate_parent_transforms)
                     .chain()
                     .in_set(TransformSystems::Propagate),
             )
             .add_systems(
                 PostUpdate,
                 (
-                    mark_dirty_trees,
-                    propagate_parent_transforms,
-                    // TODO: Adjust the internal parallel queries to make this system more efficiently share and fill CPU time.
-                    sync_simple_transforms,
+                    sync_simple_transforms::<Transform3d>,
+                    sync_simple_transforms::<Transform2d>,
                 )
-                    .chain()
+                    .after(propagate_parent_transforms)
                     .in_set(TransformSystems::Propagate),
             );
     }
diff --git a/crates/bevy_transform/src/systems.rs b/crates/bevy_transform/src/systems.rs
@@ -1,4 +1,4 @@
-use crate::components::{GlobalTransform, Transform3d, TransformTreeChanged};
+use crate::components::{GlobalTransform, Transform2d, Transform3d, TransformTreeChanged};
 use bevy_ecs::prelude::*;
 #[cfg(feature = "std")]
 pub use parallel::propagate_parent_transforms;
@@ -9,33 +9,35 @@ pub use serial::propagate_parent_transforms;
 ///
 /// Third party plugins should ensure that this is used in concert with
 /// [`propagate_parent_transforms`] and [`mark_dirty_trees`].
-pub fn sync_simple_transforms(
+pub fn sync_simple_transforms<T>(
     mut query: ParamSet<(
         Query<
-            (&Transform3d, &mut GlobalTransform),
+            (&T, &mut GlobalTransform),
             (
-                Or<(Changed<Transform3d>, Added<GlobalTransform>)>,
+                Or<(Changed<T>, Added<GlobalTransform>)>,
                 Without<ChildOf>,
                 Without<Children>,
             ),
         >,
-        Query<(Ref<Transform3d>, &mut GlobalTransform), (Without<ChildOf>, Without<Children>)>,
+        Query<(Ref<T>, &mut GlobalTransform), (Without<ChildOf>, Without<Children>)>,
     )>,
     mut orphaned: RemovedComponents<ChildOf>,
-) {
+) where
+    T: Copy + Clone + Component + Into<GlobalTransform>,
+{
     // Update changed entities.
     query
         .p0()
         .par_iter_mut()
         .for_each(|(transform, mut global_transform)| {
-            *global_transform = GlobalTransform::from(*transform);
+            *global_transform = (*transform).into();
         });
     // Update orphaned entities.
     let mut query = query.p1();
     let mut iter = query.iter_many_mut(orphaned.read());
     while let Some((transform, mut global_transform)) = iter.fetch_next() {
         if !transform.is_changed() && !global_transform.is_added() {
-            *global_transform = GlobalTransform::from(*transform);
+            *global_transform = (*transform).into();
         }
     }
 }
@@ -48,6 +50,7 @@ pub fn mark_dirty_trees(
         Entity,
         Or<(
             Changed<Transform3d>,
+            Changed<Transform2d>,
             Changed<ChildOf>,
             Added<GlobalTransform>,
         )>,
@@ -107,7 +110,11 @@ mod serial {
         >,
         mut orphaned: RemovedComponents<ChildOf>,
         transform_query: Query<
-            (Ref<Transform3d>, &mut GlobalTransform, Option<&Children>),
+            (
+                AnyOf<(Ref<Transform3d>, Ref<Transform2d>)>,
+                &mut GlobalTransform,
+                Option<&Children>,
+            ),
             With<ChildOf>,
         >,
         child_query: Query<(Entity, Ref<ChildOf>), With<GlobalTransform>>,
@@ -120,7 +127,7 @@ mod serial {
         |(entity, children, transform, mut global_transform)| {
             let changed = transform.is_changed() || global_transform.is_added() || orphaned_entities.binary_search(&entity).is_ok();
             if changed {
-                *global_transform = GlobalTransform::from(*transform);
+                *global_transform = local_to_global(transform);
             }
 
             for (child, child_of) in child_query.iter_many(children) {
@@ -175,7 +182,11 @@ mod serial {
     unsafe fn propagate_recursive(
         parent: &GlobalTransform,
         transform_query: &Query<
-            (Ref<Transform3d>, &mut GlobalTransform, Option<&Children>),
+            (
+                AnyOf<(Ref<Transform3d>, Ref<Transform2d>)>,
+                &mut GlobalTransform,
+                Option<&Children>,
+            ),
             With<ChildOf>,
         >,
         child_query: &Query<(Entity, Ref<ChildOf>), With<GlobalTransform>>,
@@ -252,6 +263,7 @@ mod serial {
 /// the serial version.
 #[cfg(feature = "std")]
 mod parallel {
+    use super::local_to_global;
     use crate::prelude::*;
     // TODO: this implementation could be used in no_std if there are equivalents of these.
     use alloc::{sync::Arc, vec::Vec};
@@ -273,7 +285,12 @@ mod parallel {
     pub fn propagate_parent_transforms(
         mut queue: Local<WorkQueue>,
         mut roots: Query<
-            (Entity, Ref<Transform3d>, &mut GlobalTransform, &Children),
+            (
+                Entity,
+                AnyOf<(Ref<Transform3d>, Ref<Transform2d>)>,
+                &mut GlobalTransform,
+                &Children,
+            ),
             (Without<ChildOf>, Changed<TransformTreeChanged>),
         >,
         nodes: NodeQuery,
@@ -282,7 +299,7 @@ mod parallel {
         roots.par_iter_mut().for_each_init(
             || queue.local_queue.borrow_local_mut(),
             |outbox, (parent, transform, mut parent_transform, children)| {
-                *parent_transform = GlobalTransform::from(*transform);
+                *parent_transform = local_to_global(transform);
 
                 // SAFETY: the parent entities passed into this function are taken from iterating
                 // over the root entity query. Queries iterate over disjoint entities, preventing
@@ -460,7 +477,8 @@ mod parallel {
 
                     // Transform prop is expensive - this helps avoid updating entire subtrees if
                     // the GlobalTransform is unchanged, at the cost of an added equality check.
-                    global_transform.set_if_neq(p_global_transform.mul_transform(*transform));
+                    let child_global = local_to_global(transform);
+                    global_transform.set_if_neq(*p_global_transform * child_global);
 
                     children.map(|children| {
                         // Only continue propagation if the entity has children.
@@ -499,7 +517,7 @@ mod parallel {
         (
             Entity,
             (
-                Ref<'static, Transform3d>,
+                AnyOf<(Ref<'static, Transform3d>, Ref<'static, Transform2d>)>,
                 Mut<'static, GlobalTransform>,
                 Ref<'static, TransformTreeChanged>,
             ),
@@ -557,6 +575,16 @@ mod parallel {
     }
 }
 
+fn local_to_global(
+    any_of: (Option<Ref<'_, Transform3d>>, Option<Ref<'_, Transform2d>>),
+) -> GlobalTransform {
+    match any_of {
+        (Some(transform_3d), _) => GlobalTransform::from(*transform_3d),
+        (_, Some(transform_2d)) => GlobalTransform::from(*transform_2d),
+        _ => unreachable!(),
+    }
+}
+
 #[cfg(test)]
 mod test {
     use alloc::{vec, vec::Vec};
