1use std::ops::Range;
14
15use bevy::camera::Viewport;
16use bevy::core_pipeline::core_3d::TransparentSortingInfo3d;
17use bevy::mesh::{MeshAttributeCompressionFlags, MeshCompressionArgs};
18use bevy::pbr::{self, MeshPipelineSystems, SetMeshViewEmptyBindGroup, ViewKeyCache};
19use bevy::{
20 camera::MainPassResolutionOverride,
21 core_pipeline::{core_3d::main_opaque_pass_3d, schedule::Core3d, Core3dSystems},
22 ecs::{
23 entity::EntityHash,
24 system::{lifetimeless::SRes, SystemParamItem},
25 },
26 math::FloatOrd,
27 mesh::MeshVertexBufferLayoutRef,
28 pbr::{
29 DrawMesh, MeshInputUniform, MeshPipeline, MeshPipelineKey, MeshPipelineViewLayoutKey,
30 MeshUniform, RenderMeshInstances, SetMeshBindGroup, SetMeshViewBindGroup,
31 },
32 platform::collections::HashSet,
33 prelude::*,
34 render::{
35 batching::{
36 gpu_preprocessing::{
37 batch_and_prepare_sorted_render_phase, BatchedInstanceBuffers,
38 IndirectParametersMetadata, UntypedPhaseIndirectParametersBuffers,
39 },
40 GetBatchData, GetFullBatchData,
41 },
42 camera::{DirtySpecializations, ExtractedCamera, PendingQueues},
43 extract_component::{ExtractComponent, ExtractComponentPlugin},
44 mesh::{allocator::MeshAllocator, RenderMesh},
45 render_asset::RenderAssets,
46 render_phase::{
47 sort_phase_system, AddRenderCommand, CachedRenderPipelinePhaseItem, DrawFunctionId,
48 DrawFunctions, PhaseItem, PhaseItemExtraIndex, SetItemPipeline, SortedPhaseItem,
49 SortedRenderPhasePlugin, ViewSortedRenderPhases,
50 },
51 render_resource::{
52 CachedRenderPipelineId, ColorTargetState, ColorWrites, Face, FragmentState,
53 PipelineCache, PrimitiveState, RenderPassDescriptor, RenderPipelineDescriptor,
54 SpecializedMeshPipeline, SpecializedMeshPipelineError, SpecializedMeshPipelines,
55 VertexState,
56 },
57 renderer::{RenderContext, ViewQuery},
58 sync_world::{MainEntity, MainEntityHashSet},
59 view::{ExtractedView, RenderVisibleEntities, RetainedViewEntity, ViewTarget},
60 Extract, Render, RenderApp, RenderDebugFlags, RenderStartup, RenderSystems,
61 },
62};
63use indexmap::IndexMap;
64use nonmax::NonMaxU32;
65
66const SHADER_ASSET_PATH: &str = "shaders/custom_stencil.wesl";
67
68fn main() {
69 App::new()
70 .add_plugins((DefaultPlugins, MeshStencilPhasePlugin))
71 .add_systems(Startup, setup)
72 .run();
73}
74
75fn setup(
76 mut commands: Commands,
77 mut meshes: ResMut<Assets<Mesh>>,
78 mut materials: ResMut<Assets<StandardMaterial>>,
79) {
80 commands.spawn((
82 Mesh3d(meshes.add(Circle::new(4.0))),
83 MeshMaterial3d(materials.add(Color::WHITE)),
84 Transform::from_rotation(Quat::from_rotation_x(-std::f32::consts::FRAC_PI_2)),
85 ));
86 commands.spawn((
90 Mesh3d(meshes.add(Cuboid::new(1.0, 1.0, 1.0))),
91 MeshMaterial3d(materials.add(Color::srgb_u8(124, 144, 255))),
92 Transform::from_xyz(0.0, 0.5, 0.0),
93 DrawStencil,
96 ));
97 commands.spawn((
98 Mesh3d(
99 meshes.add(
100 Sphere::new(0.5)
101 .mesh()
102 .build()
103 .compressed_mesh(&MeshCompressionArgs::regular())
104 .unwrap(),
105 ),
106 ),
107 MeshMaterial3d(materials.add(Color::srgb_u8(124, 255, 144))),
108 Transform::from_xyz(2.0, 0.5, 0.0),
109 DrawStencil,
112 ));
113 commands.spawn((
115 PointLight {
116 shadow_maps_enabled: true,
117 ..default()
118 },
119 Transform::from_xyz(4.0, 8.0, 4.0),
120 ));
121 commands.spawn((
123 Camera3d::default(),
124 Transform::from_xyz(-2.0, 4.5, 9.0).looking_at(Vec3::ZERO, Vec3::Y),
125 Msaa::Off,
127 ));
128}
129
130#[derive(Component, ExtractComponent, Clone, Copy, Default)]
131#[extract_app(RenderApp)]
132struct DrawStencil;
133
134struct MeshStencilPhasePlugin;
135impl Plugin for MeshStencilPhasePlugin {
136 fn build(&self, app: &mut App) {
137 app.add_plugins((
138 ExtractComponentPlugin::<DrawStencil>::default(),
139 SortedRenderPhasePlugin::<Stencil3d, MeshPipeline>::new(RenderDebugFlags::default()),
140 ));
141 let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
143 return;
144 };
145 render_app
146 .init_resource::<SpecializedMeshPipelines<StencilPipeline>>()
147 .init_resource::<DrawFunctions<Stencil3d>>()
148 .add_render_command::<Stencil3d, DrawMesh3dStencil>()
149 .init_resource::<ViewSortedRenderPhases<Stencil3d>>()
150 .init_resource::<PendingCustomMeshQueues>()
151 .add_systems(
152 RenderStartup,
153 init_stencil_pipeline.after(MeshPipelineSystems),
154 )
155 .add_systems(ExtractSchedule, extract_camera_phases)
156 .add_systems(
157 Render,
158 (
159 queue_custom_meshes.in_set(RenderSystems::QueueMeshes),
160 sort_phase_system::<Stencil3d>.in_set(RenderSystems::PhaseSort),
161 batch_and_prepare_sorted_render_phase::<Stencil3d, StencilPipeline>
162 .in_set(RenderSystems::PrepareResources),
163 ),
164 )
165 .add_systems(
166 Core3d,
167 custom_draw_system
168 .after(main_opaque_pass_3d)
169 .in_set(Core3dSystems::MainPass),
170 );
171 }
172}
173
174#[derive(Resource)]
175struct StencilPipeline {
176 mesh_pipeline: MeshPipeline,
181 shader_handle: Handle<Shader>,
184}
185
186fn init_stencil_pipeline(
187 mut commands: Commands,
188 mesh_pipeline: Res<MeshPipeline>,
189 asset_server: Res<AssetServer>,
190) {
191 commands.insert_resource(StencilPipeline {
192 mesh_pipeline: mesh_pipeline.clone(),
193 shader_handle: asset_server.load(SHADER_ASSET_PATH),
194 });
195}
196
197impl SpecializedMeshPipeline for StencilPipeline {
200 type Key = MeshPipelineKey;
201
202 fn specialize(
203 &self,
204 key: Self::Key,
205 layout: &MeshVertexBufferLayoutRef,
206 ) -> Result<RenderPipelineDescriptor, SpecializedMeshPipelineError> {
207 let mut shader_defs = Vec::new();
208 let mut vertex_attributes = Vec::new();
210 if layout.0.contains(Mesh::ATTRIBUTE_POSITION) {
211 if layout
213 .0
214 .get_attribute_compression()
215 .contains(MeshAttributeCompressionFlags::COMPRESS_POSITION)
216 {
217 shader_defs.push("VERTEX_POSITIONS_COMPRESSED".into());
218 }
219 vertex_attributes.push(Mesh::ATTRIBUTE_POSITION.at_shader_location(0));
221 }
222 let vertex_buffer_layout = layout.0.get_layout(&vertex_attributes)?;
224 let view_layout = self
225 .mesh_pipeline
226 .get_view_layout(MeshPipelineViewLayoutKey::from(key));
227 Ok(RenderPipelineDescriptor {
228 label: Some("Specialized Mesh Pipeline".into()),
229 layout: vec![
232 view_layout.main_layout,
234 view_layout.empty_layout,
236 self.mesh_pipeline.mesh_layouts.model_only.clone(),
238 ],
239 vertex: VertexState {
240 shader: self.shader_handle.clone(),
241 shader_defs: shader_defs.clone(),
242 buffers: vec![vertex_buffer_layout],
243 ..default()
244 },
245 fragment: Some(FragmentState {
246 shader: self.shader_handle.clone(),
247 shader_defs,
248 targets: vec![Some(ColorTargetState {
249 format: key.target_format(),
250 blend: None,
251 write_mask: ColorWrites::ALL,
252 })],
253 ..default()
254 }),
255 primitive: PrimitiveState {
256 topology: key.primitive_topology(),
257 strip_index_format: key.strip_index_format(),
258 cull_mode: Some(Face::Back),
259 ..default()
260 },
261 ..default()
264 })
265 }
266}
267
268type DrawMesh3dStencil = (
270 SetItemPipeline,
271 SetMeshViewBindGroup<0>,
273 SetMeshViewEmptyBindGroup<1>,
275 SetMeshBindGroup<2>,
277 DrawMesh,
279);
280
281struct Stencil3d {
289 pub sorting_info: TransparentSortingInfo3d,
292 pub distance: FloatOrd,
293 pub entity: (Entity, MainEntity),
294 pub pipeline: CachedRenderPipelineId,
295 pub draw_function: DrawFunctionId,
296 pub batch_range: Range<u32>,
297 pub extra_index: PhaseItemExtraIndex,
298 pub indexed: bool,
301}
302
303impl PhaseItem for Stencil3d {
305 #[inline]
306 fn entity(&self) -> Entity {
307 self.entity.0
308 }
309
310 #[inline]
311 fn main_entity(&self) -> MainEntity {
312 self.entity.1
313 }
314
315 #[inline]
316 fn draw_function(&self) -> DrawFunctionId {
317 self.draw_function
318 }
319
320 #[inline]
321 fn batch_range(&self) -> &Range<u32> {
322 &self.batch_range
323 }
324
325 #[inline]
326 fn batch_range_mut(&mut self) -> &mut Range<u32> {
327 &mut self.batch_range
328 }
329
330 #[inline]
331 fn extra_index(&self) -> PhaseItemExtraIndex {
332 self.extra_index.clone()
333 }
334
335 #[inline]
336 fn batch_range_and_extra_index_mut(&mut self) -> (&mut Range<u32>, &mut PhaseItemExtraIndex) {
337 (&mut self.batch_range, &mut self.extra_index)
338 }
339}
340
341impl SortedPhaseItem for Stencil3d {
342 type SortKey = FloatOrd;
343
344 #[inline]
345 fn sort_key(&self) -> Self::SortKey {
346 self.distance
347 }
348
349 #[inline]
350 fn sort(items: &mut IndexMap<(Entity, MainEntity), Stencil3d, EntityHash>) {
351 items.sort_by_key(|_, phase_item: &Stencil3d| phase_item.distance);
352 }
353
354 fn recalculate_sort_keys(
355 items: &mut IndexMap<(Entity, MainEntity), Self, EntityHash>,
356 view: &ExtractedView,
357 ) {
358 let rangefinder = view.rangefinder3d();
360 for item in items.values_mut() {
361 item.distance = FloatOrd(item.sorting_info.sort_distance(&rangefinder));
362 }
363 }
364
365 #[inline]
366 fn indexed(&self) -> bool {
367 self.indexed
368 }
369}
370
371impl CachedRenderPipelinePhaseItem for Stencil3d {
372 #[inline]
373 fn cached_pipeline(&self) -> CachedRenderPipelineId {
374 self.pipeline
375 }
376}
377
378impl GetBatchData for StencilPipeline {
379 type Param = (SRes<RenderMeshInstances>, SRes<MeshAllocator>);
380 type BatchSetCompareData = AssetId<Mesh>;
384 type BatchCompareData = ();
385 type BufferData = MeshUniform;
386
387 fn get_batch_data(
388 (mesh_instances, mesh_allocator): &SystemParamItem<Self::Param>,
389 (_entity, main_entity): (Entity, MainEntity),
390 ) -> Option<(
391 Self::BufferData,
392 Option<(Self::BatchSetCompareData, Self::BatchCompareData)>,
393 )> {
394 let RenderMeshInstances::CpuBuilding(ref mesh_instances) = **mesh_instances else {
395 error!(
396 "`get_batch_data` should never be called in GPU mesh uniform \
397 building mode"
398 );
399 return None;
400 };
401 let mesh_instance = mesh_instances.get(&main_entity)?;
402 let first_vertex_index =
403 match mesh_allocator.mesh_vertex_slice(&mesh_instance.mesh_asset_id()) {
404 Some(mesh_vertex_slice) => mesh_vertex_slice.range.start,
405 None => 0,
406 };
407 let metadata_index = mesh_allocator
408 .mesh_metadata_slice(&mesh_instance.mesh_asset_id())
409 .map(|mesh_metadata_slice| mesh_metadata_slice.range.start);
410
411 let mesh_uniform = MeshUniform::new(
412 &mesh_instance.transforms,
413 first_vertex_index,
414 mesh_instance.material_bindings_index().slot,
415 None,
416 None,
417 None,
418 Some(mesh_instance.tag()),
419 metadata_index,
420 );
421 Some((mesh_uniform, None))
422 }
423}
424
425impl GetFullBatchData for StencilPipeline {
426 type BufferInputData = MeshInputUniform;
427
428 fn get_index_and_compare_data(
429 (mesh_instances, _): &SystemParamItem<Self::Param>,
430 main_entity: MainEntity,
431 ) -> Option<(
432 NonMaxU32,
433 Option<(Self::BatchSetCompareData, Self::BatchCompareData)>,
434 )> {
435 let RenderMeshInstances::GpuBuilding(ref mesh_instances) = **mesh_instances else {
437 error!(
438 "`get_index_and_compare_data` should never be called in CPU mesh uniform building \
439 mode"
440 );
441 return None;
442 };
443 let mesh_instance = mesh_instances.get(&main_entity)?;
444 Some((
445 NonMaxU32::new(mesh_instance.gpu_specific.current_uniform_index())?,
446 mesh_instance
447 .should_batch()
448 .then_some((mesh_instance.mesh_asset_id(), ())),
449 ))
450 }
451
452 fn get_binned_batch_data(
453 (mesh_instances, mesh_allocator): &SystemParamItem<Self::Param>,
454 main_entity: MainEntity,
455 ) -> Option<Self::BufferData> {
456 let RenderMeshInstances::CpuBuilding(ref mesh_instances) = **mesh_instances else {
457 error!(
458 "`get_binned_batch_data` should never be called in GPU mesh uniform building mode"
459 );
460 return None;
461 };
462 let mesh_instance = mesh_instances.get(&main_entity)?;
463 let first_vertex_index =
464 match mesh_allocator.mesh_vertex_slice(&mesh_instance.mesh_asset_id()) {
465 Some(mesh_vertex_slice) => mesh_vertex_slice.range.start,
466 None => 0,
467 };
468 let metadata_index = mesh_allocator
469 .mesh_metadata_slice(&mesh_instance.mesh_asset_id())
470 .map(|mesh_metadata_slice| mesh_metadata_slice.range.start);
471
472 Some(MeshUniform::new(
473 &mesh_instance.transforms,
474 first_vertex_index,
475 mesh_instance.material_bindings_index().slot,
476 None,
477 None,
478 None,
479 Some(mesh_instance.tag()),
480 metadata_index,
481 ))
482 }
483
484 fn write_batch_indirect_parameters_metadata(
485 indexed: bool,
486 base_output_index: u32,
487 batch_set_index: Option<NonMaxU32>,
488 indirect_parameters_buffers: &mut UntypedPhaseIndirectParametersBuffers,
489 indirect_parameters_offset: u32,
490 ) {
491 let indirect_parameters = IndirectParametersMetadata {
495 base_output_index,
496 batch_set_index: match batch_set_index {
497 None => !0,
498 Some(batch_set_index) => u32::from(batch_set_index),
499 },
500 mesh_index: 0,
502 early_instance_count: 0,
503 late_instance_count: 0,
504 };
505
506 if indexed {
507 indirect_parameters_buffers
508 .indexed
509 .set(indirect_parameters_offset, indirect_parameters);
510 } else {
511 indirect_parameters_buffers
512 .non_indexed
513 .set(indirect_parameters_offset, indirect_parameters);
514 }
515 }
516
517 fn get_binned_index(
518 _param: &SystemParamItem<Self::Param>,
519 _query_item: MainEntity,
520 ) -> Option<NonMaxU32> {
521 None
522 }
523}
524
525fn extract_camera_phases(
529 mut stencil_phases: ResMut<ViewSortedRenderPhases<Stencil3d>>,
530 cameras: Extract<Query<(Entity, &Camera), With<Camera3d>>>,
531 mut live_entities: Local<HashSet<RetainedViewEntity>>,
532) {
533 live_entities.clear();
534 for (main_entity, camera) in &cameras {
535 if !camera.is_active {
536 continue;
537 }
538 let retained_view_entity = RetainedViewEntity::new(main_entity.into(), None, 0);
540
541 stencil_phases.prepare_for_new_frame(retained_view_entity);
542 live_entities.insert(retained_view_entity);
543 }
544
545 stencil_phases.retain(|camera_entity, _| live_entities.contains(camera_entity));
547}
548
549#[derive(Default, Deref, DerefMut, Resource)]
554struct PendingCustomMeshQueues(pub PendingQueues);
555
556fn queue_custom_meshes(
560 custom_draw_functions: Res<DrawFunctions<Stencil3d>>,
561 mut pipelines: ResMut<SpecializedMeshPipelines<StencilPipeline>>,
562 pipeline_cache: Res<PipelineCache>,
563 custom_draw_pipeline: Res<StencilPipeline>,
564 render_meshes: Res<RenderAssets<RenderMesh>>,
565 render_mesh_instances: Res<RenderMeshInstances>,
566 maybe_batched_instance_buffers: Option<
567 Res<BatchedInstanceBuffers<MeshUniform, MeshInputUniform>>,
568 >,
569 mut custom_render_phases: ResMut<ViewSortedRenderPhases<Stencil3d>>,
570 mut views: Query<(&ExtractedView, &RenderVisibleEntities)>,
571 view_key_cache: Res<ViewKeyCache>,
572 dirty_specializations: Res<DirtySpecializations>,
573 mut pending_custom_mesh_queues: ResMut<PendingCustomMeshQueues>,
574 has_marker: Query<(), With<DrawStencil>>,
575 mut mesh_instances_queued_this_iteration_scratch_space: Local<MainEntityHashSet>,
576) {
577 for (view, visible_entities) in &mut views {
578 let Some(custom_phase) = custom_render_phases.get_mut(&view.retained_view_entity) else {
579 continue;
580 };
581 let draw_custom = custom_draw_functions.read().id::<DrawMesh3dStencil>();
582
583 let Some(&view_key) = view_key_cache.get(&view.retained_view_entity) else {
584 continue;
585 };
586
587 let Some(render_visible_mesh_entities) = visible_entities.get::<Mesh3d>() else {
589 continue;
590 };
591
592 let view_pending_custom_mesh_queues =
593 pending_custom_mesh_queues.prepare_for_new_frame(view.retained_view_entity);
594
595 for &main_entity in dirty_specializations
597 .iter_to_dequeue(view.retained_view_entity, render_visible_mesh_entities)
598 {
599 custom_phase.remove(Entity::PLACEHOLDER, main_entity);
600 }
601
602 for (render_entity, visible_entity) in dirty_specializations.iter_to_queue(
603 view.retained_view_entity,
604 render_visible_mesh_entities,
605 &view_pending_custom_mesh_queues.prev_frame,
606 &mut mesh_instances_queued_this_iteration_scratch_space,
607 ) {
608 if has_marker.get(*render_entity).is_err() {
610 continue;
611 }
612 let Some(mesh_instance) = render_mesh_instances.render_mesh_queue_data(*visible_entity)
613 else {
614 view_pending_custom_mesh_queues
618 .current_frame
619 .insert((*render_entity, *visible_entity));
620 continue;
621 };
622 let Some(mesh) = render_meshes.get(mesh_instance.mesh_asset_id()) else {
623 continue;
624 };
625
626 let mut mesh_key = view_key;
630 mesh_key |= MeshPipelineKey::from_primitive_topology_and_strip_index(
631 mesh.primitive_topology(),
632 mesh.index_format(),
633 );
634
635 let pipeline_id = pipelines.specialize(
636 &pipeline_cache,
637 &custom_draw_pipeline,
638 mesh_key,
639 &mesh.layout,
640 );
641 let pipeline_id = match pipeline_id {
642 Ok(id) => id,
643 Err(err) => {
644 error!("{}", err);
645 continue;
646 }
647 };
648 custom_phase.add_retained(Stencil3d {
651 sorting_info: TransparentSortingInfo3d::Sorted {
652 mesh_center: pbr::get_mesh_instance_world_from_local(
653 *visible_entity,
654 mesh_instance.current_uniform_index,
655 &render_mesh_instances,
656 maybe_batched_instance_buffers.as_deref(),
657 )
658 .transform_point3(
659 render_meshes
660 .get(mesh_instance.mesh_asset_id())
661 .unwrap()
662 .aabb_center,
663 ),
664 depth_bias: 0.0,
665 },
666 distance: FloatOrd(0.0),
667 entity: (Entity::PLACEHOLDER, *visible_entity),
668 pipeline: pipeline_id,
669 draw_function: draw_custom,
670 batch_range: 0..1,
672 extra_index: PhaseItemExtraIndex::None,
673 indexed: mesh.indexed(),
674 });
675 }
676 }
677}
678
679fn custom_draw_system(
680 world: &World,
681 view: ViewQuery<(
682 &ExtractedCamera,
683 &ExtractedView,
684 &ViewTarget,
685 Option<&MainPassResolutionOverride>,
686 )>,
687 stencil_phases: Res<ViewSortedRenderPhases<Stencil3d>>,
688 mut ctx: RenderContext,
689) {
690 let view_entity = view.entity();
691 let (camera, extracted_view, target, resolution_override) = view.into_inner();
692
693 let Some(stencil_phase) = stencil_phases.get(&extracted_view.retained_view_entity) else {
694 return;
695 };
696
697 let mut render_pass = ctx.begin_tracked_render_pass(RenderPassDescriptor {
698 label: Some("stencil pass"),
699 color_attachments: &[Some(target.get_color_attachment())],
703 depth_stencil_attachment: None,
705 timestamp_writes: None,
706 occlusion_query_set: None,
707 multiview_mask: None,
708 });
709
710 if let Some(viewport) =
711 Viewport::from_viewport_and_override(camera.viewport.as_ref(), resolution_override)
712 {
713 render_pass.set_camera_viewport(&viewport);
714 }
715
716 if let Err(err) = stencil_phase.render(&mut render_pass, world, view_entity) {
717 error!("Error encountered while rendering the stencil phase {err:?}");
718 }
719}