From 7807f0a7eb5b73dc964856924b200a91e6142c2c Mon Sep 17 00:00:00 2001
From: Robert Swain <robert.swain@gmail.com>
Date: Mon, 13 Dec 2021 20:29:58 +0100
Subject: [PATCH] bevy_pbr2: Improve clustered-forward rendering debug
 visualizations

---
 pipelined/bevy_pbr2/src/render/pbr.wgsl | 59 +++++++++++++++----------
 1 file changed, 36 insertions(+), 23 deletions(-)

diff --git a/pipelined/bevy_pbr2/src/render/pbr.wgsl b/pipelined/bevy_pbr2/src/render/pbr.wgsl
index e2a3abd67d5da5..0d25eedd1abfcd 100644
--- a/pipelined/bevy_pbr2/src/render/pbr.wgsl
+++ b/pipelined/bevy_pbr2/src/render/pbr.wgsl
@@ -428,6 +428,10 @@ fn hsv2rgb(hue: f32, saturation: f32, value: f32) -> vec3<f32> {
 	return value * mix( vec3<f32>(1.0), rgb, vec3<f32>(saturation));
 }
 
+fn random1D(s: f32) -> f32 {
+    return fract(sin(s * 12.9898) * 43758.5453123);
+}
+
 struct FragmentInput {
     [[builtin(front_facing)]] is_front: bool;
     [[builtin(position)]] frag_coord: vec4<f32>;
@@ -580,31 +584,40 @@ fn fragment(in: FragmentInput) -> [[location(0)]] vec4<f32> {
                 emissive.rgb * output_color.a,
             output_color.a);
 
-#ifdef CLUSTERED_FORWARD_DEBUG
         // Cluster allocation debug (using 'over' alpha blending)
-        let cluster_debug_mode = 1;
-        let cluster_overlay_alpha = 1.0;
-        if (cluster_debug_mode == 0) {
-            // NOTE: This debug mode visualises the z-slices
-            var z_slice: u32 = view_z_to_z_slice(view_z);
-            // A hack to make the colors alternate a bit more
-            if ((z_slice & 1u) == 1u) {
-                z_slice = z_slice + lights.cluster_dimensions.z / 2u;
-            }
-            let slice_color = hsv2rgb(f32(z_slice) / f32(lights.cluster_dimensions.z + 1u), 1.0, 0.5);
-            output_color = vec4<f32>(
-                (1.0 - cluster_overlay_alpha) * output_color.rgb + cluster_overlay_alpha * slice_color,
-                output_color.a
-            );
-        } elseif (cluster_debug_mode == 1) {
-            // NOTE: This debug mode visualises the number of lights within the cluster that contains
-            // the fragment. It shows a sort of lighting complexity measure.
-            output_color.r = (1.0 - cluster_overlay_alpha) * output_color.r
-                + cluster_overlay_alpha * smoothStep(0.0, 16.0, f32(offset_and_count.count));
-            output_color.g = (1.0 - cluster_overlay_alpha) * output_color.g
-                + cluster_overlay_alpha * (1.0 - smoothStep(0.0, 16.0, f32(offset_and_count.count)));
+#ifdef CLUSTERED_FORWARD_DEBUG_Z_SLICES
+        // NOTE: This debug mode visualises the z-slices
+        let cluster_overlay_alpha = 0.1;
+        var z_slice: u32 = view_z_to_z_slice(view_z, is_orthographic);
+        // A hack to make the colors alternate a bit more
+        if ((z_slice & 1u) == 1u) {
+            z_slice = z_slice + lights.cluster_dimensions.z / 2u;
         }
-#endif
+        let slice_color = hsv2rgb(f32(z_slice) / f32(lights.cluster_dimensions.z + 1u), 1.0, 0.5);
+        output_color = vec4<f32>(
+            (1.0 - cluster_overlay_alpha) * output_color.rgb + cluster_overlay_alpha * slice_color,
+            output_color.a
+        );
+#endif // CLUSTERED_FORWARD_DEBUG_Z_SLICES
+#ifdef CLUSTERED_FORWARD_DEBUG_CLUSTER_LIGHT_COMPLEXITY
+        // NOTE: This debug mode visualises the number of lights within the cluster that contains
+        // the fragment. It shows a sort of lighting complexity measure.
+        let cluster_overlay_alpha = 0.1;
+        let max_light_complexity_per_cluster = 64.0;
+        output_color.r = (1.0 - cluster_overlay_alpha) * output_color.r
+            + cluster_overlay_alpha * smoothStep(0.0, max_light_complexity_per_cluster, f32(offset_and_count.count));
+        output_color.g = (1.0 - cluster_overlay_alpha) * output_color.g
+            + cluster_overlay_alpha * (1.0 - smoothStep(0.0, max_light_complexity_per_cluster, f32(offset_and_count.count)));
+#endif // CLUSTERED_FORWARD_DEBUG_CLUSTER_LIGHT_COMPLEXITY
+#ifdef CLUSTERED_FORWARD_DEBUG_CLUSTER_COHERENCY
+        // NOTE: Visualizes the cluster to which the fragment belongs
+        let cluster_overlay_alpha = 0.1;
+        let cluster_color = hsv2rgb(random1D(f32(cluster_index)), 1.0, 0.5);
+        output_color = vec4<f32>(
+            (1.0 - cluster_overlay_alpha) * output_color.rgb + cluster_overlay_alpha * cluster_color,
+            output_color.a
+        );
+#endif // CLUSTERED_FORWARD_DEBUG_CLUSTER_COHERENCY
 
         // tone_mapping
         output_color = vec4<f32>(reinhard_luminance(output_color.rgb), output_color.a);