allow extensions to StandardMaterial

Cargo.toml

@@ -1710,6 +1710,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,52 @@
+#import bevy_pbr::pbr_fragment          pbr_input_from_standard_material
+#import bevy_pbr::pbr_functions         alpha_discard
+
+#ifdef DEFERRED_PREPASS
+#import bevy_pbr::prepass_io            VertexOutput, FragmentOutput
+#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 DEFERRED_PREPASS
+    // 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

@@ -65,28 +65,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 DEFERRED_PREPASS
+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

crates/bevy_pbr/src/extended_material.rs

@@ -0,0 +1,226 @@
+use bevy_asset::{Asset, Handle};
+use bevy_reflect::TypePath;
+use bevy_render::{
+    mesh::MeshVertexBufferLayout,
+    render_asset::RenderAssets,
+    render_resource::{
+        AsBindGroup, AsBindGroupError, BindGroupLayout, RenderPipelineDescriptor, Shader,
+        ShaderRef, SpecializedMeshPipelineError, UnpreparedBindGroup,
+    },
+    renderer::RenderDevice,
+    texture::{FallbackImage, Image},
+};
+
+use crate::{Material, MaterialPipeline, MaterialPipelineKey, MeshPipeline, MeshPipelineKey};
+
+pub struct MaterialExtensionPipeline {
+    pub mesh_pipeline: MeshPipeline,
+    pub material_layout: BindGroupLayout,
+    pub vertex_shader: Option<Handle<Shader>>,
+    pub fragment_shader: Option<Handle<Shader>>,
+}
+
+pub struct MaterialExtensionKey<E: MaterialExtension> {
+    pub mesh_key: MeshPipelineKey,
+    pub bind_group_data: E::Data,
+}
+
+/// A subset of the `Material` trait for defining extensions to a base `Material`, such as the builtin `StandardMaterial`.
+/// A user type implementing the trait should be used as the `E` generic param in an `ExtendedMaterial` struct.
+pub trait MaterialExtension: Asset + AsBindGroup + Clone + Sized {
+    /// Returns this material's vertex shader. If [`ShaderRef::Default`] is returned, the default mesh vertex shader
+    /// will be used.
+    fn vertex_shader() -> ShaderRef {
+        ShaderRef::Default
+    }
+
+    /// Returns this material's fragment shader. If [`ShaderRef::Default`] is returned, the default mesh fragment shader
+    /// will be used.
+    #[allow(unused_variables)]
+    fn fragment_shader() -> ShaderRef {
+        ShaderRef::Default
+    }
+
+    /// Returns this material's prepass vertex shader. If [`ShaderRef::Default`] is returned, the default prepass vertex shader
+    /// will be used.
+    fn prepass_vertex_shader() -> ShaderRef {
+        ShaderRef::Default
+    }
+
+    /// Returns this material's prepass fragment shader. If [`ShaderRef::Default`] is returned, the default prepass fragment shader
+    /// will be used.
+    #[allow(unused_variables)]
+    fn prepass_fragment_shader() -> ShaderRef {
+        ShaderRef::Default
+    }
+
+    /// Customizes the default [`RenderPipelineDescriptor`] for a specific entity using the entity's
+    /// [`MaterialPipelineKey`] and [`MeshVertexBufferLayout`] as input.
+    /// Specialization for the base material is applied before this function is called.
+    #[allow(unused_variables)]
+    #[inline]
+    fn specialize(
+        pipeline: &MaterialExtensionPipeline,
+        descriptor: &mut RenderPipelineDescriptor,
+        layout: &MeshVertexBufferLayout,
+        key: MaterialExtensionKey<Self>,
+    ) -> Result<(), SpecializedMeshPipelineError> {
+        Ok(())
+    }
+}
+
+/// A material that extends a base [`Material`] with additional shaders and data.
+///
+/// The data from both materials will be combined and made available to the shader
+/// so that shader functions built for the base material (and referencing the base material
+/// bindings) will work as expected, and custom alterations based on custom data can also be used.
+///
+/// If the extension `E` returns a non-default result from `vertex_shader()` it will be used in place of the base
+/// material's vertex shader.
+///
+/// If the extension `E` returns a non-default result from `fragment_shader()` it will be used in place of the base
+/// fragment shader.
+///
+/// When used with `StandardMaterial` as the base, all the standard material fields are
+/// present, so the `pbr_fragment` shader functions can be called from the extension shader (see
+/// the `extended_material` example).
+#[derive(Asset, Clone, TypePath)]
+pub struct ExtendedMaterial<B: Material, E: MaterialExtension> {
+    pub base: B,
+    pub extension: E,
+}
+
+impl<B: Material, E: MaterialExtension> AsBindGroup for ExtendedMaterial<B, E> {
+    type Data = (<B as AsBindGroup>::Data, <E as AsBindGroup>::Data);
+
+    fn unprepared_bind_group(
+        &self,
+        layout: &BindGroupLayout,
+        render_device: &RenderDevice,
+        images: &RenderAssets<Image>,
+        fallback_image: &FallbackImage,
+    ) -> Result<bevy_render::render_resource::UnpreparedBindGroup<Self::Data>, AsBindGroupError>
+    {
+        // add together the bindings of the base material and the user material
+        let UnpreparedBindGroup {
+            mut bindings,
+            data: base_data,
+        } = B::unprepared_bind_group(&self.base, layout, render_device, images, fallback_image)?;
+        let extended_bindgroup = E::unprepared_bind_group(
+            &self.extension,
+            layout,
+            render_device,
+            images,
+            fallback_image,
+        )?;
+
+        bindings.extend(extended_bindgroup.bindings);
+
+        Ok(UnpreparedBindGroup {
+            bindings,
+            data: (base_data, extended_bindgroup.data),
+        })
+    }
+
+    fn bind_group_layout_entries(
+        render_device: &RenderDevice,
+    ) -> Vec<bevy_render::render_resource::BindGroupLayoutEntry>
+    where
+        Self: Sized,
+    {
+        // add together the bindings of the standard material and the user material
+        let mut entries = B::bind_group_layout_entries(render_device);
+        entries.extend(E::bind_group_layout_entries(render_device));
+        entries
+    }
+}
+
+impl<B: Material, E: MaterialExtension> Material for ExtendedMaterial<B, E> {
+    fn vertex_shader() -> bevy_render::render_resource::ShaderRef {
+        match E::vertex_shader() {
+            ShaderRef::Default => B::vertex_shader(),
+            specified => specified,
+        }
+    }
+
+    fn fragment_shader() -> bevy_render::render_resource::ShaderRef {
+        match E::fragment_shader() {
+            ShaderRef::Default => B::fragment_shader(),
+            specified => specified,
+        }
+    }
+
+    fn alpha_mode(&self) -> crate::AlphaMode {
+        B::alpha_mode(&self.base)
+    }
+
+    fn depth_bias(&self) -> f32 {
+        B::depth_bias(&self.base)
+    }
+
+    fn prepass_vertex_shader() -> bevy_render::render_resource::ShaderRef {
+        match E::prepass_vertex_shader() {
+            ShaderRef::Default => B::prepass_vertex_shader(),
+            specified => specified,
+        }
+    }
+
+    fn prepass_fragment_shader() -> bevy_render::render_resource::ShaderRef {
+        match E::prepass_fragment_shader() {
+            ShaderRef::Default => B::prepass_fragment_shader(),
+            specified => specified,
+        }
+    }
+
+    fn specialize(
+        pipeline: &MaterialPipeline<Self>,
+        descriptor: &mut RenderPipelineDescriptor,
+        layout: &MeshVertexBufferLayout,
+        key: MaterialPipelineKey<Self>,
+    ) -> Result<(), SpecializedMeshPipelineError> {
+        // call StandardMaterial specialize function
+        let MaterialPipeline::<Self> {
+            mesh_pipeline,
+            material_layout,
+            vertex_shader,
+            fragment_shader,
+            ..
+        } = pipeline.clone();
+        let base_pipeline = MaterialPipeline::<B> {
+            mesh_pipeline,
+            material_layout,
+            vertex_shader,
+            fragment_shader,
+            marker: Default::default(),
+        };
+        let base_key = MaterialPipelineKey::<B> {
+            mesh_key: key.mesh_key,
+            bind_group_data: key.bind_group_data.0,
+        };
+        B::specialize(&base_pipeline, descriptor, layout, base_key)?;
+
+        // call custom material specialize function afterwards
+        let MaterialPipeline::<Self> {
+            mesh_pipeline,
+            material_layout,
+            vertex_shader,
+            fragment_shader,
+            ..
+        } = pipeline.clone();
+
+        E::specialize(
+            &MaterialExtensionPipeline {
+                mesh_pipeline,
+                material_layout,
+                vertex_shader,
+                fragment_shader,
+            },
+            descriptor,
+            layout,
+            MaterialExtensionKey {
+                mesh_key: key.mesh_key,
+                bind_group_data: key.bind_group_data.1,
+            },
+        )
+    }
+}