From 15a248b05d4ce154b7c80a0ea8330c35abd98370 Mon Sep 17 00:00:00 2001
From: Nicola Papale <nico@nicopap.ch>
Date: Thu, 21 Jul 2022 13:29:42 +0200
Subject: [PATCH] Update Node ComputedVisibility based on UI RenderLayers

---
 crates/bevy_ui/src/entity.rs     |  2 +-
 crates/bevy_ui/src/lib.rs        | 14 +++++++++++++-
 crates/bevy_ui/src/render/mod.rs | 11 +++++++++++
 crates/bevy_ui/src/update.rs     | 27 ++++++++++++++++++++++++++-
 4 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/crates/bevy_ui/src/entity.rs b/crates/bevy_ui/src/entity.rs
index 43cc8f41425ce8..205fd30c99c66a 100644
--- a/crates/bevy_ui/src/entity.rs
+++ b/crates/bevy_ui/src/entity.rs
@@ -143,7 +143,7 @@ pub struct ButtonBundle {
 /// it will display the UI by default.
 ///
 /// [`Camera`]: bevy_render::camera::Camera
-#[derive(Component, Clone)]
+#[derive(Component, Debug, Clone)]
 pub struct UiCameraConfig {
     /// Whether to output UI to this camera view.
     ///
diff --git a/crates/bevy_ui/src/lib.rs b/crates/bevy_ui/src/lib.rs
index d72c1b83495d16..98890c327ac7b7 100644
--- a/crates/bevy_ui/src/lib.rs
+++ b/crates/bevy_ui/src/lib.rs
@@ -27,9 +27,10 @@ pub mod prelude {
 use bevy_app::prelude::*;
 use bevy_ecs::schedule::{ParallelSystemDescriptorCoercion, SystemLabel};
 use bevy_input::InputSystem;
+use bevy_render::view::VisibilitySystems;
 use bevy_transform::TransformSystem;
 use bevy_window::ModifiesWindows;
-use update::{ui_z_system, update_clipping_system};
+use update::{ui_z_system, update_clipping_system, update_render_visibility};
 
 /// The basic plugin for Bevy UI
 #[derive(Default)]
@@ -42,6 +43,11 @@ pub enum UiSystem {
     Flex,
     /// After this label, input interactions with UI entities have been updated for this frame
     Focus,
+    /// Update the [`ComputedVisibility`] component of [`Node`] entities to reflect
+    /// their visibility in accordance to UI cameras.
+    ///
+    /// [`ComputedVisibility`]: bevy_render::view::ComputedVisibility
+    RenderVisibility,
 }
 
 impl Plugin for UiPlugin {
@@ -94,6 +100,12 @@ impl Plugin for UiPlugin {
                     .before(TransformSystem::TransformPropagate)
                     .after(ModifiesWindows),
             )
+            .add_system_to_stage(
+                CoreStage::PostUpdate,
+                update_render_visibility
+                    .label(UiSystem::RenderVisibility)
+                    .after(VisibilitySystems::CheckVisibility),
+            )
             .add_system_to_stage(
                 CoreStage::PostUpdate,
                 ui_z_system
diff --git a/crates/bevy_ui/src/render/mod.rs b/crates/bevy_ui/src/render/mod.rs
index 4b911c937e2b86..590efcb401e7d7 100644
--- a/crates/bevy_ui/src/render/mod.rs
+++ b/crates/bevy_ui/src/render/mod.rs
@@ -158,6 +158,7 @@ fn get_ui_graph(render_app: &mut App) -> RenderGraph {
     ui_graph
 }
 
+#[derive(Debug)]
 pub struct ExtractedUiNode {
     pub transform: Mat4,
     pub color: Color,
@@ -224,9 +225,19 @@ const UI_CAMERA_FAR: f32 = 1000.0;
 // TODO: Evaluate if we still need this.
 const UI_CAMERA_TRANSFORM_OFFSET: f32 = -0.1;
 
+/// The UI camera used by this Camera's viewport.
+///
+/// This component is inserted into the render world in the
+/// [`extract_default_ui_camera_view`] system.
+///
+/// The component is attached to the "actual" viewport's camera.
+/// The UI camera's `ExtractedView` is attached to the entity in the
+/// `entity` field.
 #[derive(Component, Debug)]
 pub struct UiCamera {
+    /// The entity for the UI camera.
     pub entity: Entity,
+    /// UI nodes layer this camera shows.
     layers: RenderLayers,
 }
 
diff --git a/crates/bevy_ui/src/update.rs b/crates/bevy_ui/src/update.rs
index 05628555f1980b..4141fbe2824c30 100644
--- a/crates/bevy_ui/src/update.rs
+++ b/crates/bevy_ui/src/update.rs
@@ -1,6 +1,6 @@
 //! This module contains systems that update the UI when something changes
 
-use crate::{CalculatedClip, Overflow, Style};
+use crate::{entity::UiCameraConfig, CalculatedClip, Overflow, Style};
 
 use super::Node;
 use bevy_ecs::{
@@ -10,6 +10,7 @@ use bevy_ecs::{
 };
 use bevy_hierarchy::{Children, Parent};
 use bevy_math::Vec2;
+use bevy_render::view::{ComputedVisibility, RenderLayers};
 use bevy_sprite::Rect;
 use bevy_transform::components::{GlobalTransform, Transform};
 
@@ -64,6 +65,30 @@ fn update_hierarchy(
     current_global_z
 }
 
+/// Correct the `ComputedVisibility` set by [`check_visibility`] for UI nodes.
+///
+/// Since [`check_visibility`] has no concept of UI cameras, we need a "post-pass"
+/// where we re-enable UI nodes that are visible because they have a matching
+/// UI camera in their layer.
+///
+/// [`check_visibility`]: bevy_render::view::visibility::check_visibility
+pub fn update_render_visibility(
+    ui_cams: Query<&UiCameraConfig>,
+    mut nodes: Query<(&mut ComputedVisibility, &RenderLayers), With<Node>>,
+) {
+    for config in &ui_cams {
+        // Skip configs with default render layers, since `check_visibility` assumes it
+        if config.ui_render_layers == RenderLayers::default() {
+            continue;
+        }
+        for (mut visibility, node_layers) in &mut nodes {
+            if config.ui_render_layers.intersects(node_layers) {
+                visibility.set_visible_in_view();
+            }
+        }
+    }
+}
+
 /// Updates clipping for all nodes
 pub fn update_clipping_system(
     mut commands: Commands,