bevy_pbr/
extended_material.rs

1use bevy_asset::{Asset, Handle};
2use bevy_ecs::system::SystemParamItem;
3use bevy_reflect::{impl_type_path, Reflect};
4use bevy_render::{
5    mesh::MeshVertexBufferLayoutRef,
6    render_resource::{
7        AsBindGroup, AsBindGroupError, BindGroupLayout, RenderPipelineDescriptor, Shader,
8        ShaderRef, SpecializedMeshPipelineError, UnpreparedBindGroup,
9    },
10    renderer::RenderDevice,
11};
12
13use crate::{Material, MaterialPipeline, MaterialPipelineKey, MeshPipeline, MeshPipelineKey};
14
15pub struct MaterialExtensionPipeline {
16    pub mesh_pipeline: MeshPipeline,
17    pub material_layout: BindGroupLayout,
18    pub vertex_shader: Option<Handle<Shader>>,
19    pub fragment_shader: Option<Handle<Shader>>,
20}
21
22pub struct MaterialExtensionKey<E: MaterialExtension> {
23    pub mesh_key: MeshPipelineKey,
24    pub bind_group_data: E::Data,
25}
26
27/// A subset of the `Material` trait for defining extensions to a base `Material`, such as the builtin `StandardMaterial`.
28///
29/// A user type implementing the trait should be used as the `E` generic param in an `ExtendedMaterial` struct.
30pub trait MaterialExtension: Asset + AsBindGroup + Clone + Sized {
31    /// Returns this material's vertex shader. If [`ShaderRef::Default`] is returned, the base material mesh vertex shader
32    /// will be used.
33    fn vertex_shader() -> ShaderRef {
34        ShaderRef::Default
35    }
36
37    /// Returns this material's fragment shader. If [`ShaderRef::Default`] is returned, the base material mesh fragment shader
38    /// will be used.
39    #[allow(unused_variables)]
40    fn fragment_shader() -> ShaderRef {
41        ShaderRef::Default
42    }
43
44    /// Returns this material's prepass vertex shader. If [`ShaderRef::Default`] is returned, the base material prepass vertex shader
45    /// will be used.
46    fn prepass_vertex_shader() -> ShaderRef {
47        ShaderRef::Default
48    }
49
50    /// Returns this material's prepass fragment shader. If [`ShaderRef::Default`] is returned, the base material prepass fragment shader
51    /// will be used.
52    #[allow(unused_variables)]
53    fn prepass_fragment_shader() -> ShaderRef {
54        ShaderRef::Default
55    }
56
57    /// Returns this material's deferred vertex shader. If [`ShaderRef::Default`] is returned, the base material deferred vertex shader
58    /// will be used.
59    fn deferred_vertex_shader() -> ShaderRef {
60        ShaderRef::Default
61    }
62
63    /// Returns this material's prepass fragment shader. If [`ShaderRef::Default`] is returned, the base material deferred fragment shader
64    /// will be used.
65    #[allow(unused_variables)]
66    fn deferred_fragment_shader() -> ShaderRef {
67        ShaderRef::Default
68    }
69
70    /// Returns this material's [`crate::meshlet::MeshletMesh`] fragment shader. If [`ShaderRef::Default`] is returned,
71    /// the default meshlet mesh fragment shader will be used.
72    #[allow(unused_variables)]
73    #[cfg(feature = "meshlet")]
74    fn meshlet_mesh_fragment_shader() -> ShaderRef {
75        ShaderRef::Default
76    }
77
78    /// Returns this material's [`crate::meshlet::MeshletMesh`] prepass fragment shader. If [`ShaderRef::Default`] is returned,
79    /// the default meshlet mesh prepass fragment shader will be used.
80    #[allow(unused_variables)]
81    #[cfg(feature = "meshlet")]
82    fn meshlet_mesh_prepass_fragment_shader() -> ShaderRef {
83        ShaderRef::Default
84    }
85
86    /// Returns this material's [`crate::meshlet::MeshletMesh`] deferred fragment shader. If [`ShaderRef::Default`] is returned,
87    /// the default meshlet mesh deferred fragment shader will be used.
88    #[allow(unused_variables)]
89    #[cfg(feature = "meshlet")]
90    fn meshlet_mesh_deferred_fragment_shader() -> ShaderRef {
91        ShaderRef::Default
92    }
93
94    /// Customizes the default [`RenderPipelineDescriptor`] for a specific entity using the entity's
95    /// [`MaterialPipelineKey`] and [`MeshVertexBufferLayoutRef`] as input.
96    /// Specialization for the base material is applied before this function is called.
97    #[allow(unused_variables)]
98    #[inline]
99    fn specialize(
100        pipeline: &MaterialExtensionPipeline,
101        descriptor: &mut RenderPipelineDescriptor,
102        layout: &MeshVertexBufferLayoutRef,
103        key: MaterialExtensionKey<Self>,
104    ) -> Result<(), SpecializedMeshPipelineError> {
105        Ok(())
106    }
107}
108
109/// A material that extends a base [`Material`] with additional shaders and data.
110///
111/// The data from both materials will be combined and made available to the shader
112/// so that shader functions built for the base material (and referencing the base material
113/// bindings) will work as expected, and custom alterations based on custom data can also be used.
114///
115/// If the extension `E` returns a non-default result from `vertex_shader()` it will be used in place of the base
116/// material's vertex shader.
117///
118/// If the extension `E` returns a non-default result from `fragment_shader()` it will be used in place of the base
119/// fragment shader.
120///
121/// When used with `StandardMaterial` as the base, all the standard material fields are
122/// present, so the `pbr_fragment` shader functions can be called from the extension shader (see
123/// the `extended_material` example).
124#[derive(Asset, Clone, Debug, Reflect)]
125#[reflect(type_path = false)]
126pub struct ExtendedMaterial<B: Material, E: MaterialExtension> {
127    pub base: B,
128    pub extension: E,
129}
130
131impl<B, E> Default for ExtendedMaterial<B, E>
132where
133    B: Material + Default,
134    E: MaterialExtension + Default,
135{
136    fn default() -> Self {
137        Self {
138            base: B::default(),
139            extension: E::default(),
140        }
141    }
142}
143
144// We don't use the `TypePath` derive here due to a bug where `#[reflect(type_path = false)]`
145// causes the `TypePath` derive to not generate an implementation.
146impl_type_path!((in bevy_pbr::extended_material) ExtendedMaterial<B: Material, E: MaterialExtension>);
147
148impl<B: Material, E: MaterialExtension> AsBindGroup for ExtendedMaterial<B, E> {
149    type Data = (<B as AsBindGroup>::Data, <E as AsBindGroup>::Data);
150    type Param = (<B as AsBindGroup>::Param, <E as AsBindGroup>::Param);
151
152    fn unprepared_bind_group(
153        &self,
154        layout: &BindGroupLayout,
155        render_device: &RenderDevice,
156        (base_param, extended_param): &mut SystemParamItem<'_, '_, Self::Param>,
157    ) -> Result<UnpreparedBindGroup<Self::Data>, AsBindGroupError> {
158        // add together the bindings of the base material and the user material
159        let UnpreparedBindGroup {
160            mut bindings,
161            data: base_data,
162        } = B::unprepared_bind_group(&self.base, layout, render_device, base_param)?;
163        let extended_bindgroup =
164            E::unprepared_bind_group(&self.extension, layout, render_device, extended_param)?;
165
166        bindings.extend(extended_bindgroup.bindings);
167
168        Ok(UnpreparedBindGroup {
169            bindings,
170            data: (base_data, extended_bindgroup.data),
171        })
172    }
173
174    fn bind_group_layout_entries(
175        render_device: &RenderDevice,
176    ) -> Vec<bevy_render::render_resource::BindGroupLayoutEntry>
177    where
178        Self: Sized,
179    {
180        // add together the bindings of the standard material and the user material
181        let mut entries = B::bind_group_layout_entries(render_device);
182        entries.extend(E::bind_group_layout_entries(render_device));
183        entries
184    }
185}
186
187impl<B: Material, E: MaterialExtension> Material for ExtendedMaterial<B, E> {
188    fn vertex_shader() -> ShaderRef {
189        match E::vertex_shader() {
190            ShaderRef::Default => B::vertex_shader(),
191            specified => specified,
192        }
193    }
194
195    fn fragment_shader() -> ShaderRef {
196        match E::fragment_shader() {
197            ShaderRef::Default => B::fragment_shader(),
198            specified => specified,
199        }
200    }
201
202    fn alpha_mode(&self) -> crate::AlphaMode {
203        B::alpha_mode(&self.base)
204    }
205
206    fn opaque_render_method(&self) -> crate::OpaqueRendererMethod {
207        B::opaque_render_method(&self.base)
208    }
209
210    fn depth_bias(&self) -> f32 {
211        B::depth_bias(&self.base)
212    }
213
214    fn reads_view_transmission_texture(&self) -> bool {
215        B::reads_view_transmission_texture(&self.base)
216    }
217
218    fn prepass_vertex_shader() -> ShaderRef {
219        match E::prepass_vertex_shader() {
220            ShaderRef::Default => B::prepass_vertex_shader(),
221            specified => specified,
222        }
223    }
224
225    fn prepass_fragment_shader() -> ShaderRef {
226        match E::prepass_fragment_shader() {
227            ShaderRef::Default => B::prepass_fragment_shader(),
228            specified => specified,
229        }
230    }
231
232    fn deferred_vertex_shader() -> ShaderRef {
233        match E::deferred_vertex_shader() {
234            ShaderRef::Default => B::deferred_vertex_shader(),
235            specified => specified,
236        }
237    }
238
239    fn deferred_fragment_shader() -> ShaderRef {
240        match E::deferred_fragment_shader() {
241            ShaderRef::Default => B::deferred_fragment_shader(),
242            specified => specified,
243        }
244    }
245
246    #[cfg(feature = "meshlet")]
247    fn meshlet_mesh_fragment_shader() -> ShaderRef {
248        match E::meshlet_mesh_fragment_shader() {
249            ShaderRef::Default => B::meshlet_mesh_fragment_shader(),
250            specified => specified,
251        }
252    }
253
254    #[cfg(feature = "meshlet")]
255    fn meshlet_mesh_prepass_fragment_shader() -> ShaderRef {
256        match E::meshlet_mesh_prepass_fragment_shader() {
257            ShaderRef::Default => B::meshlet_mesh_prepass_fragment_shader(),
258            specified => specified,
259        }
260    }
261
262    #[cfg(feature = "meshlet")]
263    fn meshlet_mesh_deferred_fragment_shader() -> ShaderRef {
264        match E::meshlet_mesh_deferred_fragment_shader() {
265            ShaderRef::Default => B::meshlet_mesh_deferred_fragment_shader(),
266            specified => specified,
267        }
268    }
269
270    fn specialize(
271        pipeline: &MaterialPipeline<Self>,
272        descriptor: &mut RenderPipelineDescriptor,
273        layout: &MeshVertexBufferLayoutRef,
274        key: MaterialPipelineKey<Self>,
275    ) -> Result<(), SpecializedMeshPipelineError> {
276        // Call the base material's specialize function
277        let MaterialPipeline::<Self> {
278            mesh_pipeline,
279            material_layout,
280            vertex_shader,
281            fragment_shader,
282            ..
283        } = pipeline.clone();
284        let base_pipeline = MaterialPipeline::<B> {
285            mesh_pipeline,
286            material_layout,
287            vertex_shader,
288            fragment_shader,
289            marker: Default::default(),
290        };
291        let base_key = MaterialPipelineKey::<B> {
292            mesh_key: key.mesh_key,
293            bind_group_data: key.bind_group_data.0,
294        };
295        B::specialize(&base_pipeline, descriptor, layout, base_key)?;
296
297        // Call the extended material's specialize function afterwards
298        let MaterialPipeline::<Self> {
299            mesh_pipeline,
300            material_layout,
301            vertex_shader,
302            fragment_shader,
303            ..
304        } = pipeline.clone();
305
306        E::specialize(
307            &MaterialExtensionPipeline {
308                mesh_pipeline,
309                material_layout,
310                vertex_shader,
311                fragment_shader,
312            },
313            descriptor,
314            layout,
315            MaterialExtensionKey {
316                mesh_key: key.mesh_key,
317                bind_group_data: key.bind_group_data.1,
318            },
319        )
320    }
321}