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// Offsets within the `cluster_offsets_and_counts` buffer for a single cluster.
14//
15// These offsets must be monotonically nondecreasing. That is, indices are
16// always sorted into the following order: point lights, spot lights, reflection
17// probes, irradiance volumes.
18struct ClusterableObjectIndexRanges {
19    // The offset of the index of the first point light.
20    first_point_light_index_offset: u32,
21    // The offset of the index of the first spot light, which also terminates
22    // the list of point lights.
23    first_spot_light_index_offset: u32,
24    // The offset of the index of the first reflection probe, which also
25    // terminates the list of spot lights.
26    first_reflection_probe_index_offset: u32,
27    // The offset of the index of the first irradiance volumes, which also
28    // terminates the list of reflection probes.
29    first_irradiance_volume_index_offset: u32,
30    first_decal_offset: u32,
31    // One past the offset of the index of the final clusterable object for this
32    // cluster.
33    last_clusterable_object_index_offset: u32,
34}
35
36// NOTE: Keep in sync with bevy_pbr/src/light.rs
37fn view_z_to_z_slice(
38    cluster_factors: vec2<f32>,
39    z_slices: u32,
40    view_z: f32,
41    is_orthographic: bool
42) -> u32 {
43    var z_slice: u32 = 0u;
44    if is_orthographic {
45        // NOTE: view_z is correct in the orthographic case
46        z_slice = u32(floor((view_z - cluster_factors.x) * cluster_factors.y));
47    } else {
48        // NOTE: had to use -view_z to make it positive else log(negative) is nan
49        z_slice = u32(log(-view_z) * cluster_factors.x - cluster_factors.y + 1.0);
50    }
51    // NOTE: We use min as we may limit the far z plane used for clustering to be closer than
52    // the furthest thing being drawn. This means that we need to limit to the maximum cluster.
53    return min(z_slice, z_slices - 1u);
54}
55
56fn view_fragment_cluster_index(frag_coord: vec2<f32>, view_z: f32, is_orthographic: bool) -> u32 {
57    let xy = vec2<u32>(floor((frag_coord - bindings::view.viewport.xy) * bindings::lights.cluster_factors.xy));
58    let z_slice = view_z_to_z_slice(
59        bindings::lights.cluster_factors.zw,
60        bindings::lights.cluster_dimensions.z,
61        view_z,
62        is_orthographic
63    );
64    return fragment_cluster_index(vec3(xy, z_slice), bindings::lights.cluster_dimensions);
65}
66
67// Given a cluster XYZ position, returns its index in the cluster list.
68fn fragment_cluster_index(p: vec3<u32>, cluster_dimensions: vec4<u32>) -> u32 {
69    // NOTE: Restricting cluster index to avoid undefined behavior when accessing uniform buffer
70    // arrays based on the cluster index.
71    return min(
72        (p.y * cluster_dimensions.x + p.x) * cluster_dimensions.z + p.z,
73        cluster_dimensions.w - 1u
74    );
75}
76
77// this must match CLUSTER_COUNT_SIZE in light.rs
78const CLUSTER_COUNT_SIZE = 9u;
79
80// Returns the indices of clusterable objects belonging to the given cluster.
81//
82// Note that if fewer than 3 SSBO bindings are available (in WebGL 2,
83// primarily), light probes aren't clustered, and therefore both light probe
84// index ranges will be empty.
85fn unpack_clusterable_object_index_ranges(cluster_index: u32) -> ClusterableObjectIndexRanges {
86#if AVAILABLE_STORAGE_BUFFER_BINDINGS >= 3
87
88    let offset_and_counts_a = bindings::cluster_offsets_and_counts.data[cluster_index][0];
89    let offset_and_counts_b = bindings::cluster_offsets_and_counts.data[cluster_index][1];
90
91    // Sum up the counts to produce the range brackets.
92    //
93    // We could have stored the range brackets in `cluster_offsets_and_counts`
94    // directly, but doing it this way makes the logic in this path more
95    // consistent with the WebGL 2 path below.
96    let point_light_offset = offset_and_counts_a.x;
97    let spot_light_offset = point_light_offset + offset_and_counts_a.y;
98    let reflection_probe_offset = spot_light_offset + offset_and_counts_a.z;
99    let irradiance_volume_offset = reflection_probe_offset + offset_and_counts_a.w;
100    let decal_offset = irradiance_volume_offset + offset_and_counts_b.x;
101    let last_clusterable_offset = decal_offset + offset_and_counts_b.y;
102    return ClusterableObjectIndexRanges(
103        point_light_offset,
104        spot_light_offset,
105        reflection_probe_offset,
106        irradiance_volume_offset,
107        decal_offset,
108        last_clusterable_offset
109    );
110
111#else   // AVAILABLE_STORAGE_BUFFER_BINDINGS >= 3
112
113    let raw_offset_and_counts = bindings::cluster_offsets_and_counts.data[cluster_index >> 2u][cluster_index & ((1u << 2u) - 1u)];
114    //  [ 31     ..     18 | 17      ..      9 | 8       ..     0 ]
115    //  [      offset      | point light count | spot light count ]
116    let offset_and_counts = vec3<u32>(
117        (raw_offset_and_counts >> (CLUSTER_COUNT_SIZE * 2u)) & ((1u << (32u - (CLUSTER_COUNT_SIZE * 2u))) - 1u),
118        (raw_offset_and_counts >> CLUSTER_COUNT_SIZE)        & ((1u << CLUSTER_COUNT_SIZE) - 1u),
119        raw_offset_and_counts                                & ((1u << CLUSTER_COUNT_SIZE) - 1u),
120    );
121
122    // We don't cluster reflection probes or irradiance volumes on this
123    // platform, as there's no room in the UBO. Thus, those offset ranges
124    // (corresponding to `offset_d` and `offset_e` above) are empty and are
125    // simply copies of `offset_c`.
126
127    let offset_a = offset_and_counts.x;
128    let offset_b = offset_a + offset_and_counts.y;
129    let offset_c = offset_b + offset_and_counts.z;
130
131    return ClusterableObjectIndexRanges(offset_a, offset_b, offset_c, offset_c, offset_c, offset_c);
132
133#endif  // AVAILABLE_STORAGE_BUFFER_BINDINGS >= 3
134}
135
136// Returns the index of the clusterable object at the given offset.
137//
138// Note that, in the case of a light probe, the index refers to an element in
139// one of the two `light_probes` sublists, not the `clustered_lights` list.
140fn get_clusterable_object_id(index: u32) -> u32 {
141#if AVAILABLE_STORAGE_BUFFER_BINDINGS >= 3
142    return bindings::clusterable_object_index_lists.data[index];
143#else
144    // The index is correct but in clusterable_object_index_lists we pack 4 u8s into a u32
145    // This means the index into clusterable_object_index_lists is index / 4
146    let indices = bindings::clusterable_object_index_lists.data[index >> 4u][(index >> 2u) &
147        ((1u << 2u) - 1u)];
148    // And index % 4 gives the sub-index of the u8 within the u32 so we shift by 8 * sub-index
149    return (indices >> (8u * (index & ((1u << 2u) - 1u)))) & ((1u << 8u) - 1u);
150#endif
151}
152
153fn cluster_debug_visualization(
154    input_color: vec4<f32>,
155    view_z: f32,
156    is_orthographic: bool,
157    clusterable_object_index_ranges: ClusterableObjectIndexRanges,
158    cluster_index: u32,
159) -> vec4<f32> {
160    var output_color = input_color;
161
162    // Cluster allocation debug (using 'over' alpha blending)
163#ifdef CLUSTERED_FORWARD_DEBUG_Z_SLICES
164    // NOTE: This debug mode visualizes the z-slices
165    let cluster_overlay_alpha = 0.1;
166    var z_slice: u32 = view_z_to_z_slice(
167        bindings::lights.cluster_factors.zw,
168        bindings::lights.cluster_dimensions.z,
169        view_z,
170        is_orthographic
171    );
172    // A hack to make the colors alternate a bit more
173    if (z_slice & 1u) == 1u {
174        z_slice = z_slice + bindings::lights.cluster_dimensions.z / 2u;
175    }
176    let slice_color_hsv = vec3(
177        f32(z_slice) / f32(bindings::lights.cluster_dimensions.z + 1u) * PI_2,
178        1.0,
179        0.5
180    );
181    let slice_color = hsv_to_rgb(slice_color_hsv);
182    output_color = vec4<f32>(
183        (1.0 - cluster_overlay_alpha) * output_color.rgb + cluster_overlay_alpha * slice_color,
184        output_color.a
185    );
186#endif // CLUSTERED_FORWARD_DEBUG_Z_SLICES
187#ifdef CLUSTERED_FORWARD_DEBUG_CLUSTER_COMPLEXITY
188    // NOTE: This debug mode visualizes the number of clusterable objects within
189    // the cluster that contains the fragment. It shows a sort of cluster
190    // complexity measure.
191    let cluster_overlay_alpha = 0.1;
192    let max_complexity_per_cluster = 64.0;
193    let object_count = clusterable_object_index_ranges.first_reflection_probe_index_offset -
194        clusterable_object_index_ranges.first_point_light_index_offset;
195    output_color.r = (1.0 - cluster_overlay_alpha) * output_color.r + cluster_overlay_alpha *
196        smoothstep(0.0, max_complexity_per_cluster, f32(object_count));
197    output_color.g = (1.0 - cluster_overlay_alpha) * output_color.g + cluster_overlay_alpha *
198        (1.0 - smoothstep(0.0, max_complexity_per_cluster, f32(object_count)));
199#endif // CLUSTERED_FORWARD_DEBUG_CLUSTER_COMPLEXITY
200#ifdef CLUSTERED_FORWARD_DEBUG_CLUSTER_COHERENCY
201    // NOTE: Visualizes the cluster to which the fragment belongs
202    let cluster_overlay_alpha = 0.1;
203    var rng = cluster_index;
204    let cluster_color_hsv = vec3(rand_f(&rng) * PI_2, 1.0, 0.5);
205    let cluster_color = hsv_to_rgb(cluster_color_hsv);
206    output_color = vec4<f32>(
207        (1.0 - cluster_overlay_alpha) * output_color.rgb + cluster_overlay_alpha * cluster_color,
208        output_color.a
209    );
210#endif // CLUSTERED_FORWARD_DEBUG_CLUSTER_COHERENCY
211
212    return output_color;
213}