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