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 <robert.swain@gmail.com>

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,4 +1,5 @@
 #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
@@ -6,6 +7,11 @@
 #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

crates/bevy_pbr/src/extended_material.rs

@@ -0,0 +1,257 @@
+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 base material mesh vertex shader
+    /// will be used.
+    fn vertex_shader() -> ShaderRef {
+        ShaderRef::Default
+    }
+
+    /// Returns this material's fragment shader. If [`ShaderRef::Default`] is returned, the base material 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 base material 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 base material prepass fragment shader
+    /// will be used.
+    #[allow(unused_variables)]
+    fn prepass_fragment_shader() -> ShaderRef {
+        ShaderRef::Default
+    }
+
+    /// Returns this material's deferred vertex shader. If [`ShaderRef::Default`] is returned, the base material deferred vertex shader
+    /// will be used.
+    fn deferred_vertex_shader() -> ShaderRef {
+        ShaderRef::Default
+    }
+
+    /// Returns this material's prepass fragment shader. If [`ShaderRef::Default`] is returned, the base material deferred fragment shader
+    /// will be used.
+    #[allow(unused_variables)]
+    fn deferred_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 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 deferred_vertex_shader() -> bevy_render::render_resource::ShaderRef {
+        match E::deferred_vertex_shader() {
+            ShaderRef::Default => B::deferred_vertex_shader(),
+            specified => specified,
+        }
+    }
+
+    fn deferred_fragment_shader() -> bevy_render::render_resource::ShaderRef {
+        match E::deferred_fragment_shader() {
+            ShaderRef::Default => B::deferred_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 opaque_render_method(&self) -> crate::OpaqueRendererMethod {
+        B::opaque_render_method(&self.base)
+    }
+
+    fn specialize(
+        pipeline: &MaterialPipeline<Self>,
+        descriptor: &mut RenderPipelineDescriptor,
+        layout: &MeshVertexBufferLayout,
+        key: MaterialPipelineKey<Self>,
+    ) -> Result<(), SpecializedMeshPipelineError> {
+        // Call the base material's 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 the extended material's 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,
+            },
+        )
+    }
+}

crates/bevy_pbr/src/lib.rs

@@ -6,6 +6,7 @@ mod alpha;
 mod bundle;
 pub mod deferred;
 mod environment_map;
+mod extended_material;
 mod fog;
 mod light;
 mod material;
@@ -18,6 +19,7 @@ mod ssao;
 pub use alpha::*;
 pub use bundle::*;
 pub use environment_map::EnvironmentMapLight;
+pub use extended_material::*;
 pub use fog::*;
 pub use light::*;
 pub use material::*;
