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