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