@@ -73,6 +75,7 @@ pub const CLUSTERED_FORWARD_HANDLE: Handle<Shader> = Handle::weak_from_u128(1668
 pub const PBR_LIGHTING_HANDLE: Handle<Shader> = Handle::weak_from_u128(14170772752254856967);
 pub const SHADOWS_HANDLE: Handle<Shader> = Handle::weak_from_u128(11350275143789590502);
 pub const SHADOW_SAMPLING_HANDLE: Handle<Shader> = Handle::weak_from_u128(3145627513789590502);
+pub const PBR_FRAGMENT_HANDLE: Handle<Shader> = Handle::weak_from_u128(2295049283805286543);
 pub const PBR_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(4805239651767701046);
 pub const PBR_PREPASS_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(9407115064344201137);
 pub const PBR_FUNCTIONS_HANDLE: Handle<Shader> = Handle::weak_from_u128(16550102964439850292);
@@ -172,6 +175,12 @@ impl Plugin for PbrPlugin {
             "render/pbr_ambient.wgsl",
             Shader::from_wgsl
         );
+        load_internal_asset!(
+            app,
+            PBR_FRAGMENT_HANDLE,
+            "render/pbr_fragment.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/material.rs

@@ -297,7 +297,7 @@ pub struct MaterialPipeline<M: Material> {
     pub material_layout: BindGroupLayout,
     pub vertex_shader: Option<Handle<Shader>>,
     pub fragment_shader: Option<Handle<Shader>>,
-    marker: PhantomData<M>,
+    pub marker: PhantomData<M>,
 }
 
 impl<M: Material> Clone for MaterialPipeline<M> {
@@ -693,7 +693,7 @@ pub struct MaterialProperties {
 
 /// Data prepared for a [`Material`] instance.
 pub struct PreparedMaterial<T: Material> {
-    pub bindings: Vec<OwnedBindingResource>,
+    pub bindings: Vec<(u32, OwnedBindingResource)>,
     pub bind_group: BindGroup,
     pub key: T::Data,
     pub properties: MaterialProperties,

crates/bevy_pbr/src/render/pbr.wgsl

@@ -1,228 +1,43 @@
-#define_import_path bevy_pbr::fragment
-
-#import bevy_pbr::pbr_functions as pbr_functions
-#import bevy_pbr::pbr_bindings as pbr_bindings
-#import bevy_pbr::pbr_types as pbr_types
-
-#import bevy_pbr::mesh_bindings            mesh
-#import bevy_pbr::mesh_view_bindings       view, fog, screen_space_ambient_occlusion_texture
-#import bevy_pbr::mesh_view_types          FOG_MODE_OFF
-#import bevy_core_pipeline::tonemapping    screen_space_dither, powsafe, tone_mapping
-#import bevy_pbr::parallax_mapping         parallaxed_uv
-
-#import bevy_pbr::prepass_utils
-
-#ifdef SCREEN_SPACE_AMBIENT_OCCLUSION
-#import bevy_pbr::gtao_utils gtao_multibounce
-#endif
+#import bevy_pbr::pbr_functions           alpha_discard
+#import bevy_pbr::pbr_fragment            pbr_input_from_standard_material
 
 #ifdef PREPASS_PIPELINE
-#import bevy_pbr::pbr_deferred_functions deferred_gbuffer_from_pbr_input
-#import bevy_pbr::pbr_prepass_functions calculate_motion_vector
-#import bevy_pbr::prepass_io VertexOutput, FragmentOutput
-#else // PREPASS_PIPELINE
-#import bevy_pbr::forward_io VertexOutput, FragmentOutput
-#endif // PREPASS_PIPELINE
-
-#ifdef MOTION_VECTOR_PREPASS
-@group(0) @binding(2)
-var<uniform> previous_view_proj: mat4x4<f32>;
-#endif // MOTION_VECTOR_PREPASS
+#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
+#import bevy_pbr::pbr_types               STANDARD_MATERIAL_FLAGS_UNLIT_BIT
+#endif
 
 @fragment
 fn fragment(
     in: VertexOutput,
     @builtin(front_facing) is_front: bool,
 ) -> FragmentOutput {
-    var out: FragmentOutput;
+    // generate a PbrInput struct from the StandardMaterial bindings
+    var pbr_input = pbr_input_from_standard_material(in, is_front);
 
-    // calculate unlit color
-    // ---------------------
-    var unlit_color: vec4<f32> = pbr_bindings::material.base_color;
+    // alpha discard
+    pbr_input.material.base_color = alpha_discard(pbr_input.material, pbr_input.material.base_color);
 
-    let is_orthographic = view.projection[3].w == 1.0;
-    let V = pbr_functions::calculate_view(in.world_position, is_orthographic);
-#ifdef VERTEX_UVS
-    var uv = in.uv;
-#ifdef VERTEX_TANGENTS
-    if ((pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_DEPTH_MAP_BIT) != 0u) {
-        let N = in.world_normal;
-        let T = in.world_tangent.xyz;
-        let B = in.world_tangent.w * cross(N, T);
-        // Transform V from fragment to camera in world space to tangent space.
-        let Vt = vec3(dot(V, T), dot(V, B), dot(V, N));
-        uv = parallaxed_uv(
-            pbr_bindings::material.parallax_depth_scale,
-            pbr_bindings::material.max_parallax_layer_count,
-            pbr_bindings::material.max_relief_mapping_search_steps,
-            uv,
-            // Flip the direction of Vt to go toward the surface to make the
-            // parallax mapping algorithm easier to understand and reason
-            // about.
-            -Vt,
-        );
-    }
-#endif
-#endif
-
-#ifdef VERTEX_COLORS
-    unlit_color = unlit_color * in.color;
-#endif
-#ifdef VERTEX_UVS
-    if ((pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_BASE_COLOR_TEXTURE_BIT) != 0u) {
-        unlit_color = unlit_color * textureSampleBias(pbr_bindings::base_color_texture, pbr_bindings::base_color_sampler, uv, view.mip_bias);
-    }
-#endif
-
-    // gather pbr lighting data
-    // ------------------
-    var pbr_input: pbr_types::PbrInput;
-    // NOTE: Unlit bit not set means == 0 is true, so the true case is if lit
-    if ((pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_UNLIT_BIT) == 0u) {
-        // Prepare a 'processed' StandardMaterial by sampling all textures to resolve
-        // the material members
-
-        pbr_input.material.reflectance = pbr_bindings::material.reflectance;
-        pbr_input.material.flags = pbr_bindings::material.flags;
-        pbr_input.material.alpha_cutoff = pbr_bindings::material.alpha_cutoff;
-        pbr_input.frag_coord = in.position;
-        pbr_input.world_position = in.world_position;
-        pbr_input.is_orthographic = is_orthographic;
-        pbr_input.flags = mesh[in.instance_index].flags;
-
-        // emmissive
-        // TODO use .a for exposure compensation in HDR
-        var emissive: vec4<f32> = pbr_bindings::material.emissive;
-#ifdef VERTEX_UVS
-        if ((pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_EMISSIVE_TEXTURE_BIT) != 0u) {
-            emissive = vec4<f32>(emissive.rgb * textureSampleBias(pbr_bindings::emissive_texture, pbr_bindings::emissive_sampler, uv, view.mip_bias).rgb, 1.0);
-        }
-#endif
-        pbr_input.material.emissive = emissive;
-
-        // metallic and perceptual roughness
-        var metallic: f32 = pbr_bindings::material.metallic;
-        var perceptual_roughness: f32 = pbr_bindings::material.perceptual_roughness;
-#ifdef VERTEX_UVS
-        if ((pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_METALLIC_ROUGHNESS_TEXTURE_BIT) != 0u) {
-            let metallic_roughness = textureSampleBias(pbr_bindings::metallic_roughness_texture, pbr_bindings::metallic_roughness_sampler, uv, view.mip_bias);
-            // Sampling from GLTF standard channels for now
-            metallic = metallic * metallic_roughness.b;
-            perceptual_roughness = perceptual_roughness * metallic_roughness.g;
-        }
-#endif
-        pbr_input.material.metallic = metallic;
-        pbr_input.material.perceptual_roughness = perceptual_roughness;
-
-        // occlusion
-        // TODO: Split into diffuse/specular occlusion?
-        var occlusion: vec3<f32> = vec3(1.0);
-#ifdef VERTEX_UVS
-        if ((pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_OCCLUSION_TEXTURE_BIT) != 0u) {
-            occlusion = vec3(textureSampleBias(pbr_bindings::occlusion_texture, pbr_bindings::occlusion_sampler, uv, view.mip_bias).r);
-        }
-#endif
-#ifdef SCREEN_SPACE_AMBIENT_OCCLUSION
-        let ssao = textureLoad(screen_space_ambient_occlusion_texture, vec2<i32>(in.position.xy), 0i).r;
-        let ssao_multibounce = gtao_multibounce(ssao, unlit_color.rgb);
-        occlusion = min(occlusion, ssao_multibounce);
-#endif
-        pbr_input.occlusion = occlusion;
-
-        // world normal
-        pbr_input.world_normal = pbr_functions::prepare_world_normal(
-            in.world_normal,
-            (pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_DOUBLE_SIDED_BIT) != 0u,
-            is_front,
-        );
-
-        // N (normal vector)
-#ifdef LOAD_PREPASS_NORMALS
-        pbr_input.N = bevy_pbr::prepass_utils::prepass_normal(in.position, 0u);
+#ifdef PREPASS_PIPELINE
+    // write the gbuffer, lighting pass id, and optionally normal and motion_vector textures
+    let out = deferred_output(in, pbr_input);
 #else
-        pbr_input.N = pbr_functions::apply_normal_mapping(
-            pbr_bindings::material.flags,
-            pbr_input.world_normal,
-#ifdef VERTEX_TANGENTS
-#ifdef STANDARDMATERIAL_NORMAL_MAP
-            in.world_tangent,
-#endif
-#endif
-#ifdef VERTEX_UVS
-            uv,
-#endif
-            view.mip_bias,
-        );
-#endif
-
-        // V (view vector)
-        pbr_input.V = V;
-
-    } else { // if UNLIT_BIT != 0
-#ifdef PREPASS_PIPELINE
-        // in deferred mode, we need to fill some of the pbr input data even for unlit materials
-        // to pass through the gbuffer to the deferred lighting shader        
-        pbr_input = pbr_types::pbr_input_new();
-        pbr_input.flags = mesh[in.instance_index].flags;
-        pbr_input.material.flags = pbr_bindings::material.flags;
-        pbr_input.world_position = in.world_position;
-        pbr_input.world_normal = in.world_normal;
-        pbr_input.frag_coord = in.position;
-#endif
-    }
-
-    // apply alpha discard
-    // -------------------
-    // note even though this is based on the unlit color, it must be done after all texture samples for uniform control flow
-    unlit_color = pbr_functions::alpha_discard(pbr_bindings::material, unlit_color);
-    pbr_input.material.base_color = unlit_color;
-
-    // generate output
-    // ---------------
-#ifdef PREPASS_PIPELINE
-    // write the gbuffer
-    out.deferred = deferred_gbuffer_from_pbr_input(pbr_input);
-    out.deferred_lighting_pass_id = pbr_bindings::material.deferred_lighting_pass_id;
-#ifdef NORMAL_PREPASS
-    out.normal = vec4(in.world_normal * 0.5 + vec3(0.5), 1.0);
-#endif
-#ifdef MOTION_VECTOR_PREPASS
-    out.motion_vector = calculate_motion_vector(in.world_position, in.previous_world_position);
-#endif // MOTION_VECTOR_PREPASS
-
-#else // PREPASS_PIPELINE
-
     // in forward mode, we calculate the lit color immediately, and then apply some post-lighting effects here.
     // in deferred mode the lit color and these effects will be calculated in the deferred lighting shader
-    var output_color = unlit_color;
-    if ((pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_UNLIT_BIT) == 0u) {
-        output_color = pbr_functions::pbr(pbr_input);
+    var out: FragmentOutput;
+    if (pbr_input.material.flags & STANDARD_MATERIAL_FLAGS_UNLIT_BIT) == 0u {
+        out.color = apply_pbr_lighting(pbr_input);
+    } else {
+        out.color = pbr_input.material.base_color;
     }
 
-    if (fog.mode != FOG_MODE_OFF && (pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_FOG_ENABLED_BIT) != 0u) {
-        output_color = pbr_functions::apply_fog(fog, output_color, in.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(in.position.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);
+    // 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);
 #endif
-#endif
-#ifdef PREMULTIPLY_ALPHA
-    output_color = pbr_functions::premultiply_alpha(pbr_bindings::material.flags, output_color);
-#endif
-
-    // write the final pixel color
-    out.color = output_color;
-
-#endif // PREPASS_PIPELINE
 
     return out;
 }

crates/bevy_pbr/src/render/pbr_fragment.wgsl

@@ -0,0 +1,163 @@
+#define_import_path bevy_pbr::pbr_fragment
+
+#import bevy_pbr::pbr_functions as pbr_functions
+#import bevy_pbr::pbr_bindings as pbr_bindings
+#import bevy_pbr::pbr_types as pbr_types
+#import bevy_pbr::prepass_utils
+
+#import bevy_pbr::mesh_bindings            mesh
+#import bevy_pbr::mesh_view_bindings       view, screen_space_ambient_occlusion_texture
+#import bevy_pbr::parallax_mapping         parallaxed_uv
+
+#ifdef SCREEN_SPACE_AMBIENT_OCCLUSION
+#import bevy_pbr::gtao_utils gtao_multibounce
+#endif
+
+#ifdef PREPASS_PIPELINE
+#import bevy_pbr::prepass_io VertexOutput
+#else
+#import bevy_pbr::forward_io VertexOutput
+#endif
+
+// prepare a basic PbrInput from the vertex stage output, mesh binding and view binding
+fn pbr_input_from_vertex_output(
+    in: VertexOutput,
+    is_front: bool,
+    double_sided: bool,
+) -> pbr_types::PbrInput {
+    var pbr_input: pbr_types::PbrInput = pbr_types::pbr_input_new();
+
+    pbr_input.flags = mesh[in.instance_index].flags;
+    pbr_input.is_orthographic = view.projection[3].w == 1.0;
+    pbr_input.V = pbr_functions::calculate_view(in.world_position, pbr_input.is_orthographic);
+    pbr_input.frag_coord = in.position;
+    pbr_input.world_position = in.world_position;
+
+#ifdef VERTEX_COLORS
+    pbr_input.material.base_color = in.color;
+#endif
+
+    pbr_input.world_normal = pbr_functions::prepare_world_normal(
+        in.world_normal,
+        double_sided,
+        is_front,
+    );
+
+#ifdef LOAD_PREPASS_NORMALS
+    pbr_input.N = bevy_pbr::prepass_utils::prepass_normal(in.position, 0u);
+#else
+    pbr_input.N = normalize(pbr_input.world_normal);
+#endif
+
+    return pbr_input;
+}
+
+// Prepare a full PbrInput by sampling all textures to resolve
+// the material members
+fn pbr_input_from_standard_material(
+    in: VertexOutput,
+    is_front: bool,
+) -> pbr_types::PbrInput {
+    let double_sided = (pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_DOUBLE_SIDED_BIT) != 0u;
+
+    var pbr_input: pbr_types::PbrInput = pbr_input_from_vertex_output(in, is_front, double_sided);
+    pbr_input.material.flags = pbr_bindings::material.flags;
+    pbr_input.material.base_color *= pbr_bindings::material.base_color;
+    pbr_input.material.deferred_lighting_pass_id = pbr_bindings::material.deferred_lighting_pass_id;
+
+#ifdef VERTEX_UVS
+    var uv = in.uv;
+
+#ifdef VERTEX_TANGENTS
+    if ((pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_DEPTH_MAP_BIT) != 0u) {
+        let V = pbr_input.V;
+        let N = in.world_normal;
+        let T = in.world_tangent.xyz;
+        let B = in.world_tangent.w * cross(N, T);
+        // Transform V from fragment to camera in world space to tangent space.
+        let Vt = vec3(dot(V, T), dot(V, B), dot(V, N));
+        uv = parallaxed_uv(
+            pbr_bindings::material.parallax_depth_scale,
+            pbr_bindings::material.max_parallax_layer_count,
+            pbr_bindings::material.max_relief_mapping_search_steps,
+            uv,
+            // Flip the direction of Vt to go toward the surface to make the
+            // parallax mapping algorithm easier to understand and reason
+            // about.
+            -Vt,
+        );
+    }
+#endif // VERTEX_TANGENTS
+
+    if ((pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_BASE_COLOR_TEXTURE_BIT) != 0u) {
+        pbr_input.material.base_color *= textureSampleBias(pbr_bindings::base_color_texture, pbr_bindings::base_color_sampler, uv, view.mip_bias);
+    }
+#endif // VERTEX_UVS
+
+    pbr_input.material.flags = pbr_bindings::material.flags;
+
+    // NOTE: Unlit bit not set means == 0 is true, so the true case is if lit
+    if ((pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_UNLIT_BIT) == 0u) {
+
+        pbr_input.material.reflectance = pbr_bindings::material.reflectance;
+        pbr_input.material.alpha_cutoff = pbr_bindings::material.alpha_cutoff;
+
+        // emissive       
+        // TODO use .a for exposure compensation in HDR
+        var emissive: vec4<f32> = pbr_bindings::material.emissive;
+#ifdef VERTEX_UVS
+        if ((pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_EMISSIVE_TEXTURE_BIT) != 0u) {
+            emissive = vec4<f32>(emissive.rgb * textureSampleBias(pbr_bindings::emissive_texture, pbr_bindings::emissive_sampler, uv, view.mip_bias).rgb, 1.0);
+        }
+#endif
+        pbr_input.material.emissive = emissive;
+
+        // metallic and perceptual roughness
+        var metallic: f32 = pbr_bindings::material.metallic;
+        var perceptual_roughness: f32 = pbr_bindings::material.perceptual_roughness;
+#ifdef VERTEX_UVS
+        if ((pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_METALLIC_ROUGHNESS_TEXTURE_BIT) != 0u) {
+            let metallic_roughness = textureSampleBias(pbr_bindings::metallic_roughness_texture, pbr_bindings::metallic_roughness_sampler, uv, view.mip_bias);
+            // Sampling from GLTF standard channels for now
+            metallic *= metallic_roughness.b;
+            perceptual_roughness *= metallic_roughness.g;
+        }
+#endif
+        pbr_input.material.metallic = metallic;
+        pbr_input.material.perceptual_roughness = perceptual_roughness;
+
+        // occlusion
+        // TODO: Split into diffuse/specular occlusion?
+        var occlusion: vec3<f32> = vec3(1.0);
+#ifdef VERTEX_UVS
+        if ((pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_OCCLUSION_TEXTURE_BIT) != 0u) {
+            occlusion = vec3(textureSampleBias(pbr_bindings::occlusion_texture, pbr_bindings::occlusion_sampler, uv, view.mip_bias).r);
+        }
+#endif
+#ifdef SCREEN_SPACE_AMBIENT_OCCLUSION
+        let ssao = textureLoad(screen_space_ambient_occlusion_texture, vec2<i32>(in.position.xy), 0i).r;
+        let ssao_multibounce = gtao_multibounce(ssao, pbr_input.material.base_color.rgb);
+        occlusion = min(occlusion, ssao_multibounce);
+#endif
+        pbr_input.occlusion = occlusion;
+
+        // N (normal vector)
+#ifndef LOAD_PREPASS_NORMALS
+        pbr_input.N = pbr_functions::apply_normal_mapping(
+            pbr_bindings::material.flags,
+            pbr_input.world_normal,
+#ifdef VERTEX_TANGENTS
+#ifdef STANDARDMATERIAL_NORMAL_MAP
+            in.world_tangent,
+#endif
+#endif
+#ifdef VERTEX_UVS
+            uv,
+#endif
+            view.mip_bias,
+        );
+#endif
+    }
+
+    return pbr_input;
+}

crates/bevy_pbr/src/render/pbr_functions.wgsl

@@ -11,11 +11,12 @@
 #import bevy_pbr::lighting as lighting
 #import bevy_pbr::clustered_forward as clustering
 #import bevy_pbr::shadows as shadows
-#import bevy_pbr::fog as fog
+#import bevy_pbr::fog
 #import bevy_pbr::ambient as ambient
 #ifdef ENVIRONMENT_MAP
 #import bevy_pbr::environment_map
 #endif
+#import bevy_core_pipeline::tonemapping    screen_space_dither, powsafe, tone_mapping
 
 #import bevy_pbr::mesh_types      MESH_FLAGS_SHADOW_RECEIVER_BIT
 
@@ -137,7 +138,7 @@ fn calculate_view(
 }
 
 #ifndef PREPASS_FRAGMENT
-fn pbr(
+fn apply_pbr_lighting(
     in: pbr_types::PbrInput,
 ) -> vec4<f32> {
     var output_color: vec4<f32> = in.material.base_color;
@@ -247,7 +248,6 @@ fn pbr(
 }
 #endif // PREPASS_FRAGMENT
 
-#ifndef PREPASS_FRAGMENT
 fn apply_fog(fog_params: mesh_view_types::Fog, input_color: vec4<f32>, fragment_world_position: vec3<f32>, view_world_position: vec3<f32>) -> vec4<f32> {
     let view_to_world = fragment_world_position.xyz - view_world_position.xyz;
 
@@ -274,18 +274,17 @@ fn apply_fog(fog_params: mesh_view_types::Fog, input_color: vec4<f32>, fragment_
     }
 
     if fog_params.mode == mesh_view_types::FOG_MODE_LINEAR {
-        return fog::linear_fog(fog_params, input_color, distance, scattering);
+        return bevy_pbr::fog::linear_fog(fog_params, input_color, distance, scattering);
     } else if fog_params.mode == mesh_view_types::FOG_MODE_EXPONENTIAL {
-        return fog::exponential_fog(fog_params, input_color, distance, scattering);
+        return bevy_pbr::fog::exponential_fog(fog_params, input_color, distance, scattering);
     } else if fog_params.mode == mesh_view_types::FOG_MODE_EXPONENTIAL_SQUARED {
-        return fog::exponential_squared_fog(fog_params, input_color, distance, scattering);
+        return bevy_pbr::fog::exponential_squared_fog(fog_params, input_color, distance, scattering);
     } else if fog_params.mode == mesh_view_types::FOG_MODE_ATMOSPHERIC {
-        return fog::atmospheric_fog(fog_params, input_color, distance, scattering);
+        return bevy_pbr::fog::atmospheric_fog(fog_params, input_color, distance, scattering);
     } else {
         return input_color;
     }
 }
-#endif // PREPASS_FRAGMENT
 
 #ifdef PREMULTIPLY_ALPHA
 fn premultiply_alpha(standard_material_flags: u32, color: vec4<f32>) -> vec4<f32> {
@@ -338,3 +337,34 @@ fn premultiply_alpha(standard_material_flags: u32, color: vec4<f32>) -> vec4<f32
 #endif
 }
 #endif
+
+// fog, alpha premultiply
+// for non-hdr cameras, tonemapping and debanding
+fn main_pass_post_lighting_processing(
+    pbr_input: pbr_types::PbrInput,
+    input_color: vec4<f32>,
+) -> vec4<f32> {
+    var output_color = input_color;
+
+    // fog
+    if (view_bindings::fog.mode != mesh_view_types::FOG_MODE_OFF && (pbr_input.material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_FOG_ENABLED_BIT) != 0u) {
+        output_color = apply_fog(view_bindings::fog, output_color, pbr_input.world_position.xyz, view_bindings::view.world_position.xyz);
+    }
+
+#ifdef TONEMAP_IN_SHADER
+    output_color = tone_mapping(output_color, view_bindings::view.color_grading);
+#ifdef DEBAND_DITHER
+    var output_rgb = output_color.rgb;
+    output_rgb = powsafe(output_rgb, 1.0 / 2.2);
+    output_rgb += screen_space_dither(pbr_input.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
+#ifdef PREMULTIPLY_ALPHA
+    output_color = premultiply_alpha(pbr_input.material.flags, output_color);
+#endif
+    return output_color;
+}

crates/bevy_render/macros/src/as_bind_group.rs

@@ -43,7 +43,6 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
 
     let mut binding_states: Vec<BindingState> = Vec::new();
     let mut binding_impls = Vec::new();
-    let mut bind_group_entries = Vec::new();
     let mut binding_layouts = Vec::new();
     let mut attr_prepared_data_ident = None;
 
@@ -63,13 +62,16 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                     let mut buffer = #render_path::render_resource::encase::UniformBuffer::new(Vec::new());
                     let converted: #converted_shader_type = self.as_bind_group_shader_type(images);
                     buffer.write(&converted).unwrap();
-                    #render_path::render_resource::OwnedBindingResource::Buffer(render_device.create_buffer_with_data(
-                        &#render_path::render_resource::BufferInitDescriptor {
-                            label: None,
-                            usage: #render_path::render_resource::BufferUsages::COPY_DST | #render_path::render_resource::BufferUsages::UNIFORM,
-                            contents: buffer.as_ref(),
-                        },
-                    ))
+                    (
+                        #binding_index,
+                        #render_path::render_resource::OwnedBindingResource::Buffer(render_device.create_buffer_with_data(
+                            &#render_path::render_resource::BufferInitDescriptor {
+                                label: None,
+                                usage: #render_path::render_resource::BufferUsages::COPY_DST | #render_path::render_resource::BufferUsages::UNIFORM,
+                                contents: buffer.as_ref(),
+                            },
+                        ))
+                    )
                 }});
 
                 binding_layouts.push(quote!{
@@ -85,14 +87,6 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                     }
                 });
 
-                let binding_vec_index = bind_group_entries.len();
-                bind_group_entries.push(quote! {
-                    #render_path::render_resource::BindGroupEntry {
-                        binding: #binding_index,
-                        resource: bindings[#binding_vec_index].get_binding(),
-                    }
-                });
-
                 let required_len = binding_index as usize + 1;
                 if required_len > binding_states.len() {
                     binding_states.resize(required_len, BindingState::Free);
@@ -164,13 +158,6 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                         _ => {
                             // only populate bind group entries for non-uniforms
                             // uniform entries are deferred until the end
-                            let binding_vec_index = bind_group_entries.len();
-                            bind_group_entries.push(quote! {
-                                #render_path::render_resource::BindGroupEntry {
-                                    binding: #binding_index,
-                                    resource: bindings[#binding_vec_index].get_binding(),
-                                }
-                            });
                             BindingState::Occupied {
                                 binding_type,
                                 ident: field_name,
@@ -230,22 +217,28 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
 
                     if buffer {
                         binding_impls.push(quote! {
-                            #render_path::render_resource::OwnedBindingResource::Buffer({
-                                self.#field_name.clone()
-                            })
+                            (
+                                #binding_index,
+                                #render_path::render_resource::OwnedBindingResource::Buffer({
+                                    self.#field_name.clone()
+                                })
+                            )
                         });
                     } else {
                         binding_impls.push(quote! {{
                             use #render_path::render_resource::AsBindGroupShaderType;
                             let mut buffer = #render_path::render_resource::encase::StorageBuffer::new(Vec::new());
                             buffer.write(&self.#field_name).unwrap();
-                            #render_path::render_resource::OwnedBindingResource::Buffer(render_device.create_buffer_with_data(
-                                &#render_path::render_resource::BufferInitDescriptor {
-                                    label: None,
-                                    usage: #render_path::render_resource::BufferUsages::COPY_DST | #render_path::render_resource::BufferUsages::STORAGE,
-                                    contents: buffer.as_ref(),
-                                },
-                            ))
+                            (
+                                #binding_index,
+                                #render_path::render_resource::OwnedBindingResource::Buffer(render_device.create_buffer_with_data(
+                                    &#render_path::render_resource::BufferInitDescriptor {
+                                        label: None,
+                                        usage: #render_path::render_resource::BufferUsages::COPY_DST | #render_path::render_resource::BufferUsages::STORAGE,
+                                        contents: buffer.as_ref(),
+                                    },
+                                ))
+                            )
                         }});
                     }
 
