1// Reference:
2// Order-independent transparency with per-pixel ordered linked lists - https://github.com/KhronosGroup/Vulkan-Samples/blob/main/shaders/oit_linked_lists/combine.frag
3
4#import bevy_render::view::View
5#import bevy_pbr::mesh_view_types::OitFragmentNode
6
7@group(0) @binding(0) var<uniform> view: View;
8@group(0) @binding(1) var<storage, read> nodes: array<OitFragmentNode>;
9@group(0) @binding(2) var<storage, read_write> heads: array<u32>; // No need to be atomic
10@group(0) @binding(3) var<storage, read_write> atomic_counter: u32; // No need to be atomic
11
12#ifndef DEPTH_PREPASS
13@group(1) @binding(0) var depth: texture_depth_2d;
14#endif
15
16struct OitFragment {
17    color: u32,
18    depth_alpha: u32,
19}
20
21struct FullscreenVertexOutput {
22    @builtin(position) position: vec4<f32>,
23    @location(0) uv: vec2<f32>,
24};
25
26const LINKED_LIST_END_SENTINEL: u32 = 0xFFFFFFFFu;
27const SORTED_FRAGMENT_MAX_COUNT: u32 = #{SORTED_FRAGMENT_MAX_COUNT};
28
29@fragment
30fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
31    atomic_counter = 0u;
32    let screen_index = u32(floor(in.position.x) + floor(in.position.y) * view.viewport.z);
33
34    let head = heads[screen_index] - 1u;
35    if head == LINKED_LIST_END_SENTINEL {
36        // https://github.com/gfx-rs/wgpu/issues/4416
37        if true {
38            discard;
39        }
40        return vec4(0.0);
41    } else {
42#ifndef DEPTH_PREPASS
43        // If depth prepass is disabled, load depth for manual depth testing.
44        // This is necessary because early z doesn't seem to trigger in the transparent pass.
45        // This should be done during the draw pass so those fragments simply don't exist in the list,
46        // but this requires a bigger refactor
47        let d = textureLoad(depth, vec2<i32>(in.position.xy), 0);
48#else
49        let d = 0.0;
50#endif
51        let color = resolve(head, d);
52        heads[screen_index] = 0u; // LINKED_LIST_END_SENTINEL + 1u;
53        return color;
54    }
55}
56
57fn resolve(head: u32, opaque_depth: f32) -> vec4<f32> {
58    // Contains all the colors and depth for this specific fragment
59    // Fragments are sorted from front to back (depth values are in descending order)
60    // This should make insertion sort slightly faster
61    // because transparent pass sorts objects so the linked list iteration is usually in descending order.
62    var fragment_list: array<OitFragment, SORTED_FRAGMENT_MAX_COUNT>;
63    var final_color = vec4<f32>(0.0);
64
65    var packed_opaque_depth = bevy_core_pipeline::oit::pack_24bit_depth_8bit_alpha(opaque_depth, 1.0);
66
67    // fill list
68    var current_node = head;
69    var sorted_frag_count = 0u;
70    while current_node != LINKED_LIST_END_SENTINEL {
71        let fragment_node = nodes[current_node];
72        current_node = fragment_node.next;
73
74#ifndef DEPTH_PREPASS
75        // depth testing
76        if fragment_node.depth_alpha < packed_opaque_depth {
77            continue;
78        }
79#endif
80
81        if sorted_frag_count < SORTED_FRAGMENT_MAX_COUNT {
82            // There is still room in the sorted list.
83            // Insert the fragment so that the list stay sorted.
84            var i = sorted_frag_count;
85            for (; i > 0 && fragment_node.depth_alpha > fragment_list[i - 1].depth_alpha; i -= 1) {
86                fragment_list[i] = fragment_list[i - 1];
87            }
88            fragment_list[i].color = fragment_node.color;
89            fragment_list[i].depth_alpha = fragment_node.depth_alpha;
90            sorted_frag_count += 1;
91        } else if fragment_list[0].depth_alpha > fragment_node.depth_alpha {
92            // The fragment is farther than the nearest sorted one.
93            // First, make room by blending the nearest fragment from the sorted list.
94            // Then, insert the fragment in the sorted list.
95            // This is an approximation.
96            let nearest_color = bevy_pbr::rgb9e5::rgb9e5_to_vec3_(fragment_list[0].color);
97            let nearest_alpha = packed_depth_alpha_get_alpha(fragment_list[0].depth_alpha);
98            final_color = blend(final_color, vec4f(nearest_color * nearest_alpha, nearest_alpha));
99            var i = 0u;
100            for (; i < SORTED_FRAGMENT_MAX_COUNT - 1 && fragment_node.depth_alpha < fragment_list[i + 1].depth_alpha; i += 1) {
101                fragment_list[i] = fragment_list[i + 1];
102            }
103            fragment_list[i].color = fragment_node.color;
104            fragment_list[i].depth_alpha = fragment_node.depth_alpha;
105        } else {
106            // The next fragment is nearer than any of the sorted ones.
107            // Blend it early.
108            // This is an approximation.
109            let color = bevy_pbr::rgb9e5::rgb9e5_to_vec3_(fragment_node.color);
110            let alpha = packed_depth_alpha_get_alpha(fragment_node.depth_alpha);
111            final_color = blend(final_color, vec4f(color * alpha, alpha));
112        }
113    }
114
115    // blend sorted fragments
116    for (var i = 0u; i < sorted_frag_count; i += 1) {
117        let color = bevy_pbr::rgb9e5::rgb9e5_to_vec3_(fragment_list[i].color);
118        let alpha = packed_depth_alpha_get_alpha(fragment_list[i].depth_alpha);
119        var base_color = vec4(color.rgb * alpha, alpha);
120        final_color = blend(final_color, base_color);
121        if final_color.a == 1.0 {
122            break;
123        }
124    }
125
126    return final_color;
127}
128
129// OVER operator using premultiplied alpha
130// see: https://en.wikipedia.org/wiki/Alpha_compositing
131fn blend(color_a: vec4<f32>, color_b: vec4<f32>) -> vec4<f32> {
132    let final_color = color_a.rgb + (1.0 - color_a.a) * color_b.rgb;
133    let alpha = color_a.a + (1.0 - color_a.a) * color_b.a;
134    return vec4(final_color.rgb, alpha);
135}
136
137fn packed_depth_alpha_get_alpha(packed: u32) -> f32 {
138    return bevy_core_pipeline::oit::unpack_24bit_depth_8bit_alpha(packed).y;
139}