Clean up light transmission shader implementation with more proper math/parameters

crates/bevy_pbr/src/render/pbr_functions.wgsl

@@ -171,6 +171,15 @@ fn pbr(
     let metallic = in.material.metallic;
     let perceptual_roughness = in.material.perceptual_roughness;
     let roughness = perceptualRoughnessToRoughness(perceptual_roughness);
+    let ior = 1.5;
+    let thickness = 1.0;
+    var transmission = 0.0;
+    let alpha_mode = in.material.flags & STANDARD_MATERIAL_FLAGS_ALPHA_MODE_RESERVED_BITS;
+    if (alpha_mode == STANDARD_MATERIAL_FLAGS_ALPHA_MODE_BLEND) {
+        transmission = 1.0;
+    }
+
+    let transmissive_color = transmission * in.material.base_color.rgb;
 
     let occlusion = in.occlusion;
 
@@ -184,11 +193,10 @@ fn pbr(
     let reflectance = in.material.reflectance;
     let F0 = 0.16 * reflectance * reflectance * (1.0 - metallic) + output_color.rgb * metallic;
 
-    // Diffuse strength inversely related to metallicity
+    // Diffuse strength inversely related to metallicity and transmission
-    let diffuse_color = output_color.rgb * (1.0 - metallic);
+    let diffuse_color = output_color.rgb * (1.0 - metallic) * (1.0 - transmission);
 
     let R = reflect(-in.V, in.N);
-    let R2 = refract(-in.V, in.N, 0.6);
 
     let f_ab = F_AB(perceptual_roughness, NdotV);
 
@@ -253,11 +261,12 @@ fn pbr(
 
     let emissive_light = emissive.rgb * output_color.a;
 
+    // Transmitted Light
     var transmitted_light: vec3<f32> = vec3<f32>(0.0);
-    let alpha_mode = in.material.flags & STANDARD_MATERIAL_FLAGS_ALPHA_MODE_RESERVED_BITS;
+    if transmission > 0.0 {
-    if alpha_mode == STANDARD_MATERIAL_FLAGS_ALPHA_MODE_BLEND {
+        transmitted_light = transmissive_light(in.frag_coord.xyz, in.N, in.V, ior, thickness, transmissive_color).rgb;
-        transmitted_light = transmissive_light(in.frag_coord.xy, R2).rgb * in.material.base_color.rgb;
+    } else {
-        direct_light = vec3<f32>(0.0);
+        transmitted_light = vec3<f32>(0.0);
     }
 
     // Total light
@@ -278,12 +287,39 @@ fn pbr(
 }
 #endif // NORMAL_PREPASS
 
-fn transmissive_light(frag_coord: vec2<f32>, N: vec3<f32>) -> vec3<f32> {
+fn transmissive_light(frag_coord: vec3<f32>, N: vec3<f32>, V: vec3<f32>, ior: f32, thickness: f32, transmissive_color: vec3<f32>) -> vec3<f32> {
-    return textureSample(
+    // Calculate view aspect ratio, used later to scale offset so that it's proportionate
+    let aspect = view.viewport.z / view.viewport.w;
+
+    // Calculate the ratio between refaction indexes. Assume air/vacuum for the space outside the mesh
+    let eta = 1.0 / ior;
+
+    // Calculate incidence vector (opposite to view vector) and its dot product with the mesh normal
+    let I = -V;
+    let NdotI = dot(N, I);
+
+    // Calculate refracted direction using Snell's law
+    let k = 1.0 - eta * eta * (1.0 - NdotI * NdotI);
+    let T = eta * I - (eta * NdotI + sqrt(k)) * N;
+
+    // Transform refracted direction into view space
+    let view_T = (view.inverse_view * vec4<f32>(T.xyz, 0.0)).xyz;
+
+    // Calculate an offset position, by multiplying the refracted direction by the thickness and adding it to
+    // the fragment position scaled by viewport size. Use the aspect ratio calculated earlier stay proportionate
+    let offset_position = frag_coord.xy / view.viewport.zw + view_T.xy * thickness * vec2<f32>(1.0, -aspect);
+
+    // Sample the view main texture at the offset position, to get the background color
+    // TODO: use depth prepass data to reject values that are in front of the current fragment
+    let background_sample = textureSample(
         view_main_texture,
         view_main_sampler,
-        frag_coord.xy / view.viewport.zw + N.xy * 0.05,
+        offset_position,
     ).rgb;
+
+    // Calculate final color by applying transmissive color to background sample
+    // TODO: Add support for attenuationColor and attenuationDistance
+    return transmissive_color * background_sample;
 }
 
 #ifdef DEBAND_DITHER