@@ -276,14 +269,17 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                     let fallback_image = get_fallback_image(&render_path, *dimension);
 
                     binding_impls.push(quote! {
-                        #render_path::render_resource::OwnedBindingResource::TextureView({
-                            let handle: Option<&#asset_path::Handle<#render_path::texture::Image>> = (&self.#field_name).into();
-                            if let Some(handle) = handle {
-                                images.get(handle).ok_or_else(|| #render_path::render_resource::AsBindGroupError::RetryNextUpdate)?.texture_view.clone()
-                            } else {
-                                #fallback_image.texture_view.clone()
-                            }
-                        })
+                        (
+                            #binding_index,
+                            #render_path::render_resource::OwnedBindingResource::TextureView({
+                                let handle: Option<&#asset_path::Handle<#render_path::texture::Image>> = (&self.#field_name).into();
+                                if let Some(handle) = handle {
+                                    images.get(handle).ok_or_else(|| #render_path::render_resource::AsBindGroupError::RetryNextUpdate)?.texture_view.clone()
+                                } else {
+                                    #fallback_image.texture_view.clone()
+                                }
+                            })
+                        )
                     });
 
                     binding_layouts.push(quote! {
@@ -315,14 +311,17 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                     let fallback_image = get_fallback_image(&render_path, *dimension);
 
                     binding_impls.push(quote! {
-                        #render_path::render_resource::OwnedBindingResource::Sampler({
-                            let handle: Option<&#asset_path::Handle<#render_path::texture::Image>> = (&self.#field_name).into();
-                            if let Some(handle) = handle {
-                                images.get(handle).ok_or_else(|| #render_path::render_resource::AsBindGroupError::RetryNextUpdate)?.sampler.clone()
-                            } else {
-                                #fallback_image.sampler.clone()
-                            }
-                        })
+                        (
+                            #binding_index,
+                            #render_path::render_resource::OwnedBindingResource::Sampler({
+                                let handle: Option<&#asset_path::Handle<#render_path::texture::Image>> = (&self.#field_name).into();
+                                if let Some(handle) = handle {
+                                    images.get(handle).ok_or_else(|| #render_path::render_resource::AsBindGroupError::RetryNextUpdate)?.sampler.clone()
+                                } else {
+                                    #fallback_image.sampler.clone()
+                                }
+                            })
+                        )
                     });
 
                     binding_layouts.push(quote!{
@@ -340,17 +339,12 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
 
     // Produce impls for fields with uniform bindings
     let struct_name = &ast.ident;
+    let struct_name_literal = struct_name.to_string();
+    let struct_name_literal = struct_name_literal.as_str();
     let mut field_struct_impls = Vec::new();
     for (binding_index, binding_state) in binding_states.iter().enumerate() {
         let binding_index = binding_index as u32;
         if let BindingState::OccupiedMergeableUniform { uniform_fields } = binding_state {
-            let binding_vec_index = bind_group_entries.len();
-            bind_group_entries.push(quote! {
-                #render_path::render_resource::BindGroupEntry {
-                    binding: #binding_index,
-                    resource: bindings[#binding_vec_index].get_binding(),
-                }
-            });
             // single field uniform bindings for a given index can use a straightforward binding
             if uniform_fields.len() == 1 {
                 let field = &uniform_fields[0];
@@ -359,13 +353,16 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                 binding_impls.push(quote! {{
                     let mut buffer = #render_path::render_resource::encase::UniformBuffer::new(Vec::new());
                     buffer.write(&self.#field_name).unwrap();
-                    #render_path::render_resource::OwnedBindingResource::Buffer(render_device.create_buffer_with_data(
-                        &#render_path::render_resource::BufferInitDescriptor {
-                            label: None,
-                            usage: #render_path::render_resource::BufferUsages::COPY_DST | #render_path::render_resource::BufferUsages::UNIFORM,
-                            contents: buffer.as_ref(),
-                        },
-                    ))
+                    (
+                        #binding_index,
+                        #render_path::render_resource::OwnedBindingResource::Buffer(render_device.create_buffer_with_data(
+                            &#render_path::render_resource::BufferInitDescriptor {
+                                label: None,
+                                usage: #render_path::render_resource::BufferUsages::COPY_DST | #render_path::render_resource::BufferUsages::UNIFORM,
+                                contents: buffer.as_ref(),
+                            },
+                        ))
+                    )
                 }});
 
                 binding_layouts.push(quote!{
@@ -402,13 +399,16 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                     buffer.write(&#uniform_struct_name {
                         #(#field_name: &self.#field_name,)*
                     }).unwrap();
-                    #render_path::render_resource::OwnedBindingResource::Buffer(render_device.create_buffer_with_data(
-                        &#render_path::render_resource::BufferInitDescriptor {
-                            label: None,
-                            usage: #render_path::render_resource::BufferUsages::COPY_DST | #render_path::render_resource::BufferUsages::UNIFORM,
-                            contents: buffer.as_ref(),
-                        },
-                    ))
+                    (
+                        #binding_index,
+                        #render_path::render_resource::OwnedBindingResource::Buffer(render_device.create_buffer_with_data(
+                            &#render_path::render_resource::BufferInitDescriptor {
+                                label: None,
+                                usage: #render_path::render_resource::BufferUsages::COPY_DST | #render_path::render_resource::BufferUsages::UNIFORM,
+                                contents: buffer.as_ref(),
+                            },
+                        ))
+                    )
                 }});
 
                 binding_layouts.push(quote!{
@@ -443,36 +443,28 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
 
         impl #impl_generics #render_path::render_resource::AsBindGroup for #struct_name #ty_generics #where_clause {
             type Data = #prepared_data;
-            fn as_bind_group(
+
+            fn label() -> Option<&'static str> {
+                Some(#struct_name_literal)
+            }
+
+            fn unprepared_bind_group(
                 &self,
                 layout: &#render_path::render_resource::BindGroupLayout,
                 render_device: &#render_path::renderer::RenderDevice,
                 images: &#render_path::render_asset::RenderAssets<#render_path::texture::Image>,
                 fallback_image: &#render_path::texture::FallbackImage,
-            ) -> Result<#render_path::render_resource::PreparedBindGroup<Self::Data>, #render_path::render_resource::AsBindGroupError> {
+            ) -> Result<#render_path::render_resource::UnpreparedBindGroup<Self::Data>, #render_path::render_resource::AsBindGroupError> {
                 let bindings = vec![#(#binding_impls,)*];
 
-                let bind_group = {
-                    let descriptor = #render_path::render_resource::BindGroupDescriptor {
-                        entries: &[#(#bind_group_entries,)*],
-                        label: None,
-                        layout: &layout,
-                    };
-                    render_device.create_bind_group(&descriptor)
-                };
-
-                Ok(#render_path::render_resource::PreparedBindGroup {
+                Ok(#render_path::render_resource::UnpreparedBindGroup {
                     bindings,
-                    bind_group,
                     data: #get_prepared_data,
                 })
             }
 
-            fn bind_group_layout(render_device: &#render_path::renderer::RenderDevice) -> #render_path::render_resource::BindGroupLayout {
-                render_device.create_bind_group_layout(&#render_path::render_resource::BindGroupLayoutDescriptor {
-                    entries: &[#(#binding_layouts,)*],
-                    label: None,
-                })
+            fn bind_group_layout_entries(render_device: &#render_path::renderer::RenderDevice) -> Vec<#render_path::render_resource::BindGroupLayoutEntry> {
+                vec![#(#binding_layouts,)*]
             }
         }
     }))

