Fix ui depth system

crates/bevy_transform/src/hierarchy/hierarchy.rs

@@ -1,38 +1,7 @@
 use crate::components::{Children, Parent};
-use bevy_ecs::{Command, Commands, Entity, Query, Resources, World};
+use bevy_ecs::{Command, Commands, Entity, Resources, World};
 use bevy_utils::tracing::debug;
 
-pub fn run_on_hierarchy<T, S>(
-    children_query: &Query<&Children>,
-    state: &mut S,
-    entity: Entity,
-    parent_result: Option<T>,
-    mut previous_result: Option<T>,
-    run: &mut dyn FnMut(&mut S, Entity, Option<T>, Option<T>) -> Option<T>,
-) -> Option<T>
-where
-    T: Clone,
-{
-    let parent_result = run(state, entity, parent_result, previous_result);
-    previous_result = None;
-    if let Ok(children) = children_query.get(entity) {
-        for child in children.iter().cloned() {
-            previous_result = run_on_hierarchy(
-                children_query,
-                state,
-                child,
-                parent_result.clone(),
-                previous_result,
-                run,
-            );
-        }
-    } else {
-        previous_result = parent_result;
-    }
-
-    previous_result
-}
-
 #[derive(Debug)]
 pub struct DespawnRecursive {
     entity: Entity,

crates/bevy_ui/src/update.rs

@@ -1,49 +1,149 @@
 use super::Node;
 use bevy_ecs::{Entity, Query, With, Without};
-use bevy_transform::{
-    hierarchy,
-    prelude::{Children, Parent, Transform},
-};
+use bevy_transform::prelude::{Children, Parent, Transform};
 
 pub const UI_Z_STEP: f32 = 0.001;
 
 pub fn ui_z_system(
     root_node_query: Query<Entity, (With<Node>, Without<Parent>)>,
-    mut node_query: Query<(Entity, &mut Transform), With<Node>>,
+    mut node_query: Query<&mut Transform, With<Node>>,
     children_query: Query<&Children>,
 ) {
     let mut current_global_z = 0.0;
-
     for entity in root_node_query.iter() {
-        if let Some(result) = hierarchy::run_on_hierarchy(
+        current_global_z = update_hierarchy(
             &children_query,
             &mut node_query,
             entity,
-            Some(current_global_z),
-            Some(current_global_z),
-            &mut update_node_entity,
-        ) {
-            current_global_z = result;
-        }
+            current_global_z,
+            current_global_z,
+        );
     }
 }
 
-fn update_node_entity(
-    node_query: &mut Query<(Entity, &mut Transform), With<Node>>,
+fn update_hierarchy(
+    children_query: &Query<&Children>,
+    node_query: &mut Query<&mut Transform, With<Node>>,
     entity: Entity,
-    parent_result: Option<f32>,
-    previous_result: Option<f32>,
-) -> Option<f32> {
-    let mut z = UI_Z_STEP;
-    let parent_global_z = parent_result.unwrap();
-    if let Some(previous_global_z) = previous_result {
-        z += previous_global_z - parent_global_z;
-    };
-    let global_z = z + parent_global_z;
-
-    if let Ok(mut transform) = node_query.get_component_mut::<Transform>(entity) {
-        transform.translation.z = z;
+    parent_global_z: f32,
+    mut current_global_z: f32,
+) -> f32 {
+    current_global_z += UI_Z_STEP;
+    if let Ok(mut transform) = node_query.get_mut(entity) {
+        transform.translation.z = current_global_z - parent_global_z;
+    }
+    if let Ok(children) = children_query.get(entity) {
+        let current_parent_global_z = current_global_z;
+        for child in children.iter().cloned() {
+            current_global_z = update_hierarchy(
+                children_query,
+                node_query,
+                child,
+                current_parent_global_z,
+                current_global_z,
+            );
+        }
+    }
+    current_global_z
+}
+#[cfg(test)]
+mod tests {
+    use bevy_ecs::{Commands, IntoSystem, Resources, Schedule, World};
+    use bevy_transform::{components::Transform, hierarchy::BuildChildren};
+
+    use crate::Node;
+
+    use super::{ui_z_system, UI_Z_STEP};
+
+    fn node_with_transform(name: &str) -> (String, Node, Transform) {
+        (name.to_owned(), Node::default(), Transform::default())
     }
 
-    Some(global_z)
+    fn node_without_transform(name: &str) -> (String, Node) {
+        (name.to_owned(), Node::default())
+    }
+
+    fn get_steps(transform: &Transform) -> u32 {
+        (transform.translation.z / UI_Z_STEP).round() as u32
+    }
+
+    #[test]
+    fn test_ui_z_system() {
+        let mut world = World::default();
+        let mut resources = Resources::default();
+        let mut commands = Commands::default();
+        commands.set_entity_reserver(world.get_entity_reserver());
+
+        commands.spawn(node_with_transform("0"));
+
+        commands
+            .spawn(node_with_transform("1"))
+            .with_children(|parent| {
+                parent
+                    .spawn(node_with_transform("1-0"))
+                    .with_children(|parent| {
+                        parent.spawn(node_with_transform("1-0-0"));
+                        parent.spawn(node_without_transform("1-0-1"));
+                        parent.spawn(node_with_transform("1-0-2"));
+                    });
+                parent.spawn(node_with_transform("1-1"));
+                parent
+                    .spawn(node_without_transform("1-2"))
+                    .with_children(|parent| {
+                        parent.spawn(node_with_transform("1-2-0"));
+                        parent.spawn(node_with_transform("1-2-1"));
+                        parent
+                            .spawn(node_with_transform("1-2-2"))
+                            .with_children(|_| ());
+                        parent.spawn(node_with_transform("1-2-3"));
+                    });
+                parent.spawn(node_with_transform("1-3"));
+            });
+
+        commands
+            .spawn(node_without_transform("2"))
+            .with_children(|parent| {
+                parent
+                    .spawn(node_with_transform("2-0"))
+                    .with_children(|_parent| ());
+                parent
+                    .spawn(node_with_transform("2-1"))
+                    .with_children(|parent| {
+                        parent.spawn(node_with_transform("2-1-0"));
+                    });
+            });
+        commands.apply(&mut world, &mut resources);
+
+        let mut schedule = Schedule::default();
+        schedule.add_stage("update");
+        schedule.add_system_to_stage("update", ui_z_system.system());
+        schedule.initialize(&mut world, &mut resources);
+        schedule.run(&mut world, &mut resources);
+
+        let mut actual_result = world
+            .query::<(&String, &Transform)>()
+            .map(|(name, transform)| (name.clone(), get_steps(transform)))
+            .collect::<Vec<(String, u32)>>();
+        actual_result.sort_unstable_by_key(|(name, _)| name.clone());
+        let expected_result = vec![
+            ("0".to_owned(), 1),
+            ("1".to_owned(), 1),
+            ("1-0".to_owned(), 1),
+            ("1-0-0".to_owned(), 1),
+            // 1-0-1 has no transform
+            ("1-0-2".to_owned(), 3),
+            ("1-1".to_owned(), 5),
+            // 1-2 has no transform
+            ("1-2-0".to_owned(), 1),
+            ("1-2-1".to_owned(), 2),
+            ("1-2-2".to_owned(), 3),
+            ("1-2-3".to_owned(), 4),
+            ("1-3".to_owned(), 11),
+            // 2 has no transform
+            ("2-0".to_owned(), 1),
+            ("2-1".to_owned(), 2),
+            ("2-1-0".to_owned(), 1),
+        ];
+        assert_eq!(actual_result, expected_result);
+    }
 }