Add z-index support with a predictable UI stack (bevyengine#5877)

# Objective

Add consistent UI rendering and interaction where deep nodes inside two different hierarchies will never render on top of one-another by default and offer an escape hatch (z-index) for nodes to change their depth.

## The problem with current implementation

The current implementation of UI rendering is broken in that regard, mainly because [it sets the Z value of the `Transform` component based on a "global Z" space](https://github.com/bevyengine/bevy/blob/main/crates/bevy_ui/src/update.rs#L43) shared by all nodes in the UI. This doesn't account for the fact that each node's final `GlobalTransform` value will be relative to its parent. This effectively makes the depth unpredictable when two deep trees are rendered on top of one-another. 

At the moment, it's also up to each part of the UI code to sort all of the UI nodes. The solution that's offered here does the full sorting of UI node entities once and offers the result through a resource so that all systems can use it.

## Solution

### New ZIndex component
This adds a new optional `ZIndex` enum component for nodes which offers two mechanism:
- `ZIndex::Local(i32)`: Overrides the depth of the node relative to its siblings.
- `ZIndex::Global(i32)`: Overrides the depth of the node relative to the UI root. This basically allows any node in the tree to "escape" the parent and be ordered relative to the entire UI.

Note that in the current implementation, omitting `ZIndex` on a node has the same result as adding `ZIndex::Local(0)`. Additionally, the "global" stacking context is essentially a way to add your node to the root stacking context, so using `ZIndex::Local(n)` on a root node (one without parent) will share that space with all nodes using `Index::Global(n)`.

### New UiStack resource
This adds a new `UiStack` resource which is calculated from both hierarchy and `ZIndex` during UI update and contains a vector of all node entities in the UI, ordered by depth (from farthest from camera to closest). This is exposed publicly by the bevy_ui crate with the hope that it can be used for consistent ordering and to reduce the amount of sorting that needs to be done by UI systems (i.e. instead of sorting everything by `global_transform.z` in every system, this array can be iterated over).

### New z_index example
This also adds a new z_index example that showcases the new `ZIndex` component. It's also a good general demo of the new UI stack system, because making this kind of UI was very broken with the old system (e.g. nodes would render on top of each other, not respecting hierarchy or insert order at all).

![image](https://user-images.githubusercontent.com/1060971/189015985-8ea8f989-0e9d-4601-a7e0-4a27a43a53f9.png)

---

## Changelog

- Added the `ZIndex` component to bevy_ui.
- Added the `UiStack` resource to bevy_ui, and added implementation in a new `stack.rs` module.
- Removed the previous Z updating system from bevy_ui, because it was replaced with the above.
- Changed bevy_ui rendering to use UiStack instead of z ordering.
- Changed bevy_ui focus/interaction system to use UiStack instead of z ordering.
- Added a new z_index example.

## ZIndex demo
Here's a demo I wrote to test these features
https://user-images.githubusercontent.com/1060971/188329295-d7beebd6-9aee-43ab-821e-d437df5dbe8a.mp4


Co-authored-by: Carter Anderson <[email protected]>

Author: 2 people

Cargo.toml

@@ -1462,6 +1462,16 @@ description = "Demonstrates transparency for UI"
 category = "UI (User Interface)"
 wasm = true
 
+[[example]]
+name = "z_index"
+path = "examples/ui/z_index.rs"
+
+[package.metadata.example.z_index]
+name = "UI Z-Index"
+description = "Demonstrates how to control the relative depth (z-position) of UI elements"
+category = "UI (User Interface)"
+wasm = true
+
 [[example]]
 name = "ui"
 path = "examples/ui/ui.rs"

crates/bevy_ui/src/entity.rs

@@ -2,7 +2,7 @@
 
 use crate::{
     widget::{Button, ImageMode},
-    BackgroundColor, CalculatedSize, FocusPolicy, Interaction, Node, Style, UiImage,
+    BackgroundColor, CalculatedSize, FocusPolicy, Interaction, Node, Style, UiImage, ZIndex,
 };
 use bevy_ecs::{
     bundle::Bundle,
@@ -45,6 +45,8 @@ pub struct NodeBundle {
     pub visibility: Visibility,
     /// Algorithmically-computed indication of whether an entity is visible and should be extracted for rendering
     pub computed_visibility: ComputedVisibility,
+    /// Indicates the depth at which the node should appear in the UI
+    pub z_index: ZIndex,
 }
 
 impl Default for NodeBundle {
@@ -60,6 +62,7 @@ impl Default for NodeBundle {
             global_transform: Default::default(),
             visibility: Default::default(),
             computed_visibility: Default::default(),
+            z_index: Default::default(),
         }
     }
 }
@@ -97,6 +100,8 @@ pub struct ImageBundle {
     pub visibility: Visibility,
     /// Algorithmically-computed indication of whether an entity is visible and should be extracted for rendering
     pub computed_visibility: ComputedVisibility,
+    /// Indicates the depth at which the node should appear in the UI
+    pub z_index: ZIndex,
 }
 
 /// A UI node that is text
@@ -126,6 +131,8 @@ pub struct TextBundle {
     pub visibility: Visibility,
     /// Algorithmically-computed indication of whether an entity is visible and should be extracted for rendering
     pub computed_visibility: ComputedVisibility,
+    /// Indicates the depth at which the node should appear in the UI
+    pub z_index: ZIndex,
 }
 
 impl TextBundle {
@@ -174,6 +181,7 @@ impl Default for TextBundle {
             global_transform: Default::default(),
             visibility: Default::default(),
             computed_visibility: Default::default(),
+            z_index: Default::default(),
         }
     }
 }
@@ -211,6 +219,8 @@ pub struct ButtonBundle {
     pub visibility: Visibility,
     /// Algorithmically-computed indication of whether an entity is visible and should be extracted for rendering
     pub computed_visibility: ComputedVisibility,
+    /// Indicates the depth at which the node should appear in the UI
+    pub z_index: ZIndex,
 }
 
 impl Default for ButtonBundle {
@@ -227,6 +237,7 @@ impl Default for ButtonBundle {
             global_transform: Default::default(),
             visibility: Default::default(),
             computed_visibility: Default::default(),
+            z_index: Default::default(),
         }
     }
 }

crates/bevy_ui/src/focus.rs

@@ -1,7 +1,8 @@
-use crate::{entity::UiCameraConfig, CalculatedClip, Node};
+use crate::{entity::UiCameraConfig, CalculatedClip, Node, UiStack};
 use bevy_ecs::{
     entity::Entity,
     prelude::Component,
+    query::WorldQuery,
     reflect::ReflectComponent,
     system::{Local, Query, Res},
 };
@@ -11,7 +12,6 @@ use bevy_reflect::{Reflect, ReflectDeserialize, ReflectSerialize};
 use bevy_render::camera::{Camera, RenderTarget};
 use bevy_render::view::ComputedVisibility;
 use bevy_transform::components::GlobalTransform;
-use bevy_utils::FloatOrd;
 use bevy_window::Windows;
 use serde::{Deserialize, Serialize};
 use smallvec::SmallVec;
@@ -62,6 +62,19 @@ pub struct State {
     entities_to_reset: SmallVec<[Entity; 1]>,
 }
 
+/// Main query for [`ui_focus_system`]
+#[derive(WorldQuery)]
+#[world_query(mutable)]
+pub struct NodeQuery {
+    entity: Entity,
+    node: &'static Node,
+    global_transform: &'static GlobalTransform,
+    interaction: Option<&'static mut Interaction>,
+    focus_policy: Option<&'static FocusPolicy>,
+    calculated_clip: Option<&'static CalculatedClip>,
+    computed_visibility: Option<&'static ComputedVisibility>,
+}
+
 /// The system that sets Interaction for all UI elements based on the mouse cursor activity
 ///
 /// Entities with a hidden [`ComputedVisibility`] are always treated as released.
@@ -71,15 +84,8 @@ pub fn ui_focus_system(
     windows: Res<Windows>,
     mouse_button_input: Res<Input<MouseButton>>,
     touches_input: Res<Touches>,
-    mut node_query: Query<(
-        Entity,
-        &Node,
-        &GlobalTransform,
-        Option<&mut Interaction>,
-        Option<&FocusPolicy>,
-        Option<&CalculatedClip>,
-        Option<&ComputedVisibility>,
-    )>,
+    ui_stack: Res<UiStack>,
+    mut node_query: Query<NodeQuery>,
 ) {
     // reset entities that were both clicked and released in the last frame
     for entity in state.entities_to_reset.drain(..) {
@@ -91,10 +97,8 @@ pub fn ui_focus_system(
     let mouse_released =
         mouse_button_input.just_released(MouseButton::Left) || touches_input.any_just_released();
     if mouse_released {
-        for (_entity, _node, _global_transform, interaction, _focus_policy, _clip, _visibility) in
-            node_query.iter_mut()
-        {
-            if let Some(mut interaction) = interaction {
+        for node in node_query.iter_mut() {
+            if let Some(mut interaction) = node.interaction {
                 if *interaction == Interaction::Clicked {
                     *interaction = Interaction::None;
                 }
@@ -123,15 +127,21 @@ pub fn ui_focus_system(
         .find_map(|window| window.cursor_position())
         .or_else(|| touches_input.first_pressed_position());
 
-    let mut moused_over_z_sorted_nodes = node_query
-        .iter_mut()
-        .filter_map(
-            |(entity, node, global_transform, interaction, focus_policy, clip, visibility)| {
+    // prepare an iterator that contains all the nodes that have the cursor in their rect,
+    // from the top node to the bottom one. this will also reset the interaction to `None`
+    // for all nodes encountered that are no longer hovered.
+    let mut moused_over_nodes = ui_stack
+        .uinodes
+        .iter()
+        // reverse the iterator to traverse the tree from closest nodes to furthest
+        .rev()
+        .filter_map(|entity| {
+            if let Ok(node) = node_query.get_mut(*entity) {
                 // Nodes that are not rendered should not be interactable
-                if let Some(computed_visibility) = visibility {
+                if let Some(computed_visibility) = node.computed_visibility {
                     if !computed_visibility.is_visible() {
                         // Reset their interaction to None to avoid strange stuck state
-                        if let Some(mut interaction) = interaction {
+                        if let Some(mut interaction) = node.interaction {
                             // We cannot simply set the interaction to None, as that will trigger change detection repeatedly
                             if *interaction != Interaction::None {
                                 *interaction = Interaction::None;
@@ -142,12 +152,12 @@ pub fn ui_focus_system(
                     }
                 }
 
-                let position = global_transform.translation();
+                let position = node.global_transform.translation();
                 let ui_position = position.truncate();
-                let extents = node.calculated_size / 2.0;
+                let extents = node.node.size() / 2.0;
                 let mut min = ui_position - extents;
                 let mut max = ui_position + extents;
-                if let Some(clip) = clip {
+                if let Some(clip) = node.calculated_clip {
                     min = Vec2::max(min, clip.clip.min);
                     max = Vec2::min(max, clip.clip.max);
                 }
@@ -161,9 +171,9 @@ pub fn ui_focus_system(
                 };
 
                 if contains_cursor {
-                    Some((entity, focus_policy, interaction, FloatOrd(position.z)))
+                    Some(*entity)
                 } else {
-                    if let Some(mut interaction) = interaction {
+                    if let Some(mut interaction) = node.interaction {
                         if *interaction == Interaction::Hovered
                             || (cursor_position.is_none() && *interaction != Interaction::None)
                         {
@@ -172,41 +182,45 @@ pub fn ui_focus_system(
                     }
                     None
                 }
-            },
-        )
-        .collect::<Vec<_>>();
-
-    moused_over_z_sorted_nodes.sort_by_key(|(_, _, _, z)| -*z);
+            } else {
+                None
+            }
+        })
+        .collect::<Vec<Entity>>()
+        .into_iter();
 
-    let mut moused_over_z_sorted_nodes = moused_over_z_sorted_nodes.into_iter();
-    // set Clicked or Hovered on top nodes
-    for (entity, focus_policy, interaction, _) in moused_over_z_sorted_nodes.by_ref() {
-        if let Some(mut interaction) = interaction {
+    // set Clicked or Hovered on top nodes. as soon as a node with a `Block` focus policy is detected,
+    // the iteration will stop on it because it "captures" the interaction.
+    let mut iter = node_query.iter_many_mut(moused_over_nodes.by_ref());
+    while let Some(node) = iter.fetch_next() {
+        if let Some(mut interaction) = node.interaction {
             if mouse_clicked {
                 // only consider nodes with Interaction "clickable"
                 if *interaction != Interaction::Clicked {
                     *interaction = Interaction::Clicked;
                     // if the mouse was simultaneously released, reset this Interaction in the next
                     // frame
                     if mouse_released {
-                        state.entities_to_reset.push(entity);
+                        state.entities_to_reset.push(node.entity);
                     }
                 }
             } else if *interaction == Interaction::None {
                 *interaction = Interaction::Hovered;
             }
         }
 
-        match focus_policy.cloned().unwrap_or(FocusPolicy::Block) {
+        match node.focus_policy.unwrap_or(&FocusPolicy::Block) {
             FocusPolicy::Block => {
                 break;
             }
             FocusPolicy::Pass => { /* allow the next node to be hovered/clicked */ }
         }
     }
-    // reset lower nodes to None
-    for (_entity, _focus_policy, interaction, _) in moused_over_z_sorted_nodes {
-        if let Some(mut interaction) = interaction {
+    // reset `Interaction` for the remaining lower nodes to `None`. those are the nodes that remain in
+    // `moused_over_nodes` after the previous loop is exited.
+    let mut iter = node_query.iter_many_mut(moused_over_nodes);
+    while let Some(node) = iter.fetch_next() {
+        if let Some(mut interaction) = node.interaction {
             // don't reset clicked nodes because they're handled separately
             if *interaction != Interaction::Clicked && *interaction != Interaction::None {
                 *interaction = Interaction::None;

crates/bevy_ui/src/lib.rs

@@ -6,6 +6,7 @@ mod flex;
 mod focus;
 mod geometry;
 mod render;
+mod stack;
 mod ui_node;
 
 pub mod entity;
@@ -33,7 +34,9 @@ use bevy_ecs::{
 use bevy_input::InputSystem;
 use bevy_transform::TransformSystem;
 use bevy_window::ModifiesWindows;
-use update::{ui_z_system, update_clipping_system};
+use stack::ui_stack_system;
+pub use stack::UiStack;
+use update::update_clipping_system;
 
 use crate::prelude::UiCameraConfig;
 
@@ -48,6 +51,8 @@ pub enum UiSystem {
     Flex,
     /// After this label, input interactions with UI entities have been updated for this frame
     Focus,
+    /// After this label, the [`UiStack`] resource has been updated
+    Stack,
 }
 
 /// The current scale of the UI.
@@ -71,6 +76,7 @@ impl Plugin for UiPlugin {
         app.add_plugin(ExtractComponentPlugin::<UiCameraConfig>::default())
             .init_resource::<FlexSurface>()
             .init_resource::<UiScale>()
+            .init_resource::<UiStack>()
             .register_type::<AlignContent>()
             .register_type::<AlignItems>()
             .register_type::<AlignSelf>()
@@ -135,9 +141,7 @@ impl Plugin for UiPlugin {
             )
             .add_system_to_stage(
                 CoreStage::PostUpdate,
-                ui_z_system
-                    .after(UiSystem::Flex)
-                    .before(TransformSystem::TransformPropagate),
+                ui_stack_system.label(UiSystem::Stack),
             )
             .add_system_to_stage(
                 CoreStage::PostUpdate,