use bevy::{
core_pipeline::core_3d::{Opaque3d, Opaque3dBatchSetKey, Opaque3dBinKey, CORE_3D_DEPTH_FORMAT},
ecs::{component::Tick, system::StaticSystemParam},
math::{vec3, vec4},
pbr::{
DrawMesh, MeshPipeline, MeshPipelineKey, MeshPipelineViewLayoutKey, RenderMeshInstances,
SetMeshBindGroup, SetMeshViewBindGroup,
},
prelude::*,
render::{
batching::{
gpu_preprocessing::{
self, PhaseBatchedInstanceBuffers, PhaseIndirectParametersBuffers,
PreprocessWorkItem, UntypedPhaseBatchedInstanceBuffers,
},
GetBatchData, GetFullBatchData,
},
experimental::occlusion_culling::OcclusionCulling,
extract_component::{ExtractComponent, ExtractComponentPlugin},
mesh::{Indices, MeshVertexBufferLayoutRef, PrimitiveTopology, RenderMesh},
render_asset::{RenderAssetUsages, RenderAssets},
render_phase::{
AddRenderCommand, BinnedRenderPhaseType, DrawFunctions, SetItemPipeline,
ViewBinnedRenderPhases,
},
render_resource::{
ColorTargetState, ColorWrites, CompareFunction, DepthStencilState, Face, FragmentState,
FrontFace, MultisampleState, PipelineCache, PolygonMode, PrimitiveState,
RenderPipelineDescriptor, SpecializedMeshPipeline, SpecializedMeshPipelineError,
SpecializedMeshPipelines, TextureFormat, VertexState,
},
view::NoIndirectDrawing,
view::{self, ExtractedView, RenderVisibleEntities, ViewTarget, VisibilityClass},
Render, RenderApp, RenderSet,
},
};
const SHADER_ASSET_PATH: &str = "shaders/specialized_mesh_pipeline.wgsl";
fn main() {
App::new()
.add_plugins(DefaultPlugins)
.add_plugins(CustomRenderedMeshPipelinePlugin)
.add_systems(Startup, setup)
.run();
}
fn setup(mut commands: Commands, mut meshes: ResMut<Assets<Mesh>>) {
let mesh = Mesh::new(
PrimitiveTopology::TriangleList,
RenderAssetUsages::default(),
)
.with_inserted_indices(Indices::U32(vec![0, 1, 2]))
.with_inserted_attribute(
Mesh::ATTRIBUTE_POSITION,
vec![
vec3(-0.5, -0.5, 0.0),
vec3(0.5, -0.5, 0.0),
vec3(0.0, 0.25, 0.0),
],
)
.with_inserted_attribute(
Mesh::ATTRIBUTE_COLOR,
vec![
vec4(1.0, 0.0, 0.0, 1.0),
vec4(0.0, 1.0, 0.0, 1.0),
vec4(0.0, 0.0, 1.0, 1.0),
],
);
for (x, y) in [-0.5, 0.0, 0.5].into_iter().zip([-0.25, 0.5, -0.25]) {
commands.spawn((
CustomRenderedEntity,
Mesh3d(meshes.add(mesh.clone())),
Transform::from_xyz(x, y, 0.0),
));
}
commands.spawn((
Camera3d::default(),
Transform::from_xyz(0.0, 0.0, 3.0).looking_at(Vec3::ZERO, Vec3::Y),
));
}
struct CustomRenderedMeshPipelinePlugin;
impl Plugin for CustomRenderedMeshPipelinePlugin {
fn build(&self, app: &mut App) {
app.add_plugins(ExtractComponentPlugin::<CustomRenderedEntity>::default());
let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
return;
};
render_app
.init_resource::<SpecializedMeshPipelines<CustomMeshPipeline>>()
.add_render_command::<Opaque3d, DrawSpecializedPipelineCommands>()
.add_systems(Render, queue_custom_mesh_pipeline.in_set(RenderSet::Queue));
}
fn finish(&self, app: &mut App) {
let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
return;
};
render_app.init_resource::<CustomMeshPipeline>();
}
}
#[derive(Clone, Component, ExtractComponent)]
#[require(VisibilityClass)]
#[component(on_add = view::add_visibility_class::<CustomRenderedEntity>)]
struct CustomRenderedEntity;
type DrawSpecializedPipelineCommands = (
SetItemPipeline,
SetMeshViewBindGroup<0>,
SetMeshBindGroup<1>,
DrawMesh,
);
#[derive(Resource)]
struct CustomMeshPipeline {
mesh_pipeline: MeshPipeline,
shader_handle: Handle<Shader>,
}
impl FromWorld for CustomMeshPipeline {
fn from_world(world: &mut World) -> Self {
let shader_handle: Handle<Shader> = world.resource::<AssetServer>().load(SHADER_ASSET_PATH);
Self {
mesh_pipeline: MeshPipeline::from_world(world),
shader_handle,
}
}
}
impl SpecializedMeshPipeline for CustomMeshPipeline {
type Key = MeshPipelineKey;
fn specialize(
&self,
mesh_key: Self::Key,
layout: &MeshVertexBufferLayoutRef,
) -> Result<RenderPipelineDescriptor, SpecializedMeshPipelineError> {
let mut vertex_attributes = Vec::new();
if layout.0.contains(Mesh::ATTRIBUTE_POSITION) {
vertex_attributes.push(Mesh::ATTRIBUTE_POSITION.at_shader_location(0));
}
if layout.0.contains(Mesh::ATTRIBUTE_COLOR) {
vertex_attributes.push(Mesh::ATTRIBUTE_COLOR.at_shader_location(1));
}
let vertex_buffer_layout = layout.0.get_layout(&vertex_attributes)?;
Ok(RenderPipelineDescriptor {
label: Some("Specialized Mesh Pipeline".into()),
layout: vec![
self.mesh_pipeline
.get_view_layout(MeshPipelineViewLayoutKey::from(mesh_key))
.clone(),
self.mesh_pipeline.mesh_layouts.model_only.clone(),
],
push_constant_ranges: vec![],
vertex: VertexState {
shader: self.shader_handle.clone(),
shader_defs: vec![],
entry_point: "vertex".into(),
buffers: vec![vertex_buffer_layout],
},
fragment: Some(FragmentState {
shader: self.shader_handle.clone(),
shader_defs: vec![],
entry_point: "fragment".into(),
targets: vec![Some(ColorTargetState {
format: if mesh_key.contains(MeshPipelineKey::HDR) {
ViewTarget::TEXTURE_FORMAT_HDR
} else {
TextureFormat::bevy_default()
},
blend: None,
write_mask: ColorWrites::ALL,
})],
}),
primitive: PrimitiveState {
topology: mesh_key.primitive_topology(),
front_face: FrontFace::Ccw,
cull_mode: Some(Face::Back),
polygon_mode: PolygonMode::Fill,
..default()
},
depth_stencil: Some(DepthStencilState {
format: CORE_3D_DEPTH_FORMAT,
depth_write_enabled: true,
depth_compare: CompareFunction::GreaterEqual,
stencil: default(),
bias: default(),
}),
multisample: MultisampleState {
count: mesh_key.msaa_samples(),
..MultisampleState::default()
},
zero_initialize_workgroup_memory: false,
})
}
}
fn queue_custom_mesh_pipeline(
pipeline_cache: Res<PipelineCache>,
custom_mesh_pipeline: Res<CustomMeshPipeline>,
(mut opaque_render_phases, opaque_draw_functions): (
ResMut<ViewBinnedRenderPhases<Opaque3d>>,
Res<DrawFunctions<Opaque3d>>,
),
mut specialized_mesh_pipelines: ResMut<SpecializedMeshPipelines<CustomMeshPipeline>>,
views: Query<(
&RenderVisibleEntities,
&ExtractedView,
&Msaa,
Has<NoIndirectDrawing>,
Has<OcclusionCulling>,
)>,
(render_meshes, render_mesh_instances): (
Res<RenderAssets<RenderMesh>>,
Res<RenderMeshInstances>,
),
param: StaticSystemParam<<MeshPipeline as GetBatchData>::Param>,
mut phase_batched_instance_buffers: ResMut<
PhaseBatchedInstanceBuffers<Opaque3d, <MeshPipeline as GetBatchData>::BufferData>,
>,
mut phase_indirect_parameters_buffers: ResMut<PhaseIndirectParametersBuffers<Opaque3d>>,
mut change_tick: Local<Tick>,
) {
let system_param_item = param.into_inner();
let UntypedPhaseBatchedInstanceBuffers {
ref mut data_buffer,
ref mut work_item_buffers,
ref mut late_indexed_indirect_parameters_buffer,
ref mut late_non_indexed_indirect_parameters_buffer,
..
} = phase_batched_instance_buffers.buffers;
let draw_function_id = opaque_draw_functions
.read()
.id::<DrawSpecializedPipelineCommands>();
for (view_visible_entities, view, msaa, no_indirect_drawing, gpu_occlusion_culling) in
views.iter()
{
let Some(opaque_phase) = opaque_render_phases.get_mut(&view.retained_view_entity) else {
continue;
};
let work_item_buffer = gpu_preprocessing::get_or_create_work_item_buffer::<Opaque3d>(
work_item_buffers,
view.retained_view_entity,
no_indirect_drawing,
gpu_occlusion_culling,
);
gpu_preprocessing::init_work_item_buffers(
work_item_buffer,
late_indexed_indirect_parameters_buffer,
late_non_indexed_indirect_parameters_buffer,
);
let view_key = MeshPipelineKey::from_msaa_samples(msaa.samples())
| MeshPipelineKey::from_hdr(view.hdr);
let mut mesh_batch_set_info = None;
for &(render_entity, visible_entity) in
view_visible_entities.get::<CustomRenderedEntity>().iter()
{
let Some(mesh_instance) = render_mesh_instances.render_mesh_queue_data(visible_entity)
else {
continue;
};
let Some(mesh) = render_meshes.get(mesh_instance.mesh_asset_id) else {
continue;
};
let mut mesh_key = view_key;
mesh_key |= MeshPipelineKey::from_primitive_topology(mesh.primitive_topology());
if mesh_batch_set_info.is_none() {
mesh_batch_set_info = Some(MeshBatchSetInfo {
indirect_parameters_index: phase_indirect_parameters_buffers
.buffers
.allocate(mesh.indexed(), 1),
is_indexed: mesh.indexed(),
});
}
let mesh_info = mesh_batch_set_info.unwrap();
let Some(input_index) =
MeshPipeline::get_binned_index(&system_param_item, visible_entity)
else {
continue;
};
let output_index = data_buffer.add() as u32;
let pipeline_id = specialized_mesh_pipelines
.specialize(
&pipeline_cache,
&custom_mesh_pipeline,
mesh_key,
&mesh.layout,
)
.expect("Failed to specialize mesh pipeline");
let next_change_tick = change_tick.get() + 1;
change_tick.set(next_change_tick);
opaque_phase.add(
Opaque3dBatchSetKey {
draw_function: draw_function_id,
pipeline: pipeline_id,
material_bind_group_index: None,
vertex_slab: default(),
index_slab: None,
lightmap_slab: None,
},
Opaque3dBinKey {
asset_id: AssetId::<Mesh>::invalid().untyped(),
},
(render_entity, visible_entity),
mesh_instance.current_uniform_index,
BinnedRenderPhaseType::BatchableMesh,
*change_tick,
);
work_item_buffer.push(
mesh.indexed(),
PreprocessWorkItem {
input_index: input_index.into(),
output_or_indirect_parameters_index: if no_indirect_drawing {
output_index
} else {
mesh_info.indirect_parameters_index
},
},
);
}
if let Some(mesh_info) = mesh_batch_set_info {
phase_indirect_parameters_buffers
.buffers
.add_batch_set(mesh_info.is_indexed, mesh_info.indirect_parameters_index);
}
}
}
#[derive(Clone, Copy)]
struct MeshBatchSetInfo {
indirect_parameters_index: u32,
is_indexed: bool,
}