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