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