1#define_import_path bevy_pbr::clustered_forward
2
3#import bevy_pbr::mesh_view_bindings as Bindings
4
5// NOTE: Keep in sync with bevy_pbr/src/light.rs
6fn view_z_to_z_slice(view_z: f32, is_orthographic: bool) -> u32 {
7    var z_slice: u32 = 0u;
8    if (is_orthographic) {
9        // NOTE: view_z is correct in the orthographic case
10        z_slice = u32(floor((view_z - Bindings::lights.cluster_factors.z) * Bindings::lights.cluster_factors.w));
11    } else {
12        // NOTE: had to use -view_z to make it positive else log(negative) is nan
13        z_slice = u32(log(-view_z) * Bindings::lights.cluster_factors.z - Bindings::lights.cluster_factors.w + 1.0);
14    }
15    // NOTE: We use min as we may limit the far z plane used for clustering to be closeer than
16    // the furthest thing being drawn. This means that we need to limit to the maximum cluster.
17    return min(z_slice, Bindings::lights.cluster_dimensions.z - 1u);
18}
19
20fn fragment_cluster_index(frag_coord: vec2<f32>, view_z: f32, is_orthographic: bool) -> u32 {
21    let xy = vec2<u32>(floor(frag_coord * Bindings::lights.cluster_factors.xy));
22    let z_slice = view_z_to_z_slice(view_z, is_orthographic);
23    // NOTE: Restricting cluster index to avoid undefined behavior when accessing uniform buffer
24    // arrays based on the cluster index.
25    return min(
26        (xy.y * Bindings::lights.cluster_dimensions.x + xy.x) * Bindings::lights.cluster_dimensions.z + z_slice,
27        Bindings::lights.cluster_dimensions.w - 1u
28    );
29}
30
31// this must match CLUSTER_COUNT_SIZE in light.rs
32const CLUSTER_COUNT_SIZE = 9u;
33fn unpack_offset_and_counts(cluster_index: u32) -> vec3<u32> {
34#ifdef NO_STORAGE_BUFFERS_SUPPORT
35    let offset_and_counts = Bindings::cluster_offsets_and_counts.data[cluster_index >> 2u][cluster_index & ((1u << 2u) - 1u)];
36    //  [ 31     ..     18 | 17      ..      9 | 8       ..     0 ]
37    //  [      offset      | point light count | spot light count ]
38    return vec3<u32>(
39        (offset_and_counts >> (CLUSTER_COUNT_SIZE * 2u)) & ((1u << (32u - (CLUSTER_COUNT_SIZE * 2u))) - 1u),
40        (offset_and_counts >> CLUSTER_COUNT_SIZE)        & ((1u << CLUSTER_COUNT_SIZE) - 1u),
41        offset_and_counts                                & ((1u << CLUSTER_COUNT_SIZE) - 1u),
42    );
43#else
44    return Bindings::cluster_offsets_and_counts.data[cluster_index].xyz;
45#endif
46}
47
48fn get_light_id(index: u32) -> u32 {
49#ifdef NO_STORAGE_BUFFERS_SUPPORT
50    // The index is correct but in cluster_light_index_lists we pack 4 u8s into a u32
51    // This means the index into cluster_light_index_lists is index / 4
52    let indices = Bindings::cluster_light_index_lists.data[index >> 4u][(index >> 2u) & ((1u << 2u) - 1u)];
53    // And index % 4 gives the sub-index of the u8 within the u32 so we shift by 8 * sub-index
54    return (indices >> (8u * (index & ((1u << 2u) - 1u)))) & ((1u << 8u) - 1u);
55#else
56    return Bindings::cluster_light_index_lists.data[index];
57#endif
58}
59
60fn cluster_debug_visualization(
61    output_color: vec4<f32>,
62    view_z: f32,
63    is_orthographic: bool,
64    offset_and_counts: vec3<u32>,
65    cluster_index: u32,
66) -> vec4<f32> {
67    // Cluster allocation debug (using 'over' alpha blending)
68#ifdef CLUSTERED_FORWARD_DEBUG_Z_SLICES
69    // NOTE: This debug mode visualises the z-slices
70    let cluster_overlay_alpha = 0.1;
71    var z_slice: u32 = view_z_to_z_slice(view_z, is_orthographic);
72    // A hack to make the colors alternate a bit more
73    if ((z_slice & 1u) == 1u) {
74        z_slice = z_slice + Bindings::lights.cluster_dimensions.z / 2u;
75    }
76    let slice_color = hsv2rgb(f32(z_slice) / f32(Bindings::lights.cluster_dimensions.z + 1u), 1.0, 0.5);
77    output_color = vec4<f32>(
78        (1.0 - cluster_overlay_alpha) * output_color.rgb + cluster_overlay_alpha * slice_color,
79        output_color.a
80    );
81#endif // CLUSTERED_FORWARD_DEBUG_Z_SLICES
82#ifdef CLUSTERED_FORWARD_DEBUG_CLUSTER_LIGHT_COMPLEXITY
83    // NOTE: This debug mode visualises the number of lights within the cluster that contains
84    // the fragment. It shows a sort of lighting complexity measure.
85    let cluster_overlay_alpha = 0.1;
86    let max_light_complexity_per_cluster = 64.0;
87    output_color.r = (1.0 - cluster_overlay_alpha) * output_color.r
88        + cluster_overlay_alpha * smoothStep(0.0, max_light_complexity_per_cluster, f32(offset_and_counts[1] + offset_and_counts[2]));
89    output_color.g = (1.0 - cluster_overlay_alpha) * output_color.g
90        + cluster_overlay_alpha * (1.0 - smoothStep(0.0, max_light_complexity_per_cluster, f32(offset_and_counts[1] + offset_and_counts[2])));
91#endif // CLUSTERED_FORWARD_DEBUG_CLUSTER_LIGHT_COMPLEXITY
92#ifdef CLUSTERED_FORWARD_DEBUG_CLUSTER_COHERENCY
93    // NOTE: Visualizes the cluster to which the fragment belongs
94    let cluster_overlay_alpha = 0.1;
95    let cluster_color = hsv2rgb(random1D(f32(cluster_index)), 1.0, 0.5);
96    output_color = vec4<f32>(
97        (1.0 - cluster_overlay_alpha) * output_color.rgb + cluster_overlay_alpha * cluster_color,
98        output_color.a
99    );
100#endif // CLUSTERED_FORWARD_DEBUG_CLUSTER_COHERENCY
101
102    return output_color;
103}