crates/bevy_render/src/render_resource/bind_group.rs

@@ -9,7 +9,10 @@ use crate::{
 pub use bevy_render_macros::AsBindGroup;
 use encase::ShaderType;
 use std::ops::Deref;
-use wgpu::BindingResource;
+use wgpu::{
+    BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor, BindGroupLayoutEntry,
+    BindingResource,
+};
 
 define_atomic_id!(BindGroupId);
 render_resource_wrapper!(ErasedBindGroup, wgpu::BindGroup);
@@ -262,6 +265,11 @@ pub trait AsBindGroup {
     /// Data that will be stored alongside the "prepared" bind group.
     type Data: Send + Sync;
 
+    /// label
+    fn label() -> Option<&'static str> {
+        None
+    }
+
     /// Creates a bind group for `self` matching the layout defined in [`AsBindGroup::bind_group_layout`].
     fn as_bind_group(
         &self,
@@ -269,10 +277,56 @@ pub trait AsBindGroup {
         render_device: &RenderDevice,
         images: &RenderAssets<Image>,
         fallback_image: &FallbackImage,
-    ) -> Result<PreparedBindGroup<Self::Data>, AsBindGroupError>;
+    ) -> Result<PreparedBindGroup<Self::Data>, AsBindGroupError> {
+        let UnpreparedBindGroup { bindings, data } =
+            Self::unprepared_bind_group(self, layout, render_device, images, fallback_image)?;
+
+        let entries = bindings
+            .iter()
+            .map(|(index, binding)| BindGroupEntry {
+                binding: *index,
+                resource: binding.get_binding(),
+            })
+            .collect::<Vec<_>>();
+
+        let bind_group = render_device.create_bind_group(&BindGroupDescriptor {
+            label: Self::label(),
+            layout,
+            entries: &entries,
+        });
+
+        Ok(PreparedBindGroup {
+            bindings,
+            bind_group,
+            data,
+        })
+    }
+
+    /// Returns a vec of (binding index, `OwnedBindingResource`).
+    /// In cases where `OwnedBindingResource` is not available (as for bindless texture arrays currently),
+    /// an implementor may define `as_bind_group` directly. This may prevent certain features
+    /// from working correctly.
+    fn unprepared_bind_group(
+        &self,
+        layout: &BindGroupLayout,
+        render_device: &RenderDevice,
+        images: &RenderAssets<Image>,
+        fallback_image: &FallbackImage,
+    ) -> Result<UnpreparedBindGroup<Self::Data>, AsBindGroupError>;
 
     /// Creates the bind group layout matching all bind groups returned by [`AsBindGroup::as_bind_group`]
     fn bind_group_layout(render_device: &RenderDevice) -> BindGroupLayout
+    where
+        Self: Sized,
+    {
+        render_device.create_bind_group_layout(&BindGroupLayoutDescriptor {
+            label: Self::label(),
+            entries: &Self::bind_group_layout_entries(render_device),
+        })
+    }
+
+    /// Returns a vec of bind group layout entries
+    fn bind_group_layout_entries(render_device: &RenderDevice) -> Vec<BindGroupLayoutEntry>
     where
         Self: Sized;
 }
