docs/bevy/examples/shader_advanced/manual_material.rs 11.4 K · 335 lines · raw · history

1 //! A simple 3D scene with light shining over a cube sitting on a plane.
2
3 use bevy::{
4 asset::{AsAssetId, AssetEventSystems},
5 core_pipeline::core_3d::Opaque3d,
6 ecs::system::{
7 lifetimeless::{SRes, SResMut},
8 SystemParamItem,
9 },
10 material::{key::ErasedMeshPipelineKey, MaterialProperties},
11 pbr::{
12 base_specialize, DrawMaterial, EntitiesNeedingSpecialization, MainPassOpaqueDrawFunction,
13 MaterialFragmentShader, MeshPipelineKey, PreparedMaterial, RenderMaterialInstance,
14 RenderMaterialInstances,
15 },
16 platform::collections::hash_map::Entry,
17 prelude::*,
18 render::{
19 camera::{DirtySpecializationSystems, DirtySpecializations},
20 erased_render_asset::{ErasedRenderAsset, ErasedRenderAssetPlugin, PrepareAssetError},
21 material_bind_groups::{
22 MaterialBindGroupAllocator, MaterialBindGroupAllocators, RenderMaterialBindings,
23 },
24 render_asset::RenderAssets,
25 render_phase::DrawFunctions,
26 render_resource::{
27 binding_types::{sampler, texture_2d},
28 AsBindGroup, BindGroupBuilder, BindGroupLayoutDescriptor, BindGroupLayoutEntries,
29 Sampler, SamplerBindingType, SamplerDescriptor, ShaderStages, TextureSampleType,
30 TextureViewDimension, UnpreparedBindingResource, UnpreparedBindingResources,
31 },
32 renderer::RenderDevice,
33 sync_world::MainEntity,
34 texture::GpuImage,
35 Extract, RenderApp, RenderStartup,
36 },
37 utils::Parallel,
38 };
39 use std::{any::TypeId, sync::Arc};
40
41 const SHADER_ASSET_PATH: &str = "shaders/manual_material.wesl";
42
43 fn main() {
44 App::new()
45 .add_plugins((DefaultPlugins, ImageMaterialPlugin))
46 .add_systems(Startup, setup)
47 .run();
48 }
49
50 struct ImageMaterialPlugin;
51
52 impl Plugin for ImageMaterialPlugin {
53 fn build(&self, app: &mut App) {
54 app.init_asset::<ImageMaterial>()
55 .add_plugins(ErasedRenderAssetPlugin::<ImageMaterial>::default())
56 .add_systems(
57 PostUpdate,
58 check_entities_needing_specialization.after(AssetEventSystems),
59 )
60 .init_resource::<EntitiesNeedingSpecialization<ImageMaterial>>();
61
62 let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
63 return;
64 };
65
66 render_app
67 .add_systems(RenderStartup, init_image_material_resources)
68 .add_systems(
69 ExtractSchedule,
70 (
71 extract_image_materials,
72 extract_image_materials_needing_specialization
73 .in_set(DirtySpecializationSystems::CheckForChanges),
74 extract_image_materials_that_need_specializations_removed
75 .in_set(DirtySpecializationSystems::CheckForRemovals),
76 ),
77 );
78 }
79 }
80
81 fn init_image_material_resources(
82 mut commands: Commands,
83 render_device: Res<RenderDevice>,
84 mut bind_group_allocators: ResMut<MaterialBindGroupAllocators>,
85 ) {
86 let bind_group_layout = BindGroupLayoutDescriptor::new(
87 "image_material_layout",
88 &BindGroupLayoutEntries::sequential(
89 ShaderStages::FRAGMENT,
90 (
91 texture_2d(TextureSampleType::Float { filterable: false }),
92 sampler(SamplerBindingType::NonFiltering),
93 ),
94 ),
95 );
96 let sampler = render_device.create_sampler(&SamplerDescriptor::default());
97 commands.insert_resource(ImageMaterialBindGroupLayout(bind_group_layout.clone()));
98 commands.insert_resource(ImageMaterialBindGroupSampler(sampler));
99
100 bind_group_allocators.insert(
101 TypeId::of::<ImageMaterial>(),
102 MaterialBindGroupAllocator::new(
103 &render_device,
104 "image_material_allocator",
105 None,
106 bind_group_layout,
107 None,
108 ),
109 );
110 }
111
112 #[derive(Resource)]
113 struct ImageMaterialBindGroupLayout(BindGroupLayoutDescriptor);
114
115 #[derive(Resource)]
116 struct ImageMaterialBindGroupSampler(Sampler);
117
118 #[derive(Component)]
119 struct ImageMaterial3d(Handle<ImageMaterial>);
120
121 impl AsAssetId for ImageMaterial3d {
122 type Asset = ImageMaterial;
123
124 fn as_asset_id(&self) -> AssetId<Self::Asset> {
125 self.0.id()
126 }
127 }
128
129 #[derive(Asset, TypePath, AsBindGroup, Debug, Clone)]
130 struct ImageMaterial {
131 image: Handle<Image>,
132 }
133
134 impl ErasedRenderAsset for ImageMaterial {
135 type SourceAsset = ImageMaterial;
136 type ErasedAsset = PreparedMaterial;
137 type Param = (
138 SRes<DrawFunctions<Opaque3d>>,
139 SRes<ImageMaterialBindGroupLayout>,
140 SRes<AssetServer>,
141 SResMut<MaterialBindGroupAllocators>,
142 SResMut<RenderMaterialBindings>,
143 SRes<RenderAssets<GpuImage>>,
144 SRes<ImageMaterialBindGroupSampler>,
145 );
146
147 fn prepare_asset(
148 source_asset: Self::SourceAsset,
149 asset_id: AssetId<Self::SourceAsset>,
150 (
151 opaque_draw_functions,
152 material_layout,
153 asset_server,
154 bind_group_allocators,
155 render_material_bindings,
156 gpu_images,
157 image_material_sampler,
158 ): &mut SystemParamItem<Self::Param>,
159 ) -> std::result::Result<Self::ErasedAsset, PrepareAssetError<Self::SourceAsset>> {
160 let material_layout = material_layout.0.clone();
161 let draw_function_id = opaque_draw_functions.read().id::<DrawMaterial>();
162 let bind_group_allocator = bind_group_allocators
163 .get_mut(&TypeId::of::<ImageMaterial>())
164 .unwrap();
165 let Some(image) = gpu_images.get(&source_asset.image) else {
166 return Err(PrepareAssetError::RetryNextUpdate(source_asset));
167 };
168 let mut unprepared = BindGroupBuilder {
169 binding_resources: UnpreparedBindingResources(vec![
170 (
171 0,
172 UnpreparedBindingResource::TextureView(
173 TextureViewDimension::D2,
174 image.texture_view.clone(),
175 ),
176 ),
177 (
178 1,
179 UnpreparedBindingResource::Sampler(
180 SamplerBindingType::NonFiltering,
181 image_material_sampler.0.clone(),
182 ),
183 ),
184 ]),
185 data_buffer: vec![],
186 };
187 let binding = match render_material_bindings.entry(asset_id.into()) {
188 Entry::Occupied(mut occupied_entry) => {
189 bind_group_allocator.free(*occupied_entry.get());
190 let new_binding =
191 bind_group_allocator.allocate_unprepared(&mut unprepared, &material_layout);
192 *occupied_entry.get_mut() = new_binding;
193 new_binding
194 }
195 Entry::Vacant(vacant_entry) => *vacant_entry.insert(
196 bind_group_allocator.allocate_unprepared(&mut unprepared, &material_layout),
197 ),
198 };
199
200 let mut properties = MaterialProperties {
201 material_layout: Some(material_layout),
202 mesh_pipeline_key_bits: ErasedMeshPipelineKey::new(MeshPipelineKey::empty()),
203 base_specialize: Some(base_specialize),
204 ..Default::default()
205 };
206 properties.add_draw_function(MainPassOpaqueDrawFunction, draw_function_id);
207 properties.add_shader(MaterialFragmentShader, asset_server.load(SHADER_ASSET_PATH));
208
209 Ok(PreparedMaterial {
210 binding,
211 properties: Arc::new(properties),
212 })
213 }
214 }
215
216 /// set up a simple 3D scene
217 fn setup(
218 mut commands: Commands,
219 mut meshes: ResMut<Assets<Mesh>>,
220 mut materials: ResMut<Assets<ImageMaterial>>,
221 asset_server: Res<AssetServer>,
222 ) {
223 // cube
224 commands.spawn((
225 Mesh3d(meshes.add(Cuboid::new(2.0, 2.0, 2.0))),
226 ImageMaterial3d(materials.add(ImageMaterial {
227 image: asset_server.load("branding/icon.png"),
228 })),
229 Transform::from_xyz(0.0, 0.5, 0.0),
230 ));
231 // light
232 commands.spawn((
233 PointLight {
234 shadow_maps_enabled: true,
235 ..default()
236 },
237 Transform::from_xyz(4.0, 8.0, 4.0),
238 ));
239 // camera
240 commands.spawn((
241 Camera3d::default(),
242 Transform::from_xyz(-2.5, 4.5, 9.0).looking_at(Vec3::ZERO, Vec3::Y),
243 ));
244 }
245
246 fn extract_image_materials(
247 mut material_instances: ResMut<RenderMaterialInstances>,
248 changed_meshes_query: Extract<
249 Query<
250 (Entity, &ViewVisibility, &ImageMaterial3d),
251 Or<(Changed<ViewVisibility>, Changed<ImageMaterial3d>)>,
252 >,
253 >,
254 ) {
255 let last_change_tick = material_instances.current_change_tick;
256
257 for (entity, view_visibility, material) in &changed_meshes_query {
258 if view_visibility.get() {
259 material_instances.instances.insert(
260 entity.into(),
261 RenderMaterialInstance {
262 asset_id: material.0.id().untyped(),
263 last_change_tick,
264 },
265 );
266 } else {
267 material_instances
268 .instances
269 .remove(&MainEntity::from(entity));
270 }
271 }
272 }
273
274 fn check_entities_needing_specialization(
275 needs_specialization: Query<
276 Entity,
277 (
278 Or<(
279 Changed<Mesh3d>,
280 AssetChanged<Mesh3d>,
281 Changed<ImageMaterial3d>,
282 AssetChanged<ImageMaterial3d>,
283 )>,
284 With<ImageMaterial3d>,
285 ),
286 >,
287 mut par_local: Local<Parallel<Vec<Entity>>>,
288 mut entities_needing_specialization: ResMut<EntitiesNeedingSpecialization<ImageMaterial>>,
289 mut removed_mesh_3d_components: RemovedComponents<Mesh3d>,
290 mut removed_mesh_material_3d_components: RemovedComponents<ImageMaterial3d>,
291 ) {
292 entities_needing_specialization.changed.clear();
293 entities_needing_specialization.removed.clear();
294
295 // Gather all entities that need their specializations regenerated.
296 needs_specialization
297 .par_iter()
298 .for_each(|entity| par_local.borrow_local_mut().push(entity));
299 par_local.drain_into(&mut entities_needing_specialization.changed);
300
301 // All entities that removed their `Mesh3d` or `ImageMaterial3d` components
302 // need to have their specializations removed as well.
303 for entity in removed_mesh_3d_components
304 .read()
305 .chain(removed_mesh_material_3d_components.read())
306 {
307 entities_needing_specialization.removed.push(entity);
308 }
309 }
310
311 fn extract_image_materials_needing_specialization(
312 entities_needing_specialization: Extract<Res<EntitiesNeedingSpecialization<ImageMaterial>>>,
313 mut dirty_specializations: ResMut<DirtySpecializations>,
314 ) {
315 // Drain the list of entities needing specialization from the main world
316 // into the render-world `DirtySpecializations` table.
317 for entity in entities_needing_specialization.changed.iter() {
318 dirty_specializations
319 .changed_renderables
320 .insert(MainEntity::from(*entity));
321 }
322 }
323
324 /// A system that adds entities that were judged to need their specializations
325 /// removed to the appropriate table in [`DirtySpecializations`].
326 fn extract_image_materials_that_need_specializations_removed(
327 entities_needing_specialization: Extract<Res<EntitiesNeedingSpecialization<ImageMaterial>>>,
328 mut dirty_specializations: ResMut<DirtySpecializations>,
329 ) {
330 for entity in entities_needing_specialization.removed.iter() {
331 dirty_specializations
332 .removed_renderables
333 .insert(MainEntity::from(*entity));
334 }
335 }