1#define_import_path bevy_pbr::shadows
2
3#import bevy_pbr::{
4    mesh_view_types::POINT_LIGHT_FLAGS_SPOT_LIGHT_Y_NEGATIVE,
5    mesh_view_bindings as view_bindings,
6    shadow_sampling::{
7        SPOT_SHADOW_TEXEL_SIZE, sample_shadow_cubemap, sample_shadow_cubemap_pcss,
8        sample_shadow_map, sample_shadow_map_pcss,
9    }
10}
11
12#import bevy_render::{
13    color_operations::hsv_to_rgb,
14    maths::{orthonormalize, PI_2}
15}
16
17const flip_z: vec3<f32> = vec3<f32>(1.0, 1.0, -1.0);
18
19fn fetch_point_shadow(
20    light_id: u32,
21    frag_position: vec4<f32>,
22    surface_normal: vec3<f32>,
23    frag_coord_xy: vec2<f32>,
24) -> f32 {
25    let light = &view_bindings::clustered_lights.data[light_id];
26
27    // because the shadow maps align with the axes and the frustum planes are at 45 degrees
28    // we can get the worldspace depth by taking the largest absolute axis
29    let surface_to_light = (*light).position_radius.xyz - frag_position.xyz;
30    let surface_to_light_abs = abs(surface_to_light);
31    let distance_to_light = max(surface_to_light_abs.x, max(surface_to_light_abs.y, surface_to_light_abs.z));
32
33    // The normal bias here is already scaled by the texel size at 1 world unit from the light.
34    // The texel size increases proportionally with distance from the light so multiplying by
35    // distance to light scales the normal bias to the texel size at the fragment distance.
36    let normal_offset = (*light).shadow_normal_bias * distance_to_light * surface_normal.xyz;
37    let depth_offset = (*light).shadow_depth_bias * normalize(surface_to_light.xyz);
38    let offset_position = frag_position.xyz + normal_offset + depth_offset;
39
40    // similar largest-absolute-axis trick as above, but now with the offset fragment position
41    let frag_ls = offset_position.xyz - (*light).position_radius.xyz ;
42    let abs_position_ls = abs(frag_ls);
43    let major_axis_magnitude = max(abs_position_ls.x, max(abs_position_ls.y, abs_position_ls.z));
44
45    // NOTE: These simplifications come from multiplying:
46    // projection * vec4(0, 0, -major_axis_magnitude, 1.0)
47    // and keeping only the terms that have any impact on the depth.
48    // Projection-agnostic approach:
49    let zw = -major_axis_magnitude * (*light).light_custom_data.xy + (*light).light_custom_data.zw;
50    let depth = zw.x / zw.y;
51
52    // If soft shadows are enabled, use the PCSS path. Cubemaps assume a
53    // left-handed coordinate space, so we have to flip the z-axis when
54    // sampling.
55    if ((*light).soft_shadow_size > 0.0) {
56        return sample_shadow_cubemap_pcss(
57            frag_ls * flip_z,
58            distance_to_light,
59            depth,
60            light_id,
61            (*light).soft_shadow_size,
62            frag_coord_xy,
63        );
64    }
65
66    // Do the lookup, using HW PCF and comparison. Cubemaps assume a left-handed
67    // coordinate space, so we have to flip the z-axis when sampling.
68    return sample_shadow_cubemap(frag_ls * flip_z, distance_to_light, depth, light_id, frag_coord_xy);
69}
70
71fn fetch_spot_shadow(
72    light_id: u32,
73    frag_position: vec4<f32>,
74    surface_normal: vec3<f32>,
75    near_z: f32,
76    frag_coord_xy: vec2<f32>,
77) -> f32 {
78    let light = &view_bindings::clustered_lights.data[light_id];
79
80    let surface_to_light = (*light).position_radius.xyz - frag_position.xyz;
81
82    // construct the light view matrix
83    var spot_dir = vec3<f32>((*light).light_custom_data.x, 0.0, (*light).light_custom_data.y);
84    // reconstruct spot dir from x/z and y-direction flag
85    spot_dir.y = sqrt(max(0.0, 1.0 - spot_dir.x * spot_dir.x - spot_dir.z * spot_dir.z));
86    if (((*light).flags & POINT_LIGHT_FLAGS_SPOT_LIGHT_Y_NEGATIVE) != 0u) {
87        spot_dir.y = -spot_dir.y;
88    }
89
90    // view matrix z_axis is the reverse of transform.forward()
91    let fwd = -spot_dir;
92    let distance_to_light = dot(fwd, surface_to_light);
93    let offset_position =
94        -surface_to_light
95        + ((*light).shadow_depth_bias * normalize(surface_to_light))
96        + (surface_normal.xyz * (*light).shadow_normal_bias) * distance_to_light;
97
98    let light_inv_rot = orthonormalize(fwd);
99
100    // because the matrix is a pure rotation matrix, the inverse is just the transpose, and to calculate
101    // the product of the transpose with a vector we can just post-multiply instead of pre-multiplying.
102    // this allows us to keep the matrix construction code identical between CPU and GPU.
103    let projected_position = offset_position * light_inv_rot;
104
105    // divide xy by perspective matrix "f" and by -projected.z (projected.z is -projection matrix's w)
106    // to get ndc coordinates
107    let f_div_minus_z = 1.0 / ((*light).spot_light_tan_angle * -projected_position.z);
108    let shadow_xy_ndc = projected_position.xy * f_div_minus_z;
109    // convert to uv coordinates
110    let shadow_uv = shadow_xy_ndc * vec2<f32>(0.5, -0.5) + vec2<f32>(0.5, 0.5);
111
112    let depth = near_z / -projected_position.z;
113
114    // If soft shadows are enabled, use the PCSS path.
115    let array_index = i32(light_id) + view_bindings::lights.spot_light_shadowmap_offset;
116    if ((*light).soft_shadow_size > 0.0) {
117        return sample_shadow_map_pcss(
118            shadow_uv,
119            depth,
120            array_index,
121            frag_coord_xy,
122            SPOT_SHADOW_TEXEL_SIZE,
123            (*light).soft_shadow_size,
124        );
125    }
126
127    return sample_shadow_map(
128        shadow_uv,
129        depth,
130        array_index,
131        frag_coord_xy,
132        SPOT_SHADOW_TEXEL_SIZE,
133    );
134}
135
136fn get_cascade_index(light_id: u32, view_z: f32) -> u32 {
137    let light = &view_bindings::lights.directional_lights[light_id];
138
139    for (var i: u32 = 0u; i < (*light).num_cascades; i = i + 1u) {
140        if (-view_z < (*light).cascades[i].far_bound) {
141            return i;
142        }
143    }
144    return (*light).num_cascades;
145}
146
147// Converts from world space to the uv position in the light's shadow map.
148//
149// The depth is stored in the return value's z coordinate. If the return value's
150// w coordinate is 0.0, then we landed outside the shadow map entirely.
151fn world_to_directional_light_local(
152    light_id: u32,
153    cascade_index: u32,
154    offset_position: vec4<f32>
155) -> vec4<f32> {
156    let light = &view_bindings::lights.directional_lights[light_id];
157    let cascade = &(*light).cascades[cascade_index];
158
159    let offset_position_clip = (*cascade).clip_from_world * offset_position;
160    if (offset_position_clip.w <= 0.0) {
161        return vec4(0.0);
162    }
163    let offset_position_ndc = offset_position_clip.xyz / offset_position_clip.w;
164    // No shadow outside the orthographic projection volume
165    if (any(offset_position_ndc.xy < vec2<f32>(-1.0)) || offset_position_ndc.z < 0.0
166            || any(offset_position_ndc > vec3<f32>(1.0))) {
167        return vec4(0.0);
168    }
169
170    // compute texture coordinates for shadow lookup, compensating for the Y-flip difference
171    // between the NDC and texture coordinates
172    let flip_correction = vec2<f32>(0.5, -0.5);
173    let light_local = offset_position_ndc.xy * flip_correction + vec2<f32>(0.5, 0.5);
174
175    let depth = offset_position_ndc.z;
176
177    return vec4(light_local, depth, 1.0);
178}
179
180fn sample_directional_cascade(
181    light_id: u32,
182    cascade_index: u32,
183    frag_position: vec4<f32>,
184    surface_normal: vec3<f32>,
185    frag_coord_xy: vec2<f32>,
186) -> f32 {
187    let light = &view_bindings::lights.directional_lights[light_id];
188    let cascade = &(*light).cascades[cascade_index];
189
190    // The normal bias is scaled to the texel size.
191    let normal_offset = (*light).shadow_normal_bias * (*cascade).texel_size * surface_normal.xyz;
192    let depth_offset = (*light).shadow_depth_bias * (*light).direction_to_light.xyz;
193    let offset_position = vec4<f32>(frag_position.xyz + normal_offset + depth_offset, frag_position.w);
194
195    let light_local = world_to_directional_light_local(light_id, cascade_index, offset_position);
196    if (light_local.w == 0.0) {
197        return 1.0;
198    }
199
200    let array_index = i32((*light).depth_texture_base_index + cascade_index);
201    let texel_size = (*cascade).texel_size;
202
203    // If soft shadows are enabled, use the PCSS path.
204    if ((*light).soft_shadow_size > 0.0) {
205        return sample_shadow_map_pcss(
206            light_local.xy,
207            light_local.z,
208            array_index,
209            frag_coord_xy,
210            texel_size,
211            (*light).soft_shadow_size,
212        );
213    }
214
215    return sample_shadow_map(
216        light_local.xy,
217        light_local.z,
218        array_index,
219        frag_coord_xy,
220        texel_size,
221    );
222}
223
224fn fetch_directional_shadow(
225    light_id: u32,
226    frag_position: vec4<f32>,
227    surface_normal: vec3<f32>,
228    view_z: f32,
229    frag_coord_xy: vec2<f32>,
230) -> f32 {
231    let light = &view_bindings::lights.directional_lights[light_id];
232    let cascade_index = get_cascade_index(light_id, view_z);
233
234    if (cascade_index >= (*light).num_cascades) {
235        return 1.0;
236    }
237
238    var shadow = sample_directional_cascade(
239        light_id,
240        cascade_index,
241        frag_position,
242        surface_normal,
243        frag_coord_xy,
244    );
245
246    // Blend with the next cascade, if there is one.
247    let next_cascade_index = cascade_index + 1u;
248    if (next_cascade_index < (*light).num_cascades) {
249        let this_far_bound = (*light).cascades[cascade_index].far_bound;
250        let next_near_bound = (1.0 - (*light).cascades_overlap_proportion) * this_far_bound;
251        if (-view_z >= next_near_bound) {
252            let next_shadow = sample_directional_cascade(
253                light_id,
254                next_cascade_index,
255                frag_position,
256                surface_normal,
257                frag_coord_xy,
258            );
259            shadow = mix(shadow, next_shadow, (-view_z - next_near_bound) / (this_far_bound - next_near_bound));
260        }
261    }
262    return shadow;
263}
264
265fn cascade_debug_visualization(
266    output_color: vec3<f32>,
267    light_id: u32,
268    view_z: f32,
269) -> vec3<f32> {
270    let overlay_alpha = 0.95;
271    let cascade_index = get_cascade_index(light_id, view_z);
272    let cascade_color_hsv = vec3(
273        f32(cascade_index) / f32(#{MAX_CASCADES_PER_LIGHT}u + 1u) * PI_2,
274        1.0,
275        0.5
276    );
277    let cascade_color = hsv_to_rgb(cascade_color_hsv);
278    return vec3<f32>(
279        (1.0 - overlay_alpha) * output_color.rgb + overlay_alpha * cascade_color
280    );
281}