allow extensions to StandardMaterial (#7820)

# Objective

allow extending `Material`s (including the built in `StandardMaterial`)
with custom vertex/fragment shaders and additional data, to easily get
pbr lighting with custom modifications, or otherwise extend a base
material.

# Solution

- added `ExtendedMaterial<B: Material, E: MaterialExtension>` which
contains a base material and a user-defined extension.
- added example `extended_material` showing how to use it
- modified AsBindGroup to have "unprepared" functions that return raw
resources / layout entries so that the extended material can combine
them

note: doesn't currently work with array resources, as i can't figure out
how to make the OwnedBindingResource::get_binding() work, as wgpu
requires a `&'a[&'a TextureView]` and i have a `Vec<TextureView>`.

# Migration Guide

manual implementations of `AsBindGroup` will need to be adjusted, the
changes are pretty straightforward and can be seen in the diff for e.g.
the `texture_binding_array` example.

---------

Co-authored-by: Robert Swain <[email protected]>

Author: robtfm

Cargo.toml

@@ -1720,6 +1720,16 @@ description = "A shader and a material that uses it"
 category = "Shaders"
 wasm = true
 
+[[example]]
+name = "extended_material"
+path = "examples/shader/extended_material.rs"
+
+[package.metadata.example.extended_material]
+name = "Extended Material"
+description = "A custom shader that builds on the standard material"
+category = "Shaders"
+wasm = true
+
 [[example]]
 name = "shader_prepass"
 path = "examples/shader/shader_prepass.rs"

assets/shaders/extended_material.wgsl

@@ -0,0 +1,53 @@
+#import bevy_pbr::pbr_fragment          pbr_input_from_standard_material
+#import bevy_pbr::pbr_functions         alpha_discard
+
+#ifdef PREPASS_PIPELINE
+#import bevy_pbr::prepass_io            VertexOutput, FragmentOutput
+#import bevy_pbr::pbr_deferred_functions  deferred_output
+#else
+#import bevy_pbr::forward_io            VertexOutput, FragmentOutput
+#import bevy_pbr::pbr_functions         apply_pbr_lighting, main_pass_post_lighting_processing
+#endif
+
+struct MyExtendedMaterial {
+    quantize_steps: u32,
+}
+
+@group(1) @binding(100)
+var<uniform> my_extended_material: MyExtendedMaterial;
+
+@fragment
+fn fragment(
+    in: VertexOutput,
+    @builtin(front_facing) is_front: bool,
+) -> FragmentOutput {
+    // generate a PbrInput struct from the StandardMaterial bindings
+    var pbr_input = pbr_input_from_standard_material(in, is_front);
+
+    // we can optionally modify the input before lighting and alpha_discard is applied
+    pbr_input.material.base_color.b = pbr_input.material.base_color.r;
+
+    // alpha discard
+    pbr_input.material.base_color = alpha_discard(pbr_input.material, pbr_input.material.base_color);
+
+#ifdef PREPASS_PIPELINE
+    // in deferred mode we can't modify anything after that, as lighting is run in a separate fullscreen shader.
+    let out = deferred_output(in, pbr_input);
+#else
+    var out: FragmentOutput;
+    // apply lighting
+    out.color = apply_pbr_lighting(pbr_input);
+
+    // we can optionally modify the lit color before post-processing is applied
+    out.color = vec4<f32>(vec4<u32>(out.color * f32(my_extended_material.quantize_steps))) / f32(my_extended_material.quantize_steps);
+
+    // apply in-shader post processing (fog, alpha-premultiply, and also tonemapping, debanding if the camera is non-hdr)
+    // note this does not include fullscreen postprocessing effects like bloom.
+    out.color = main_pass_post_lighting_processing(pbr_input, out.color);
+
+    // we can optionally modify the final result here
+    out.color = out.color * 2.0;
+#endif
+
+    return out;
+}

crates/bevy_pbr/src/deferred/deferred_lighting.wgsl

@@ -67,28 +67,12 @@ fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
         pbr_input.occlusion = min(pbr_input.occlusion, ssao_multibounce);
 #endif // SCREEN_SPACE_AMBIENT_OCCLUSION
 
-        output_color = pbr_functions::pbr(pbr_input);
+        output_color = pbr_functions::apply_pbr_lighting(pbr_input);
     } else {
         output_color = pbr_input.material.base_color;
     }
 
