1#define_import_path bevy_pbr::transmission
2
3#import bevy_pbr::{
4    lighting,
5    prepass_utils,
6    utils::interleaved_gradient_noise,
7    utils,
8    mesh_view_bindings as view_bindings,
9};
10
11#import bevy_render::maths::PI
12
13#ifdef TONEMAP_IN_SHADER
14#import bevy_core_pipeline::tonemapping::approximate_inverse_tone_mapping
15#endif
16
17fn specular_transmissive_light(world_position: vec4<f32>, frag_coord: vec3<f32>, view_z: f32, N: vec3<f32>, V: vec3<f32>, F0: vec3<f32>, ior: f32, thickness: f32, perceptual_roughness: f32, specular_transmissive_color: vec3<f32>, transmitted_environment_light_specular: vec3<f32>) -> vec3<f32> {
18    // Calculate the ratio between refaction indexes. Assume air/vacuum for the space outside the mesh
19    let eta = 1.0 / ior;
20
21    // Calculate incidence vector (opposite to view vector) and its dot product with the mesh normal
22    let I = -V;
23    let NdotI = dot(N, I);
24
25    // Calculate refracted direction using Snell's law
26    let k = 1.0 - eta * eta * (1.0 - NdotI * NdotI);
27    let T = eta * I - (eta * NdotI + sqrt(k)) * N;
28
29    // Calculate the exit position of the refracted ray, by propagating refacted direction through thickness
30    let exit_position = world_position.xyz + T * thickness;
31
32    // Transform exit_position into clip space
33    let clip_exit_position = view_bindings::view.clip_from_world  * vec4<f32>(exit_position, 1.0);
34
35    // Scale / offset position so that coordinate is in right space for sampling transmissive background texture
36    let offset_position = (clip_exit_position.xy / clip_exit_position.w) * vec2<f32>(0.5, -0.5) + 0.5;
37
38    // Fetch background color
39    var background_color: vec4<f32>;
40    if perceptual_roughness == 0.0 {
41        // If the material has zero roughness, we can use a faster approach without the blur
42        background_color = fetch_transmissive_background_non_rough(offset_position, frag_coord);
43    } else {
44        background_color = fetch_transmissive_background(offset_position, frag_coord, view_z, perceptual_roughness);
45    }
46
47    // Compensate for exposure, since the background color is coming from an already exposure-adjusted texture
48    background_color = vec4(background_color.rgb / view_bindings::view.exposure, background_color.a);
49
50    // Dot product of the refracted direction with the exit normal (Note: We assume the exit normal is the entry normal but inverted)
51    let MinusNdotT = dot(-N, T);
52
53    // Calculate 1.0 - fresnel factor (how much light is _NOT_ reflected, i.e. how much is transmitted)
54    let F = vec3(1.0) - lighting::fresnel(F0, MinusNdotT);
55
56    // Calculate final color by applying fresnel multiplied specular transmissive color to a mix of background color and transmitted specular environment light
57    return F * specular_transmissive_color * mix(transmitted_environment_light_specular, background_color.rgb, background_color.a);
58}
59
60fn fetch_transmissive_background_non_rough(offset_position: vec2<f32>, frag_coord: vec3<f32>) -> vec4<f32> {
61    var background_color = textureSampleLevel(
62        view_bindings::view_transmission_texture,
63        view_bindings::view_transmission_sampler,
64        offset_position,
65        0.0
66    );
67
68#ifdef DEPTH_PREPASS
69#ifndef WEBGL2
70    // Use depth prepass data to reject values that are in front of the current fragment
71    if prepass_utils::prepass_depth(vec4<f32>(offset_position * view_bindings::view.viewport.zw, 0.0, 0.0), 0u) > frag_coord.z {
72        background_color.a = 0.0;
73    }
74#endif
75#endif
76
77#ifdef TONEMAP_IN_SHADER
78    background_color = approximate_inverse_tone_mapping(background_color, view_bindings::view.color_grading);
79#endif
80
81    return background_color;
82}
83
84fn fetch_transmissive_background(offset_position: vec2<f32>, frag_coord: vec3<f32>, view_z: f32, perceptual_roughness: f32) -> vec4<f32> {
85    // Calculate view aspect ratio, used to scale offset so that it's proportionate
86    let aspect = view_bindings::view.viewport.z / view_bindings::view.viewport.w;
87
88    // Calculate how “blurry” the transmission should be.
89    // Blur is more or less eyeballed to look approximately “right”, since the “correct”
90    // approach would involve projecting many scattered rays and figuring out their individual
91    // exit positions. IRL, light rays can be scattered when entering/exiting a material (due to
92    // roughness) or inside the material (due to subsurface scattering). Here, we only consider
93    // the first scenario.
94    //
95    // Blur intensity is:
96    // - proportional to the square of `perceptual_roughness`
97    // - proportional to the inverse of view z
98    let blur_intensity = (perceptual_roughness * perceptual_roughness) / view_z;
99
100#ifdef SCREEN_SPACE_SPECULAR_TRANSMISSION_BLUR_TAPS
101    let num_taps = #{SCREEN_SPACE_SPECULAR_TRANSMISSION_BLUR_TAPS}; // Controlled by the `Camera3d::screen_space_specular_transmission_quality` property
102#else
103    let num_taps = 8; // Fallback to 8 taps, if not specified
104#endif
105    let num_spirals = i32(ceil(f32(num_taps) / 8.0));
106#ifdef TEMPORAL_JITTER
107    let random_angle = interleaved_gradient_noise(frag_coord.xy, view_bindings::globals.frame_count);
108#else
109    let random_angle = interleaved_gradient_noise(frag_coord.xy, 0u);
110#endif
111    // Pixel checkerboard pattern (helps make the interleaved gradient noise pattern less visible)
112    let pixel_checkboard = (
113#ifdef TEMPORAL_JITTER
114        // 0 or 1 on even/odd pixels, alternates every frame
115        (i32(frag_coord.x) + i32(frag_coord.y) + i32(view_bindings::globals.frame_count)) % 2
116#else
117        // 0 or 1 on even/odd pixels
118        
119#else
120        // 0 or 1 on even/odd pixels
121        (i32(frag_coord.x) + i32(frag_coord.y)) % 2
122#endif
123    );
124
125    var result = vec4<f32>(0.0);
126    for (var i: i32 = 0; i < num_taps; i = i + 1) {
127        let current_spiral = (i >> 3u);
128        let angle = (random_angle + f32(current_spiral) / f32(num_spirals)) * 2.0 * PI;
129        let m = vec2(sin(angle), cos(angle));
130        let rotation_matrix = mat2x2(
131            m.y, -m.x,
132            m.x, m.y
133        );
134
135        // Get spiral offset
136        var spiral_offset: vec2<f32>;
137        switch i & 7 {
138            // https://www.iryoku.com/next-generation-post-processing-in-call-of-duty-advanced-warfare (slides 120-135)
139            // TODO: Figure out a more reasonable way of doing this, as WGSL
140            // seems to only allow constant indexes into constant arrays at the moment.
141            // The downstream shader compiler should be able to optimize this into a single
142            // constant when unrolling the for loop, but it's still not ideal.
143            case 0: { spiral_offset = utils::SPIRAL_OFFSET_0_; } // Note: We go even first and then odd, so that the lowest
144            case 1: { spiral_offset = utils::SPIRAL_OFFSET_2_; } // quality possible (which does 4 taps) still does a full spiral
145            case 2: { spiral_offset = utils::SPIRAL_OFFSET_4_; } // instead of just the first half of it
146            case 3: { spiral_offset = utils::SPIRAL_OFFSET_6_; }
147            case 4: { spiral_offset = utils::SPIRAL_OFFSET_1_; }
148            case 5: { spiral_offset = utils::SPIRAL_OFFSET_3_; }
149            case 6: { spiral_offset = utils::SPIRAL_OFFSET_5_; }
150            case 7: { spiral_offset = utils::SPIRAL_OFFSET_7_; }
151            default: {}
152        }
153
154        // Make each consecutive spiral slightly smaller than the previous one
155        spiral_offset *= 1.0 - (0.5 * f32(current_spiral + 1) / f32(num_spirals));
156
157        // Rotate and correct for aspect ratio
158        let rotated_spiral_offset = (rotation_matrix * spiral_offset) * vec2(1.0, aspect);
159
160        // Calculate final offset position, with blur and spiral offset
161        let modified_offset_position = offset_position + rotated_spiral_offset * blur_intensity * (1.0 - f32(pixel_checkboard) * 0.1);
162
163        // Sample the view transmission texture at the offset position + noise offset, to get the background color
164        var sample = textureSampleLevel(
165            view_bindings::view_transmission_texture,
166            view_bindings::view_transmission_sampler,
167            modified_offset_position,
168            0.0
169        );
170
171#ifdef DEPTH_PREPASS
172#ifndef WEBGL2
173        // Use depth prepass data to reject values that are in front of the current fragment
174        if prepass_utils::prepass_depth(vec4<f32>(modified_offset_position * view_bindings::view.viewport.zw, 0.0, 0.0), 0u) > frag_coord.z {
175            sample = vec4<f32>(0.0);
176        }
177#endif
178#endif
179
180        // As blur intensity grows higher, gradually limit *very bright* color RGB values towards a
181        // maximum length of 1.0 to prevent stray “firefly” pixel artifacts. This can potentially make
182        // very strong emissive meshes appear much dimmer, but the artifacts are noticeable enough to
183        // warrant this treatment.
184        let normalized_rgb = normalize(sample.rgb);
185        result += vec4(min(sample.rgb, normalized_rgb / saturate(blur_intensity / 2.0)), sample.a);
186    }
187
188    result /= f32(num_taps);
189
190#ifdef TONEMAP_IN_SHADER
191    result = approximate_inverse_tone_mapping(result, view_bindings::view.color_grading);
192#endif
193
194    return result;
195}