add ambient lighting hook (#5428)

# Objective

add a hook for ambient occlusion to the pbr shader

## Solution

add a hook for ambient occlusion to the pbr shader


Co-authored-by: atlas dostal <[email protected]>

Author: robtfm

crates/bevy_pbr/src/lib.rs

@@ -75,6 +75,8 @@ pub const PBR_PREPASS_SHADER_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 9407115064344201137);
 pub const PBR_FUNCTIONS_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 16550102964439850292);
+pub const PBR_AMBIENT_HANDLE: HandleUntyped =
+    HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 2441520459096337034);
 pub const SHADOW_SHADER_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 1836745567947005696);
 
@@ -132,6 +134,12 @@ impl Plugin for PbrPlugin {
             "render/pbr_functions.wgsl",
             Shader::from_wgsl
         );
+        load_internal_asset!(
+            app,
+            PBR_AMBIENT_HANDLE,
+            "render/pbr_ambient.wgsl",
+            Shader::from_wgsl
+        );
         load_internal_asset!(app, PBR_SHADER_HANDLE, "render/pbr.wgsl", Shader::from_wgsl);
         load_internal_asset!(
             app,

crates/bevy_pbr/src/render/pbr.wgsl

@@ -5,6 +5,7 @@
 #import bevy_pbr::utils
 #import bevy_pbr::clustered_forward
 #import bevy_pbr::lighting
+#import bevy_pbr::pbr_ambient
 #import bevy_pbr::shadows
 #import bevy_pbr::fog
 #import bevy_pbr::pbr_functions
@@ -66,8 +67,6 @@ fn fragment(in: FragmentInput) -> @location(0) vec4<f32> {
             occlusion = textureSample(occlusion_texture, occlusion_sampler, in.uv).r;
         }
 #endif
-        pbr_input.occlusion = occlusion;
-
         pbr_input.frag_coord = in.frag_coord;
         pbr_input.world_position = in.world_position;
         pbr_input.world_normal = prepare_world_normal(
@@ -91,6 +90,8 @@ fn fragment(in: FragmentInput) -> @location(0) vec4<f32> {
 #endif
         );
         pbr_input.V = calculate_view(in.world_position, pbr_input.is_orthographic);
+        pbr_input.occlusion = occlusion;
+
         output_color = pbr(pbr_input);
     } else {
         output_color = alpha_discard(material, output_color);

crates/bevy_pbr/src/render/pbr_ambient.wgsl

@@ -0,0 +1,19 @@
+#define_import_path bevy_pbr::pbr_ambient
+
+// A precomputed `NdotV` is provided because it is computed regardless,
+// but `world_normal` and the view vector `V` are provided separately for more advanced uses.
+fn ambient_light(
+    world_position: vec4<f32>,
+    world_normal: vec3<f32>,
+    V: vec3<f32>,
+    NdotV: f32,
+    diffuse_color: vec3<f32>,
+    specular_color: vec3<f32>,
+    perceptual_roughness: f32,
+    occlusion: f32,
+) -> vec3<f32> {
+    let diffuse_ambient = EnvBRDFApprox(diffuse_color, 1.0, NdotV);
+    let specular_ambient = EnvBRDFApprox(specular_color, perceptual_roughness, NdotV);
+
+    return (diffuse_ambient + specular_ambient) * lights.ambient_color.rgb * occlusion;
+}

crates/bevy_pbr/src/render/pbr_functions.wgsl

@@ -233,13 +233,10 @@ fn pbr(
         light_accum = light_accum + light_contrib * shadow;
     }
 
-    let diffuse_ambient = EnvBRDFApprox(diffuse_color, 1.0, NdotV);
-    let specular_ambient = EnvBRDFApprox(F0, perceptual_roughness, NdotV);
+    let ambient_contrib = ambient_light(in.world_position, in.N, in.V, NdotV, diffuse_color, F0, perceptual_roughness, occlusion);
 
     output_color = vec4<f32>(
-        light_accum
-            + (diffuse_ambient + specular_ambient) * lights.ambient_color.rgb * occlusion
-            + emissive.rgb * output_color.a,
+        light_accum + ambient_contrib + emissive.rgb * output_color.a,
         output_color.a
     );