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,
+ },
+ )
+ }
+}