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