1//! A postprocessing pass that performs screen-space reflections.
2
3import bevy_core_pipeline::fullscreen_vertex_shader::FullscreenVertexOutput;
4import package::{
5    render::{
6        clustered_forward,
7        pbr_lighting as lighting,
8        pbr_lighting::{LAYER_BASE, LAYER_CLEARCOAT},
9        mesh_view_bindings::{
10            view,
11            globals,
12            depth_prepass_texture,
13            deferred_prepass_texture,
14            ssr_settings
15        },
16        pbr_functions,
17        utils,
18        view_transformations::{
19            depth_ndc_to_view_z,
20            frag_coord_to_ndc,
21            ndc_to_frag_coord,
22            ndc_to_uv,
23            position_view_to_ndc,
24            position_world_to_ndc,
25            position_world_to_view,
26        },
27    },
28    deferred::{
29        functions::pbr_input_from_deferred_gbuffer,
30        types as pbr_deferred_types,
31    },
32    prepass::utils as prepass_utils,
33    ssr::raymarch::{
34        depth_ray_march_from_cs,
35        depth_ray_march_march,
36        depth_ray_march_new_from_depth,
37        depth_ray_march_to_ws_dir,
38    },
39};
40import bevy_render::{
41    view::View,
42    maths::orthonormalize,
43};
44
45@if(ENVIRONMENT_MAP) import package::light_probe::environment_map;
46
47/// The texture representing the color framebuffer.
48@group(2) @binding(0) var color_texture: texture_2d<f32>;
49
50/// The sampler that lets us sample from the color framebuffer.
51@group(2) @binding(1) var color_sampler: sampler;
52
53// Group 1, bindings 2 and 3 are in `raymarch.wesl`.
54
55@group(2) @binding(4) var stbn_texture: texture_2d_array<f32>;
56
57struct BrdfSample {
58    wi: vec3<f32>,
59    value_over_pdf: vec3<f32>,
60}
61
62fn sample_specular_brdf(wo: vec3<f32>, roughness: f32, F0: vec3<f32>, urand: vec2<f32>, N: vec3<f32>) -> BrdfSample {
63    var brdf_sample: BrdfSample;
64
65    // Use VNDF sampling for the half-vector.
66    let wi = lighting::sample_visible_ggx(urand, roughness, N, wo);
67    let H = normalize(wo + wi);
68    let NdotL = max(dot(N, wi), 0.0001);
69    let NdotV = max(dot(N, wo), 0.0001);
70    let VdotH = max(dot(wo, H), 0.0001);
71
72    let F = lighting::F_Schlick_vec(F0, 1.0, VdotH);
73
74    // Height-correlated Smith G2 / G1(V)
75    let a2 = roughness * roughness;
76    let lambdaV = NdotL * sqrt((NdotV - a2 * NdotV) * NdotV + a2);
77    let lambdaL = NdotV * sqrt((NdotL - a2 * NdotL) * NdotL + a2);
78
79    brdf_sample.wi = wi;
80    brdf_sample.value_over_pdf = F * (NdotV * NdotL + lambdaV) / (lambdaV + lambdaL);
81
82    return brdf_sample;
83}
84
85/// Returns the reflected color in the RGB channel and the specular occlusion in
86/// the alpha channel.
87///
88/// The general approach here is similar to [1]. We first project the reflection
89/// ray into screen space. Then we perform uniform steps along that screen-space
90/// reflected ray, converting each step to view space.
91///
92/// The arguments are:
93///
94/// * `R_world`: The reflection vector in world space.
95///
96/// * `P_world`: The current position in world space.
97///
98/// * `jitter`: Jitter to apply to the first step of the linear search; 0..=1
99///   range.
100///
101/// [1]: https://lettier.github.io/3d-game-shaders-for-beginners/screen-space-reflection.html
102fn evaluate_ssr(R_world: vec3<f32>, P_world: vec3<f32>, jitter: f32) -> vec4<f32> {
103    let depth_size = vec2<f32>(textureDimensions(depth_prepass_texture));
104
105    var raymarch = depth_ray_march_new_from_depth(depth_size);
106    depth_ray_march_from_cs(&raymarch, position_world_to_ndc(P_world));
107    depth_ray_march_to_ws_dir(&raymarch, normalize(R_world));
108    raymarch.linear_steps = ssr_settings.linear_steps;
109    raymarch.bisection_steps = ssr_settings.bisection_steps;
110    raymarch.use_secant = ssr_settings.use_secant != 0u;
111    raymarch.depth_thickness_linear_z = ssr_settings.thickness;
112    raymarch.jitter = jitter;
113    raymarch.march_behind_surfaces = false;
114
115    let raymarch_result = depth_ray_march_march(&raymarch);
116    if (raymarch_result.hit) {
117        return vec4(
118            textureSampleLevel(color_texture, color_sampler, raymarch_result.hit_uv, 0.0).rgb,
119            0.0
120        );
121    }
122
123    return vec4(0.0, 0.0, 0.0, 1.0);
124}
125
126@fragment
127fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
128    // Sample the depth.
129    var frag_coord = in.position;
130    frag_coord.z = prepass_utils::prepass_depth(in.position, 0u);
131
132    // Load the G-buffer data.
133    let fragment = textureLoad(color_texture, vec2<i32>(frag_coord.xy), 0);
134    let gbuffer = textureLoad(deferred_prepass_texture, vec2<i32>(frag_coord.xy), 0);
135    let pbr_input = pbr_input_from_deferred_gbuffer(frag_coord, gbuffer);
136
137    // Don't do anything if the surface is too rough or too smooth
138    let perceptual_roughness = pbr_input.material.perceptual_roughness;
139
140    var min_fade: f32;
141    if (ssr_settings.min_perceptual_roughness >= ssr_settings.min_perceptual_roughness_fully_active) {
142        min_fade = step(ssr_settings.min_perceptual_roughness, perceptual_roughness);
143    } else {
144        min_fade = smoothstep(
145            ssr_settings.min_perceptual_roughness,
146            ssr_settings.min_perceptual_roughness_fully_active,
147            perceptual_roughness
148        );
149    }
150
151    var max_fade: f32;
152    if (ssr_settings.max_perceptual_roughness_starts_to_fade >= ssr_settings.max_perceptual_roughness) {
153        max_fade = step(perceptual_roughness, ssr_settings.max_perceptual_roughness);
154    } else {
155        max_fade = 1.0 - smoothstep(
156            ssr_settings.max_perceptual_roughness_starts_to_fade,
157            ssr_settings.max_perceptual_roughness,
158            perceptual_roughness
159        );
160    }
161
162    var fade = saturate(min_fade) * saturate(max_fade);
163
164    let ndc_position = frag_coord_to_ndc(vec4(in.position.xy, frag_coord.z, 1.0));
165    let uv = ndc_to_uv(ndc_position.xy);
166    let dist = min(uv, vec2(1.0) - uv);
167    var fade_xy: vec2<f32>;
168    if (ssr_settings.edge_fadeout_no_longer_active >= ssr_settings.edge_fadeout_fully_active) {
169        fade_xy = step(vec2(ssr_settings.edge_fadeout_no_longer_active), dist);
170    } else {
171        fade_xy = smoothstep(
172            vec2(ssr_settings.edge_fadeout_no_longer_active),
173            vec2(ssr_settings.edge_fadeout_fully_active),
174            dist
175        );
176    }
177    fade *= fade_xy.x * fade_xy.y;
178
179    // Unpack the PBR input.
180    var specular_occlusion = pbr_input.specular_occlusion;
181    let world_position = pbr_input.world_position.xyz;
182    let N = pbr_input.N;
183    let V = pbr_input.V;
184
185    // Build a basis for sampling the BRDF.
186    let tangent_to_world = orthonormalize(N);
187
188    let roughness = lighting::perceptualRoughnessToRoughness(perceptual_roughness);
189    let F0 = pbr_functions::calculate_F0(pbr_input.material.base_color.rgb, pbr_input.material.metallic, pbr_input.material.reflectance);
190
191    // Get some random numbers. If the spatio-temporal blue noise (STBN) texture
192    // is available (i.e. not the 1x1 placeholder), we use it. Otherwise, we
193    // fall back to procedural noise.
194    let stbn_dims = textureDimensions(stbn_texture);
195    var urand: vec2<f32>;
196    var raymarch_jitter: f32;
197    if (all(stbn_dims > vec2(1u))) {
198        let stbn_layers = textureNumLayers(stbn_texture);
199        let stbn_noise = textureLoad(
200            stbn_texture,
201            vec2<u32>(in.position.xy) % stbn_dims,
202            i32(globals.frame_count % u32(stbn_layers)),
203            0
204        );
205        urand = stbn_noise.xy;
206        // Use the third channel for jitter to avoid correlation with BRDF sampling.
207        raymarch_jitter = stbn_noise.z;
208    } else {
209        // Fallback to PCG-based procedural noise.
210        // We use a XOR-sum of products with large primes to decorrelate the
211        // seed from the screen-space coordinates and frame count, avoiding
212        // visible "crawling" artifacts.
213        var state = (u32(in.position.x) * 2131358057u) ^
214                    (u32(in.position.y) * 3416869721u) ^
215                    (globals.frame_count * 1199786941u);
216        urand = utils::rand_vec2f(&state);
217        raymarch_jitter = utils::rand_f(&state);
218    }
219
220    // Sample the BRDF.
221    let N_tangent = vec3(0.0, 0.0, 1.0);
222    let V_tangent = V * tangent_to_world;
223
224    let brdf_sample = sample_specular_brdf(V_tangent, roughness, F0, urand, N_tangent);
225    let R_stochastic = tangent_to_world * brdf_sample.wi;
226    let brdf_sample_value_over_pdf = brdf_sample.value_over_pdf;
227
228    // Do the raymarching.
229    let ssr_specular = evaluate_ssr(R_stochastic, world_position, raymarch_jitter);
230    var indirect_light = (ssr_specular.rgb * brdf_sample_value_over_pdf) * fade;
231
232    // Sample the environment map if necessary.
233    //
234    // This will take the specular part of the environment map into account.
235    //
236    // TODO: Merge this with the duplicated code in `apply_pbr_lighting`.
237@if(ENVIRONMENT_MAP) {
238    // Unpack values required for environment mapping.
239    let base_color = pbr_input.material.base_color.rgb;
240    let metallic = pbr_input.material.metallic;
241    let reflectance = pbr_input.material.reflectance;
242    let specular_transmission = pbr_input.material.specular_transmission;
243    let diffuse_transmission = pbr_input.material.diffuse_transmission;
244
245    // Do the above calculations again for the clearcoat layer. Remember that
246    // the clearcoat can have its own roughness and its own normal.
247    @if(STANDARD_MATERIAL_CLEARCOAT)
248    let clearcoat = pbr_input.material.clearcoat;
249    @if(STANDARD_MATERIAL_CLEARCOAT)
250    let clearcoat_perceptual_roughness = pbr_input.material.clearcoat_perceptual_roughness;
251    @if(STANDARD_MATERIAL_CLEARCOAT)
252    let clearcoat_roughness = lighting::perceptualRoughnessToRoughness(clearcoat_perceptual_roughness);
253    @if(STANDARD_MATERIAL_CLEARCOAT)
254    let clearcoat_N = pbr_input.clearcoat_N;
255    @if(STANDARD_MATERIAL_CLEARCOAT)
256    let clearcoat_NdotV = max(dot(clearcoat_N, pbr_input.V), 0.0001);
257    @if(STANDARD_MATERIAL_CLEARCOAT)
258    let clearcoat_R = reflect(-pbr_input.V, clearcoat_N);
259
260    // Calculate various other values needed for environment mapping.
261    let env_roughness = lighting::perceptualRoughnessToRoughness(perceptual_roughness);
262    let diffuse_color = pbr_functions::calculate_diffuse_color(
263        base_color,
264        metallic,
265        specular_transmission,
266        diffuse_transmission
267    );
268    let NdotV = max(dot(N, V), 0.0001);
269    let F_ab = lighting::F_AB(perceptual_roughness, NdotV);
270    let F0_dielectric = pbr_functions::calculate_F0_dielectric(reflectance);
271
272    // Don't add stochastic noise to hits that sample the prefiltered env map.
273    // The prefiltered env map already accounts for roughness.
274    let R = reflect(-V, N);
275
276    // Pack all the values into a structure.
277    var lighting_input: lighting::LightingInput;
278    lighting_input.layers[LAYER_BASE].NdotV = NdotV;
279    lighting_input.layers[LAYER_BASE].N = N;
280    lighting_input.layers[LAYER_BASE].R = R;
281    lighting_input.layers[LAYER_BASE].perceptual_roughness = perceptual_roughness;
282    lighting_input.layers[LAYER_BASE].roughness = env_roughness;
283    lighting_input.P = world_position.xyz;
284    lighting_input.V = V;
285    lighting_input.diffuse_color = diffuse_color;
286    lighting_input.metallic = metallic;
287    lighting_input.F0_dielectric = F0_dielectric;
288    lighting_input.F0_metallic = base_color;
289    lighting_input.F_ab = F_ab;
290    @if(STANDARD_MATERIAL_CLEARCOAT) {
291    lighting_input.layers[LAYER_CLEARCOAT].NdotV = clearcoat_NdotV;
292    lighting_input.layers[LAYER_CLEARCOAT].N = clearcoat_N;
293    lighting_input.layers[LAYER_CLEARCOAT].R = clearcoat_R;
294    lighting_input.layers[LAYER_CLEARCOAT].perceptual_roughness = clearcoat_perceptual_roughness;
295    lighting_input.layers[LAYER_CLEARCOAT].roughness = clearcoat_roughness;
296    lighting_input.clearcoat_strength = clearcoat;
297    }
298
299    // Determine which cluster we're in. We'll need this to find the right
300    // reflection probe.
301    let cluster_index = clustered_forward::view_fragment_cluster_index(
302        frag_coord.xy, frag_coord.z, false);
303    var clusterable_object_index_ranges =
304        clustered_forward::unpack_clusterable_object_index_ranges(cluster_index);
305
306    // Sample the environment map.
307    //
308    // We pass `true` for `found_diffuse_indirect` here because we only want
309    // the specular part; the diffuse part was already accumulated in the
310    // main PBR pass.
311    let environment_light = environment_map::environment_map_light(
312        &lighting_input, &clusterable_object_index_ranges, true);
313
314    // Accumulate the environment map light.
315    // Note that we multiply by (1.0 - ssr_specular.a * fade) to occlude the
316    // environment map if SSR hits.
317    indirect_light += (view.exposure * environment_light.specular * specular_occlusion) * (1.0 - (1.0 - ssr_specular.a) * fade);
318}
319
320    // Write the results.
321    return vec4(fragment.rgb + indirect_light, fragment.a);
322}