From fcf6067a107b4d6236ecd7046ea18f8f2e2c1918 Mon Sep 17 00:00:00 2001
From: IceSentry <IceSentry@users.noreply.github.com>
Date: Fri, 18 Oct 2024 21:16:17 -0400
Subject: [PATCH] Fix OIT depth test (#15991)

# Objective

- The depth test was only using the final depth but it should be testing
each fragments
- Fully transparent fragments should not be added to the list

## Solution

- Test each fragment after sorting

## Testing

before:
TODO

after:
TODO
---
 .../bevy_core_pipeline/src/oit/oit_draw.wgsl  |  6 +++++
 .../src/oit/resolve/oit_resolve.wgsl          | 24 ++++++++++++-------
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/crates/bevy_core_pipeline/src/oit/oit_draw.wgsl b/crates/bevy_core_pipeline/src/oit/oit_draw.wgsl
index 80827cd498..739d63e9ea 100644
--- a/crates/bevy_core_pipeline/src/oit/oit_draw.wgsl
+++ b/crates/bevy_core_pipeline/src/oit/oit_draw.wgsl
@@ -5,6 +5,12 @@
 #ifdef OIT_ENABLED
 // Add the fragment to the oit buffer
 fn oit_draw(position: vec4f, color: vec4f) {
+    // Don't add fully transparent fragments to the list
+    // because we don't want to have to sort them in the resolve pass
+    // TODO should this be comparing with < espilon ?
+    if color.a == 0.0 {
+        return;
+    }
     // get the index of the current fragment relative to the screen size
     let screen_index = i32(floor(position.x) + floor(position.y) * view.viewport.z);
     // get the size of the buffer.
diff --git a/crates/bevy_core_pipeline/src/oit/resolve/oit_resolve.wgsl b/crates/bevy_core_pipeline/src/oit/resolve/oit_resolve.wgsl
index bbdb8a035f..41ebb457ad 100644
--- a/crates/bevy_core_pipeline/src/oit/resolve/oit_resolve.wgsl
+++ b/crates/bevy_core_pipeline/src/oit/resolve/oit_resolve.wgsl
@@ -34,16 +34,13 @@ fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
         }
         return vec4(0.0);
     } else {
-        let result = sort(screen_index, buffer_size);
-        reset_indices(screen_index);
-
-        // Manually do depth testing.
+        // Load depth for manual depth testing.
         // This is necessary because early z doesn't seem to trigger in the transparent pass.
-        // Once we have a per pixel linked list it should be done much earlier
+        // This should be done during the draw pass so those fragments simply don't exist in the list,
+        // but this requires a bigger refactor
         let d = textureLoad(depth, vec2<i32>(in.position.xy), 0);
-        if d > result.depth {
-            discard;
-        }
+        let result = sort(screen_index, buffer_size, d);
+        reset_indices(screen_index);
 
         return result.color;
     }
@@ -61,7 +58,7 @@ struct SortResult {
     depth: f32,
 }
 
-fn sort(screen_index: i32, buffer_size: i32) -> SortResult {
+fn sort(screen_index: i32, buffer_size: i32, opaque_depth: f32) -> SortResult {
     var counter = atomicLoad(&layer_ids[screen_index]);
 
     // fill list
@@ -90,10 +87,19 @@ fn sort(screen_index: i32, buffer_size: i32) -> SortResult {
     // resolve blend
     var final_color = vec4(0.0);
     for (var i = 0; i <= counter; i += 1) {
+        // depth testing
+        // This needs to happen here because we can only stop iterating if the fragment is
+        // occluded by something opaque and the fragments need to be sorted first
+        if fragment_list[i].depth < opaque_depth {
+            break;
+        }
         let color = fragment_list[i].color;
         let alpha = fragment_list[i].alpha;
         var base_color = vec4(color.rgb * alpha, alpha);
         final_color = blend(final_color, base_color);
+        if final_color.a == 1.0 {
+            break;
+        }
     }
     var result: SortResult;
     result.color = final_color;