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