1// A postprocessing pass that performs screen-space reflections.
2
3#define_import_path bevy_pbr::ssr
4
5#import bevy_core_pipeline::fullscreen_vertex_shader::FullscreenVertexOutput
6#import bevy_pbr::{
7    lighting,
8    lighting::{LAYER_BASE, LAYER_CLEARCOAT},
9    mesh_view_bindings::{view, depth_prepass_texture, deferred_prepass_texture, ssr_settings},
10    pbr_deferred_functions::pbr_input_from_deferred_gbuffer,
11    pbr_deferred_types,
12    pbr_functions,
13    prepass_utils,
14    raymarch::{
15        depth_ray_march_from_cs,
16        depth_ray_march_march,
17        depth_ray_march_new_from_depth,
18        depth_ray_march_to_ws_dir,
19    },
20    utils,
21    view_transformations::{
22        depth_ndc_to_view_z,
23        frag_coord_to_ndc,
24        ndc_to_frag_coord,
25        ndc_to_uv,
26        position_view_to_ndc,
27        position_world_to_ndc,
28        position_world_to_view,
29    },
30}
31#import bevy_render::view::View
32
33#ifdef ENVIRONMENT_MAP
34#import bevy_pbr::environment_map
35#endif
36
37// The texture representing the color framebuffer.
38@group(1) @binding(0) var color_texture: texture_2d<f32>;
39
40// The sampler that lets us sample from the color framebuffer.
41@group(1) @binding(1) var color_sampler: sampler;
42
43// Group 1, bindings 2 and 3 are in `raymarch.wgsl`.
44
45// Returns the reflected color in the RGB channel and the specular occlusion in
46// the alpha channel.
47//
48// The general approach here is similar to [1]. We first project the reflection
49// ray into screen space. Then we perform uniform steps along that screen-space
50// reflected ray, converting each step to view space.
51//
52// The arguments are:
53//
54// * `R_world`: The reflection vector in world space.
55//
56// * `P_world`: The current position in world space.
57//
58// [1]: https://lettier.github.io/3d-game-shaders-for-beginners/screen-space-reflection.html
59fn evaluate_ssr(R_world: vec3<f32>, P_world: vec3<f32>) -> vec4<f32> {
60    let depth_size = vec2<f32>(textureDimensions(depth_prepass_texture));
61
62    var raymarch = depth_ray_march_new_from_depth(depth_size);
63    depth_ray_march_from_cs(&raymarch, position_world_to_ndc(P_world));
64    depth_ray_march_to_ws_dir(&raymarch, normalize(R_world));
65    raymarch.linear_steps = ssr_settings.linear_steps;
66    raymarch.bisection_steps = ssr_settings.bisection_steps;
67    raymarch.use_secant = ssr_settings.use_secant != 0u;
68    raymarch.depth_thickness_linear_z = ssr_settings.thickness;
69    raymarch.jitter = 1.0;  // Disable jitter for now.
70    raymarch.march_behind_surfaces = false;
71
72    let raymarch_result = depth_ray_march_march(&raymarch);
73    if (raymarch_result.hit) {
74        return vec4(
75            textureSampleLevel(color_texture, color_sampler, raymarch_result.hit_uv, 0.0).rgb,
76            0.0
77        );
78    }
79
80    return vec4(0.0, 0.0, 0.0, 1.0);
81}
82
83@fragment
84fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
85    // Sample the depth.
86    var frag_coord = in.position;
87    frag_coord.z = prepass_utils::prepass_depth(in.position, 0u);
88
89    // Load the G-buffer data.
90    let fragment = textureLoad(color_texture, vec2<i32>(frag_coord.xy), 0);
91    let gbuffer = textureLoad(deferred_prepass_texture, vec2<i32>(frag_coord.xy), 0);
92    let pbr_input = pbr_input_from_deferred_gbuffer(frag_coord, gbuffer);
93
94    // Don't do anything if the surface is too rough, since we can't blur or do
95    // temporal accumulation yet.
96    let perceptual_roughness = pbr_input.material.perceptual_roughness;
97    if (perceptual_roughness > ssr_settings.perceptual_roughness_threshold) {
98        return fragment;
99    }
100
101    // Unpack the PBR input.
102    var specular_occlusion = pbr_input.specular_occlusion;
103    let world_position = pbr_input.world_position.xyz;
104    let N = pbr_input.N;
105    let V = pbr_input.V;
106
107    // Calculate the reflection vector.
108    let R = reflect(-V, N);
109
110    // Do the raymarching.
111    let ssr_specular = evaluate_ssr(R, world_position);
112    var indirect_light = ssr_specular.rgb;
113    specular_occlusion *= ssr_specular.a;
114
115    // Sample the environment map if necessary.
116    //
117    // This will take the specular part of the environment map into account if
118    // the ray missed. Otherwise, it only takes the diffuse part.
119    //
120    // TODO: Merge this with the duplicated code in `apply_pbr_lighting`.
121#ifdef ENVIRONMENT_MAP
122    // Unpack values required for environment mapping.
123    let base_color = pbr_input.material.base_color.rgb;
124    let metallic = pbr_input.material.metallic;
125    let reflectance = pbr_input.material.reflectance;
126    let specular_transmission = pbr_input.material.specular_transmission;
127    let diffuse_transmission = pbr_input.material.diffuse_transmission;
128    let diffuse_occlusion = pbr_input.diffuse_occlusion;
129
130#ifdef STANDARD_MATERIAL_CLEARCOAT
131    // Do the above calculations again for the clearcoat layer. Remember that
132    // the clearcoat can have its own roughness and its own normal.
133    let clearcoat = pbr_input.material.clearcoat;
134    let clearcoat_perceptual_roughness = pbr_input.material.clearcoat_perceptual_roughness;
135    let clearcoat_roughness = lighting::perceptualRoughnessToRoughness(clearcoat_perceptual_roughness);
136    let clearcoat_N = pbr_input.clearcoat_N;
137    let clearcoat_NdotV = max(dot(clearcoat_N, pbr_input.V), 0.0001);
138    let clearcoat_R = reflect(-pbr_input.V, clearcoat_N);
139#endif  // STANDARD_MATERIAL_CLEARCOAT
140
141    // Calculate various other values needed for environment mapping.
142    let roughness = lighting::perceptualRoughnessToRoughness(perceptual_roughness);
143    let diffuse_color = pbr_functions::calculate_diffuse_color(
144        base_color,
145        metallic,
146        specular_transmission,
147        diffuse_transmission
148    );
149    let NdotV = max(dot(N, V), 0.0001);
150    let F_ab = lighting::F_AB(perceptual_roughness, NdotV);
151    let F0 = pbr_functions::calculate_F0(base_color, metallic, reflectance);
152
153    // Pack all the values into a structure.
154    var lighting_input: lighting::LightingInput;
155    lighting_input.layers[LAYER_BASE].NdotV = NdotV;
156    lighting_input.layers[LAYER_BASE].N = N;
157    lighting_input.layers[LAYER_BASE].R = R;
158    lighting_input.layers[LAYER_BASE].perceptual_roughness = perceptual_roughness;
159    lighting_input.layers[LAYER_BASE].roughness = roughness;
160    lighting_input.P = world_position.xyz;
161    lighting_input.V = V;
162    lighting_input.diffuse_color = diffuse_color;
163    lighting_input.F0_ = F0;
164    lighting_input.F_ab = F_ab;
165#ifdef STANDARD_MATERIAL_CLEARCOAT
166    lighting_input.layers[LAYER_CLEARCOAT].NdotV = clearcoat_NdotV;
167    lighting_input.layers[LAYER_CLEARCOAT].N = clearcoat_N;
168    lighting_input.layers[LAYER_CLEARCOAT].R = clearcoat_R;
169    lighting_input.layers[LAYER_CLEARCOAT].perceptual_roughness = clearcoat_perceptual_roughness;
170    lighting_input.layers[LAYER_CLEARCOAT].roughness = clearcoat_roughness;
171    lighting_input.clearcoat_strength = clearcoat;
172#endif  // STANDARD_MATERIAL_CLEARCOAT
173
174    // Sample the environment map.
175    let environment_light = environment_map::environment_map_light(&lighting_input, false);
176
177    // Accumulate the environment map light.
178    indirect_light += view.exposure *
179        (environment_light.diffuse * diffuse_occlusion +
180        environment_light.specular * specular_occlusion);
181#endif
182
183    // Write the results.
184    return vec4(fragment.rgb + indirect_light, 1.0);
185}