1import bevy_pbr::render::mesh_view_bindings::{view, oit_nodes_capacity, oit_nodes, oit_heads, oit_atomic_counter, oit_settings};
2import bevy_pbr::render::mesh_view_types::OitFragmentNode;
3import bevy_pbr::prepass::utils as prepass_utils;
4
5@if(OIT_ENABLED)
6// Add the fragment to the oit buffer
7fn oit_draw(position: vec4f, color: vec4f) {
8    @if(DEPTH_PREPASS)
9    if position.z < prepass_utils::prepass_depth(position, 0u) {
10        return;
11    }
12    // Don't add fully transparent fragments to the list
13    // because we don't want to have to sort them in the resolve pass
14    if color.r < oit_settings.alpha_threshold
15        && color.g < oit_settings.alpha_threshold
16        && color.b < oit_settings.alpha_threshold
17        && color.a < oit_settings.alpha_threshold
18    {
19        return;
20    }
21    // get the index of the current fragment relative to the screen size
22    let screen_index = u32(floor(position.x) + floor(position.y) * view.viewport.z);
23
24    var new_node_index = atomicAdd(&oit_atomic_counter, 1u);
25    // exit early if we've reached the maximum amount of fragments nodes
26    if new_node_index >= oit_nodes_capacity {
27        // TODO for tail blending we should return the color here
28        return;
29    }
30
31    var node: OitFragmentNode;
32    // In `oit_heads` buffer, index starts from 1, end sentinel is 0 so that we can avoid writing `u32::MAX` from CPU. wgpu guarantees buffers are zero-initialized.
33    node.next = atomicExchange(&oit_heads[screen_index], new_node_index + 1u) - 1u;
34    node.color = bevy_pbr::render::rgb9e5::vec3_to_rgb9e5_(color.rgb);
35    node.depth_alpha = pack_24bit_depth_8bit_alpha(position.z, color.a);
36    oit_nodes[new_node_index] = node;
37}
38
39/// The packing scheme puts depth in the higher bits so that
40///    depth(a) < depth(b) <=> packed(a) < packed(b)
41/// irregardless of alpha(a) and alpha(b)
42/// The property is used to optimize the resolve step
43fn pack_24bit_depth_8bit_alpha(depth: f32, alpha: f32) -> u32 {
44    let depth_bits = u32(saturate(depth) * f32(0xFFFFFFu) + 0.5);
45    let alpha_bits = u32(saturate(alpha) * f32(0xFFu) + 0.5);
46    return (depth_bits << 8u) | alpha_bits;
47}
48
49fn unpack_24bit_depth_8bit_alpha(packed: u32) -> vec2<f32> {
50    let depth_bits = packed >> 8u;
51    let alpha_bits = packed & 0xFFu;
52    return vec2(f32(depth_bits) / f32(0xFFFFFFu), f32(alpha_bits) / f32(0xFFu));
53}