1#import bevy_render::view::View
2
3@group(0) @binding(0) var<uniform> view: View;
4@group(0) @binding(1) var<storage, read_write> layers: array<vec2<u32>>;
5@group(0) @binding(2) var<storage, read_write> layer_ids: array<atomic<i32>>;
6
7@group(1) @binding(0) var depth: texture_depth_2d;
8
9struct OitFragment {
10    color: vec3<f32>,
11    alpha: f32,
12    depth: f32,
13}
14// Contains all the colors and depth for this specific fragment
15var<private> fragment_list: array<OitFragment, #{LAYER_COUNT}>;
16
17struct FullscreenVertexOutput {
18    @builtin(position) position: vec4<f32>,
19    @location(0) uv: vec2<f32>,
20};
21
22@fragment
23fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
24    let buffer_size = i32(view.viewport.z * view.viewport.w);
25    let screen_index = i32(floor(in.position.x) + floor(in.position.y) * view.viewport.z);
26
27    let counter = atomicLoad(&layer_ids[screen_index]);
28    if counter == 0 {
29        reset_indices(screen_index);
30
31        // https://github.com/gfx-rs/wgpu/issues/4416
32        if true {
33            discard;
34        }
35        return vec4(0.0);
36    } else {
37        // Load depth for manual depth testing.
38        // This is necessary because early z doesn't seem to trigger in the transparent pass.
39        // This should be done during the draw pass so those fragments simply don't exist in the list,
40        // but this requires a bigger refactor
41        let d = textureLoad(depth, vec2<i32>(in.position.xy), 0);
42        let result = sort(screen_index, buffer_size, d);
43        reset_indices(screen_index);
44
45        return result.color;
46    }
47}
48
49// Resets all indices to 0.
50// This means we don't have to clear the entire layers buffer
51fn reset_indices(screen_index: i32) {
52    atomicStore(&layer_ids[screen_index], 0);
53    layers[screen_index] = vec2(0u);
54}
55
56struct SortResult {
57    color: vec4f,
58    depth: f32,
59}
60
61fn sort(screen_index: i32, buffer_size: i32, opaque_depth: f32) -> SortResult {
62    var counter = atomicLoad(&layer_ids[screen_index]);
63
64    // fill list
65    for (var i = 0; i < counter; i += 1) {
66        let fragment = layers[screen_index + buffer_size * i];
67        // unpack color/alpha/depth
68        let color = bevy_pbr::rgb9e5::rgb9e5_to_vec3_(fragment.x);
69        let depth_alpha = bevy_core_pipeline::oit::unpack_24bit_depth_8bit_alpha(fragment.y);
70        fragment_list[i].color = color;
71        fragment_list[i].alpha = depth_alpha.y;
72        fragment_list[i].depth = depth_alpha.x;
73    }
74
75    // bubble sort the list based on the depth
76    for (var i = counter; i >= 0; i -= 1) {
77        for (var j = 0; j < i; j += 1) {
78            if fragment_list[j].depth < fragment_list[j + 1].depth {
79                // swap
80                let temp = fragment_list[j + 1];
81                fragment_list[j + 1] = fragment_list[j];
82                fragment_list[j] = temp;
83            }
84        }
85    }
86
87    // resolve blend
88    var final_color = vec4(0.0);
89    for (var i = 0; i <= counter; i += 1) {
90        // depth testing
91        // This needs to happen here because we can only stop iterating if the fragment is
92        // occluded by something opaque and the fragments need to be sorted first
93        if fragment_list[i].depth < opaque_depth {
94            break;
95        }
96        let color = fragment_list[i].color;
97        let alpha = fragment_list[i].alpha;
98        var base_color = vec4(color.rgb * alpha, alpha);
99        final_color = blend(final_color, base_color);
100        if final_color.a == 1.0 {
101            break;
102        }
103    }
104    var result: SortResult;
105    result.color = final_color;
106    result.depth = fragment_list[0].depth;
107
108    return result;
109}
110
111// OVER operator using premultiplied alpha
112// see: https://en.wikipedia.org/wiki/Alpha_compositing
113fn blend(color_a: vec4<f32>, color_b: vec4<f32>) -> vec4<f32> {
114    let final_color = color_a.rgb + (1.0 - color_a.a) * color_b.rgb;
115    let alpha = color_a.a + (1.0 - color_a.a) * color_b.a;
116    return vec4(final_color.rgb, alpha);
117}