@@ -285,14 +339,21 @@ pub enum AsBindGroupError {
 
 /// A prepared bind group returned as a result of [`AsBindGroup::as_bind_group`].
 pub struct PreparedBindGroup<T> {
-    pub bindings: Vec<OwnedBindingResource>,
+    pub bindings: Vec<(u32, OwnedBindingResource)>,
     pub bind_group: BindGroup,
     pub data: T,
 }
 
+/// a map containing `OwnedBindingResource`s, keyed by the target binding index
+pub struct UnpreparedBindGroup<T> {
+    pub bindings: Vec<(u32, OwnedBindingResource)>,
+    pub data: T,
+}
+
 /// An owned binding resource of any type (ex: a [`Buffer`], [`TextureView`], etc).
 /// This is used by types like [`PreparedBindGroup`] to hold a single list of all
 /// render resources used by bindings.
+#[derive(Debug)]
 pub enum OwnedBindingResource {
     Buffer(Buffer),
     TextureView(TextureView),

crates/bevy_sprite/src/mesh2d/material.rs

@@ -462,7 +462,7 @@ pub struct Material2dBindGroupId(Option<BindGroupId>);
 
 /// Data prepared for a [`Material2d`] instance.
 pub struct PreparedMaterial2d<T: Material2d> {
-    pub bindings: Vec<OwnedBindingResource>,
+    pub bindings: Vec<(u32, OwnedBindingResource)>,
     pub bind_group: BindGroup,
     pub key: T::Data,
 }

examples/README.md

@@ -292,6 +292,7 @@ Example | Description
 [Array Texture](../examples/shader/array_texture.rs) | A shader that shows how to reuse the core bevy PBR shading functionality in a custom material that obtains the base color from an array texture.
 [Compute - Game of Life](../examples/shader/compute_shader_game_of_life.rs) | A compute shader that simulates Conway's Game of Life
 [Custom Vertex Attribute](../examples/shader/custom_vertex_attribute.rs) | A shader that reads a mesh's custom vertex attribute
+[Extended Material](../examples/shader/extended_material.rs) | A custom shader that builds on the standard material
 [Instancing](../examples/shader/shader_instancing.rs) | A shader that renders a mesh multiple times in one draw call
 [Material](../examples/shader/shader_material.rs) | A shader and a material that uses it
 [Material - GLSL](../examples/shader/shader_material_glsl.rs) | A shader that uses the GLSL shading language

examples/shader/extended_material.rs

@@ -0,0 +1,92 @@
+//! Demonstrates using a custom extension to the `StandardMaterial` to modify the results of the builtin pbr shader.
+
+use bevy::reflect::TypePath;
+use bevy::{
+    pbr::{ExtendedMaterial, MaterialExtension, OpaqueRendererMethod},
+    prelude::*,
+    render::render_resource::*,
+};
+
+fn main() {
+    App::new()
+        .add_plugins(DefaultPlugins)
+        .add_plugins(MaterialPlugin::<
+            ExtendedMaterial<StandardMaterial, MyExtension>,
+        >::default())
+        .add_systems(Startup, setup)
+        .add_systems(Update, rotate_things)
+        .run();
+}
+
+fn setup(
+    mut commands: Commands,
+    mut meshes: ResMut<Assets<Mesh>>,
+    mut materials: ResMut<Assets<ExtendedMaterial<StandardMaterial, MyExtension>>>,
+) {
+    // sphere
+    commands.spawn(MaterialMeshBundle {
+        mesh: meshes.add(
+            Mesh::try_from(shape::Icosphere {
+                radius: 1.0,
+                subdivisions: 5,
+            })
+            .unwrap(),
+        ),
+        transform: Transform::from_xyz(0.0, 0.5, 0.0),
+        material: materials.add(ExtendedMaterial {
+            base: StandardMaterial {
+                base_color: Color::RED,
+                // can be used in forward or deferred mode.
+                opaque_render_method: OpaqueRendererMethod::Auto,
+                // in deferred mode, only the PbrInput can be modified (uvs, color and other material properties),
+                // in forward mode, the output can also be modified after lighting is applied.
+                // see the fragment shader `extended_material.wgsl` for more info.
+                // Note: to run in deferred mode, you must also add a `DeferredPrepass` component to the camera and either
+                // change the above to `OpaqueRendererMethod::Deferred` or add the `DefaultOpaqueRendererMethod` resource.
+                ..Default::default()
+            },
+            extension: MyExtension { quantize_steps: 3 },
+        }),
+        ..default()
+    });
+
+    // light
+    commands.spawn((PointLightBundle::default(), Rotate));
+
+    // camera
+    commands.spawn(Camera3dBundle {
+        transform: Transform::from_xyz(-2.0, 2.5, 5.0).looking_at(Vec3::ZERO, Vec3::Y),
+        ..default()
+    });
+}
+
+#[derive(Component)]
+struct Rotate;
+
+fn rotate_things(mut q: Query<&mut Transform, With<Rotate>>, time: Res<Time>) {
+    for mut t in q.iter_mut() {
+        t.translation = Vec3::new(
+            time.elapsed_seconds().sin(),
+            0.5,
+            time.elapsed_seconds().cos(),
+        ) * 4.0;
+    }
+}
+
+#[derive(Asset, AsBindGroup, TypePath, Debug, Clone)]
+struct MyExtension {
+    // We need to ensure that the bindings of the base material and the extension do not conflict,
+    // so we start from binding slot 100, leaving slots 0-99 for the base material.
+    #[uniform(100)]
+    quantize_steps: u32,
+}
+
+impl MaterialExtension for MyExtension {
+    fn fragment_shader() -> ShaderRef {
+        "shaders/extended_material.wgsl".into()
+    }
+
+    fn deferred_fragment_shader() -> ShaderRef {
+        "shaders/extended_material.wgsl".into()
+    }
+}

examples/shader/texture_binding_array.rs

@@ -141,36 +141,45 @@ impl AsBindGroup for BindlessMaterial {
         })
     }
 