-    // fog
-    if (fog.mode != FOG_MODE_OFF && (pbr_input.material.flags & STANDARD_MATERIAL_FLAGS_FOG_ENABLED_BIT) != 0u) {
-        output_color = pbr_functions::apply_fog(fog, output_color, pbr_input.world_position.xyz, view.world_position.xyz);
-    }
-
-#ifdef TONEMAP_IN_SHADER
-    output_color = tone_mapping(output_color, view.color_grading);
-#ifdef DEBAND_DITHER
-    var output_rgb = output_color.rgb;
-    output_rgb = powsafe(output_rgb, 1.0 / 2.2);
-    output_rgb = output_rgb + screen_space_dither(frag_coord.xy);
-    // This conversion back to linear space is required because our output texture format is
-    // SRGB; the GPU will assume our output is linear and will apply an SRGB conversion.
-    output_rgb = powsafe(output_rgb, 2.2);
-    output_color = vec4(output_rgb, output_color.a);
-#endif
-#endif
+    output_color = pbr_functions::main_pass_post_lighting_processing(pbr_input, output_color);
 
     return output_color;
 }

crates/bevy_pbr/src/deferred/pbr_deferred_functions.wgsl

@@ -1,11 +1,17 @@
 #define_import_path bevy_pbr::pbr_deferred_functions
+
 #import bevy_pbr::pbr_types PbrInput, standard_material_new, STANDARD_MATERIAL_FLAGS_FOG_ENABLED_BIT, STANDARD_MATERIAL_FLAGS_UNLIT_BIT
 #import bevy_pbr::pbr_deferred_types as deferred_types
 #import bevy_pbr::pbr_functions as pbr_functions
 #import bevy_pbr::rgb9e5 as rgb9e5
 #import bevy_pbr::mesh_view_bindings as view_bindings
 #import bevy_pbr::mesh_view_bindings view
 #import bevy_pbr::utils octahedral_encode, octahedral_decode
+#import bevy_pbr::prepass_io VertexOutput, FragmentOutput
+
+#ifdef MOTION_VECTOR_PREPASS
+    #import bevy_pbr::pbr_prepass_functions calculate_motion_vector
+#endif
 
 // ---------------------------
 // from https://github.com/DGriffin91/bevy_coordinate_systems/blob/main/src/transformations.wgsl
@@ -126,4 +132,23 @@ fn pbr_input_from_deferred_gbuffer(frag_coord: vec4<f32>, gbuffer: vec4<u32>) ->
     return pbr;
 }
 
+#ifdef PREPASS_PIPELINE
+fn deferred_output(in: VertexOutput, pbr_input: PbrInput) -> FragmentOutput {
+    var out: FragmentOutput;
 
+    // gbuffer
+    out.deferred = deferred_gbuffer_from_pbr_input(pbr_input);
+    // lighting pass id (used to determine which lighting shader to run for the fragment)
+    out.deferred_lighting_pass_id = pbr_input.material.deferred_lighting_pass_id;
+    // normal if required
+#ifdef NORMAL_PREPASS
+    out.normal = vec4(in.world_normal * 0.5 + vec3(0.5), 1.0);
+#endif
+    // motion vectors if required
+#ifdef MOTION_VECTOR_PREPASS
+    out.motion_vector = calculate_motion_vector(in.world_position, in.previous_world_position);
+#endif
+
+    return out;
+}
+#endif