1import package::light_probe::{light_probe_iterator_new, light_probe_iterator_next};
2import package::render::mesh_view_bindings as bindings;
3import package::render::mesh_view_bindings::light_probes;
4import package::render::mesh_view_types::{
5    LIGHT_PROBE_FLAG_AFFECTS_LIGHTMAPPED_MESH_DIFFUSE, LIGHT_PROBE_FLAG_PARALLAX_CORRECT
6};
7import package::render::pbr_lighting::{F_Schlick_vec, LightingInput, LayerLightingInput, LAYER_BASE, LAYER_CLEARCOAT};
8import package::render::clustered_forward::ClusterableObjectIndexRanges;
9import bevy_render::maths::quat_rotate;
10
11/// The maximum representable value in a 32-bit floating point number.
12const FLOAT_MAX: f32 = 3.40282347e+38;
13
14struct EnvironmentMapLight {
15    diffuse: vec3<f32>,
16    specular: vec3<f32>,
17};
18
19struct EnvironmentMapRadiances {
20    irradiance: vec3<f32>,
21    radiance: vec3<f32>,
22}
23
24struct MultiscatterResult {
25    FssEss: vec3<f32>,
26    FmsEms: vec3<f32>,
27    Edss: vec3<f32>,
28}
29
30/// Computes the direction at which to sample the reflection probe.
31///
32/// * `light_from_world` is the matrix that transforms world space into light
33///   probe space (a 1×1×1 cube centered on the origin).
34///
35/// * `parallax_correction_bounds` is the half-extents of the simulated
36///   reflection boundaries used for parallax correction, in light probe space.
37///   It's ignored if the `parallax_correct` parameter is false.
38///
39/// * `parallax_correct` is true if parallax correction is to be applied and
40///   false otherwise.
41fn compute_cubemap_sample_dir(
42    world_ray_origin: vec3<f32>,
43    world_ray_direction: vec3<f32>,
44    light_from_world: mat4x4<f32>,
45    parallax_correction_bounds: vec3<f32>,
46    parallax_correct: bool,
47    view_rotation: vec4<f32>,
48) -> vec3<f32> {
49    var sample_dir: vec3<f32>;
50
51    // If we're supposed to parallax correct, then intersect with the light cube.
52    if (parallax_correct) {
53        // Compute the direction of the ray bouncing off the surface, in light
54        // probe space.
55        // Recall that light probe space is a 1×1×1 cube centered at the origin.
56        let ray_origin = (light_from_world * vec4(world_ray_origin, 1.0)).xyz;
57        let ray_direction = (light_from_world * vec4(world_ray_direction, 0.0)).xyz;
58
59        // Solve for the intersection of that ray with each side of the cube.
60        // Since our light probe is a 1×1×1 cube centered at the origin in light
61        // probe space, the faces of the cube are at X = ±0.5, Y = ±0.5, and Z =
62        // ±0.5.
63        var t0 = (-parallax_correction_bounds - ray_origin) / ray_direction;
64        var t1 = (parallax_correction_bounds - ray_origin) / ray_direction;
65
66        // We're shooting the rays forward, so we need to rule out negative time
67        // values. So, if t is negative, make it a large value so that we won't
68        // choose it below.
69        // We would use infinity here but WGSL forbids it:
70        // https://github.com/gfx-rs/wgpu/issues/5515
71        t0 = select(vec3(FLOAT_MAX), t0, t0 >= vec3(0.0));
72        t1 = select(vec3(FLOAT_MAX), t1, t1 >= vec3(0.0));
73
74        // Choose the minimum valid time value to find the intersection of the
75        // first cube face.
76        let t_min = min(t0, t1);
77        let t = min(min(t_min.x, t_min.y), t_min.z);
78
79        // Compute the sample direction. (It doesn't have to be normalized.)
80        sample_dir = ray_origin + ray_direction * t;
81    } else {
82        // We treat the reflection as infinitely far away in the non-parallax
83        // case, so the ray origin is irrelevant.
84        sample_dir = (light_from_world * vec4(world_ray_direction, 0.0)).xyz;
85    }
86
87    // Rotating the world space ray direction by the environment map light view rotation,
88    // it is equivalent to rotating the environment cubemap itself.
89    sample_dir = quat_rotate(view_rotation, sample_dir);
90
91    // Cubemaps are left-handed, so we negate the Z coordinate.
92    sample_dir.z = -sample_dir.z;
93    return sample_dir;
94}
95
96// Define two versions of this function, one for the case in which there are
97// multiple light probes and one for the case in which only the view light probe
98// is present.
99
100@if(MULTIPLE_LIGHT_PROBES_IN_ARRAY)
101
102fn compute_radiances(
103    input: LayerLightingInput,
104    clusterable_object_index_ranges: ptr<function, ClusterableObjectIndexRanges>,
105    world_position: vec3<f32>,
106    found_diffuse_indirect: bool,
107) -> EnvironmentMapRadiances {
108    // Unpack.
109    let N = input.N;
110    let R = input.R;
111    let perceptual_roughness = input.perceptual_roughness;
112    let roughness = input.roughness;
113
114    var radiances: EnvironmentMapRadiances;
115
116    // Find all reflection probes that contain the fragment. We're going to
117    // accumulate all the radiance and irradiance from them in a weighted sum.
118    var iterator = light_probe_iterator_new(
119        world_position,
120        /*is_irradiance_volume=*/ false,
121        clusterable_object_index_ranges,
122    );
123
124    var total_weight = 0.0;
125    radiances.irradiance = vec3(0.0);
126    radiances.radiance = vec3(0.0);
127
128    while (true) {
129        var query_result = light_probe_iterator_next(&iterator);
130        var view_rotation = vec4<f32>(0.0, 0.0, 0.0, 1.0);
131
132        // If we reached the end of the light probe list, and we didn't find
133        // enough reflection probes to reach a weight of 1.0, use the view
134        // environment map if applicable. This allows for e.g. nice transitions
135        // between the interior of a building and the outdoor environment map.
136        if (query_result.texture_index < 0 && total_weight < 0.9999) {
137            query_result.texture_index = light_probes.view_cubemap_index;
138            query_result.intensity = light_probes.intensity_for_view;
139            view_rotation = light_probes.view_rotation;
140            query_result.light_from_world = mat4x4(
141                vec4(1.0, 0.0, 0.0, 0.0),
142                vec4(0.0, 1.0, 0.0, 0.0),
143                vec4(0.0, 0.0, 1.0, 0.0),
144                vec4(0.0, 0.0, 0.0, 1.0)
145            );
146            query_result.parallax_correction_bounds = vec3(0.0);
147            if light_probes.view_environment_map_affects_lightmapped_mesh_diffuse != 0u {
148                query_result.flags = LIGHT_PROBE_FLAG_AFFECTS_LIGHTMAPPED_MESH_DIFFUSE;
149            } else {
150                query_result.flags = 0u;
151            }
152            query_result.weight = 1.0 - total_weight;
153        }
154
155        // If we reached the end, we're done.
156        if (query_result.texture_index < 0) {
157            break;
158        }
159
160        // Split-sum approximation for image based lighting: https://cdn2.unrealengine.com/Resources/files/2013SiggraphPresentationsNotes-26915738.pdf
161        let radiance_level = perceptual_roughness * f32(textureNumLevels(
162            bindings::specular_environment_maps[query_result.texture_index]) - 1u);
163
164        // If we're lightmapped, and we shouldn't accumulate diffuse light from the
165        // environment map, note that.
166        var enable_diffuse = !found_diffuse_indirect;
167        @if(LIGHTMAP)
168        enable_diffuse = enable_diffuse &&
169            (query_result.flags & LIGHT_PROBE_FLAG_AFFECTS_LIGHTMAPPED_MESH_DIFFUSE) != 0u;
170
171        let parallax_correct = (query_result.flags & LIGHT_PROBE_FLAG_PARALLAX_CORRECT) != 0u;
172
173        if (enable_diffuse) {
174            let irradiance_sample_dir = compute_cubemap_sample_dir(
175                world_position,
176                N,
177                query_result.light_from_world,
178                query_result.parallax_correction_bounds,
179                parallax_correct,
180                view_rotation,
181            );
182            radiances.irradiance += textureSampleLevel(
183                bindings::diffuse_environment_maps[query_result.texture_index],
184                bindings::environment_map_sampler,
185                irradiance_sample_dir,
186                0.0).rgb * query_result.intensity * query_result.weight;
187        }
188
189        var radiance_sample_dir = radiance_sample_direction(N, R, roughness);
190        radiance_sample_dir = compute_cubemap_sample_dir(
191            world_position,
192            radiance_sample_dir,
193            query_result.light_from_world,
194            query_result.parallax_correction_bounds,
195            parallax_correct,
196            view_rotation,
197        );
198        radiances.radiance +=
199            textureSampleLevel(
200                bindings::specular_environment_maps[query_result.texture_index],
201                bindings::environment_map_sampler,
202                radiance_sample_dir,
203                radiance_level).rgb * query_result.intensity * query_result.weight;
204
205        total_weight += query_result.weight;
206    }
207
208    if (total_weight != 0.0) {
209        radiances.irradiance /= total_weight;
210        radiances.radiance /= total_weight;
211    }
212
213    return radiances;
214}
215
216@else   // MULTIPLE_LIGHT_PROBES_IN_ARRAY
217
218fn compute_radiances(
219    input: LayerLightingInput,
220    clusterable_object_index_ranges: ptr<function, ClusterableObjectIndexRanges>,
221    world_position: vec3<f32>,
222    found_diffuse_indirect: bool,
223) -> EnvironmentMapRadiances {
224    // Unpack.
225    let N = input.N;
226    let R = input.R;
227    let perceptual_roughness = input.perceptual_roughness;
228    let roughness = input.roughness;
229
230    var radiances: EnvironmentMapRadiances;
231
232    // If we have no light probe, bail.
233    if (light_probes.view_cubemap_index < 0) {
234        radiances.irradiance = vec3(0.0);
235        radiances.radiance = vec3(0.0);
236        return radiances;
237    }
238
239    // Split-sum approximation for image based lighting: https://cdn2.unrealengine.com/Resources/files/2013SiggraphPresentationsNotes-26915738.pdf
240    // Technically we could use textureNumLevels(specular_environment_map) - 1 here, but we use a uniform
241    // because textureNumLevels() does not work on WebGL2
242    let radiance_level = perceptual_roughness * f32(light_probes.smallest_specular_mip_level_for_view);
243
244    let intensity = light_probes.intensity_for_view;
245
246    // If we're lightmapped, and we shouldn't accumulate diffuse light from the
247    // environment map, note that.
248    var enable_diffuse = !found_diffuse_indirect;
249    @if(LIGHTMAP)
250    enable_diffuse = enable_diffuse &&
251        light_probes.view_environment_map_affects_lightmapped_mesh_diffuse;
252
253    if (enable_diffuse) {
254        var irradiance_sample_dir = N;
255        // Rotating the world space ray direction by the environment map light view rotation,
256        // it is equivalent to rotating the environment cubemap itself.
257        irradiance_sample_dir = quat_rotate(light_probes.view_rotation, irradiance_sample_dir);
258        // Cube maps are left-handed so we negate the z coordinate.
259        irradiance_sample_dir.z = -irradiance_sample_dir.z;
260        radiances.irradiance = textureSampleLevel(
261            bindings::diffuse_environment_map,
262            bindings::environment_map_sampler,
263            irradiance_sample_dir,
264            0.0).rgb * intensity;
265    }
266
267    var radiance_sample_dir = radiance_sample_direction(N, R, roughness);
268    // Rotating the world space ray direction by the environment map light view rotation,
269    // it is equivalent to rotating the environment cubemap itself.
270    radiance_sample_dir = quat_rotate(light_probes.view_rotation, radiance_sample_dir);
271    // Cube maps are left-handed so we negate the z coordinate.
272    radiance_sample_dir.z = -radiance_sample_dir.z;
273    radiances.radiance = textureSampleLevel(
274        bindings::specular_environment_map,
275        bindings::environment_map_sampler,
276        radiance_sample_dir,
277        radiance_level).rgb * intensity;
278
279    return radiances;
280}
281
282@if(STANDARD_MATERIAL_CLEARCOAT)
283
284/// Adds the environment map light from the clearcoat layer to that of the base
285/// layer.
286fn environment_map_light_clearcoat(
287    out: ptr<function, EnvironmentMapLight>,
288    input: ptr<function, LightingInput>,
289    clusterable_object_index_ranges: ptr<function, ClusterableObjectIndexRanges>,
290    found_diffuse_indirect: bool,
291) {
292    // Unpack.
293    let world_position = (*input).P;
294    let clearcoat_NdotV = (*input).layers[LAYER_CLEARCOAT].NdotV;
295    let clearcoat_strength = (*input).clearcoat_strength;
296
297    // Calculate the Fresnel term `Fc` for the clearcoat layer.
298    // 0.04 is a hardcoded value for F0 from the Filament spec.
299    let clearcoat_F0 = vec3<f32>(0.04);
300    let Fc = F_Schlick_vec(clearcoat_F0, 1.0, clearcoat_NdotV) * clearcoat_strength;
301    let inv_Fc = 1.0 - Fc;
302
303    let clearcoat_radiances = compute_radiances(
304        (*input).layers[LAYER_CLEARCOAT],
305        clusterable_object_index_ranges,
306        world_position,
307        found_diffuse_indirect,
308    );
309
310    // Composite the clearcoat layer on top of the existing one.
311    // These formulas are from Filament:
312    // <https://google.github.io/filament/Filament.md.html#lighting/imagebasedlights/clearcoat>
313    (*out).diffuse *= inv_Fc;
314    (*out).specular = (*out).specular * inv_Fc * inv_Fc + clearcoat_radiances.radiance * Fc;
315}
316
317/// Multiscattering approximation: https://www.jcgt.org/published/0008/01/03/paper.pdf
318///
319/// We initially used this (https://bruop.github.io/ibl) reference with Roughness Dependent
320/// Fresnel, but it made fresnel very bright so we reverted to the "typical" fresnel term.
321fn compute_multiscatter(
322    F0: vec3<f32>,
323    F_ab: vec2<f32>,
324    Ems: f32,
325    specular_occlusion: f32,
326) -> MultiscatterResult {
327    let FssEss = (F0 * F_ab.x + F_ab.y) * specular_occlusion;
328    let Favg = F0 + (1.0 - F0) / 21.0;
329    let FmsEms = FssEss * Favg / (1.0 - Ems * Favg) * Ems;
330    let Edss = 1.0 - (FssEss + FmsEms);
331
332    return MultiscatterResult(FssEss, FmsEms, Edss);
333}
334
335fn environment_map_light(
336    input: ptr<function, LightingInput>,
337    clusterable_object_index_ranges: ptr<function, ClusterableObjectIndexRanges>,
338    found_diffuse_indirect: bool,
339) -> EnvironmentMapLight {
340    // Unpack.
341    let roughness = (*input).layers[LAYER_BASE].roughness;
342    let diffuse_color = (*input).diffuse_color;
343    let metallic = (*input).metallic;
344    let NdotV = (*input).layers[LAYER_BASE].NdotV;
345    let F_ab = (*input).F_ab;
346    let F0_dielectric = (*input).F0_dielectric;
347    let F0_metallic = (*input).F0_metallic;
348    let world_position = (*input).P;
349
350    var out: EnvironmentMapLight;
351
352    let radiances = compute_radiances(
353        (*input).layers[LAYER_BASE],
354        clusterable_object_index_ranges,
355        world_position,
356        found_diffuse_indirect,
357    );
358
359    if (all(radiances.irradiance == vec3(0.0)) && all(radiances.radiance == vec3(0.0))) {
360        out.diffuse = vec3(0.0);
361        out.specular = vec3(0.0);
362        return out;
363    }
364
365    // No real world material has specular values under 0.02, so we use this range as a
366    // "pre-baked specular occlusion" that extinguishes the fresnel term, for artistic control.
367    // See: https://google.github.io/filament/Filament.md.html#specularocclusion
368    let F0_surface = mix(F0_dielectric, F0_metallic, metallic);
369    let specular_occlusion = saturate(dot(F0_surface, vec3(50.0 * 0.33)));
370
371    // Compute per-material (dielectric and metallic separately) then mix the results.
372    // We can't use F0 directly as the multiscattering term is nonlinear.
373    let Ems = 1.0 - (F_ab.x + F_ab.y);
374
375    let ms_dielectric = compute_multiscatter(F0_dielectric, F_ab, Ems, specular_occlusion);
376    let ms_metallic = compute_multiscatter(F0_metallic, F_ab, Ems, specular_occlusion);
377
378    let FssEss = mix(ms_dielectric.FssEss, ms_metallic.FssEss, metallic);
379    let FmsEms = mix(ms_dielectric.FmsEms, ms_metallic.FmsEms, metallic);
380    let kD = diffuse_color * ms_dielectric.Edss;
381
382    if (!found_diffuse_indirect) {
383        out.diffuse = (FmsEms + kD) * radiances.irradiance;
384    } else {
385        out.diffuse = vec3(0.0);
386    }
387
388    out.specular = FssEss * radiances.radiance;
389
390    @if(STANDARD_MATERIAL_CLEARCOAT)
391    environment_map_light_clearcoat(
392        &out,
393        input,
394        clusterable_object_index_ranges,
395        found_diffuse_indirect,
396    );
397
398    return out;
399}
400
401/// "Moving Frostbite to Physically Based Rendering 3.0", listing 22
402/// https://seblagarde.wordpress.com/wp-content/uploads/2015/07/course_notes_moving_frostbite_to_pbr_v32.pdf#page=70
403fn radiance_sample_direction(N: vec3<f32>, R: vec3<f32>, roughness: f32) -> vec3<f32> {
404    let smoothness = saturate(1.0 - roughness);
405    let lerp_factor = smoothness * (sqrt(smoothness) + roughness);
406    return mix(N, R, lerp_factor);
407}