-    fn bind_group_layout(render_device: &RenderDevice) -> BindGroupLayout
+    fn unprepared_bind_group(
+        &self,
+        _: &BindGroupLayout,
+        _: &RenderDevice,
+        _: &RenderAssets<Image>,
+        _: &FallbackImage,
+    ) -> Result<UnpreparedBindGroup<Self::Data>, AsBindGroupError> {
+        // we implement as_bind_group directly because
+        panic!("bindless texture arrays can't be owned")
+        // or rather, they can be owned, but then you can't make a `&'a [&'a TextureView]` from a vec of them in get_binding().
+    }
+
+    fn bind_group_layout_entries(_: &RenderDevice) -> Vec<BindGroupLayoutEntry>
     where
         Self: Sized,
     {
-        render_device.create_bind_group_layout(&BindGroupLayoutDescriptor {
-            label: "bindless_material_layout".into(),
-            entries: &[
-                // @group(1) @binding(0) var textures: binding_array<texture_2d<f32>>;
-                BindGroupLayoutEntry {
-                    binding: 0,
-                    visibility: ShaderStages::FRAGMENT,
-                    ty: BindingType::Texture {
-                        sample_type: TextureSampleType::Float { filterable: true },
-                        view_dimension: TextureViewDimension::D2,
-                        multisampled: false,
-                    },
-                    count: NonZeroU32::new(MAX_TEXTURE_COUNT as u32),
+        vec![
+            // @group(1) @binding(0) var textures: binding_array<texture_2d<f32>>;
+            BindGroupLayoutEntry {
+                binding: 0,
+                visibility: ShaderStages::FRAGMENT,
+                ty: BindingType::Texture {
+                    sample_type: TextureSampleType::Float { filterable: true },
+                    view_dimension: TextureViewDimension::D2,
+                    multisampled: false,
                 },
-                // @group(1) @binding(1) var nearest_sampler: sampler;
-                BindGroupLayoutEntry {
-                    binding: 1,
-                    visibility: ShaderStages::FRAGMENT,
-                    ty: BindingType::Sampler(SamplerBindingType::Filtering),
-                    count: None,
-                    // Note: as textures, multiple samplers can also be bound onto one binding slot.
-                    // One may need to pay attention to the limit of sampler binding amount on some platforms.
-                    // count: NonZeroU32::new(MAX_TEXTURE_COUNT as u32),
-                },
-            ],
-        })
+                count: NonZeroU32::new(MAX_TEXTURE_COUNT as u32),
+            },
+            // @group(1) @binding(1) var nearest_sampler: sampler;
+            BindGroupLayoutEntry {
+                binding: 1,
+                visibility: ShaderStages::FRAGMENT,
+                ty: BindingType::Sampler(SamplerBindingType::Filtering),
+                count: None,
+                // Note: as textures, multiple samplers can also be bound onto one binding slot.
+                // One may need to pay attention to the limit of sampler binding amount on some platforms.
+                // count: NonZeroU32::new(MAX_TEXTURE_COUNT as u32),
+            },
+        ]
     }
 }