Skip to content

Instantly share code, notes, and snippets.

@viridia

viridia/node_shader_pipeline.rs

Created July 15, 2024 02:16
Show Gist options
  • Save viridia/f3abbe83fcf970f6e4ec5c43416c5f96 to your computer and use it in GitHub Desktop.
Save viridia/f3abbe83fcf970f6e4ec5c43416c5f96 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
node_shader_pipeline.rs
use bevy::{
core_pipeline::core_3d::{Opaque3d, Opaque3dBinKey, Transparent3d, CORE_3D_DEPTH_FORMAT},
pbr::{
DrawMesh, MeshPipeline, MeshPipelineKey, MeshPipelineViewLayoutKey, RenderMeshInstances,
SetMeshBindGroup, SetMeshViewBindGroup,
},
prelude::*,
render::{
extract_component::{ExtractComponent, ExtractComponentPlugin},
render_asset::RenderAssets,
render_phase::{
AddRenderCommand, BinnedRenderPhaseType, DrawFunctions, SetItemPipeline,
ViewBinnedRenderPhases,
},
render_resource::{
BlendState, ColorTargetState, ColorWrites, CompareFunction, DepthBiasState,
DepthStencilState, Face, FragmentState, FrontFace, MultisampleState, PipelineCache,
PolygonMode, PrimitiveState, RenderPipelineDescriptor, SpecializedRenderPipeline,
SpecializedRenderPipelines, StencilState, TextureFormat, VertexBufferLayout,
VertexFormat, VertexState, VertexStepMode,
},
texture::BevyDefault,
view::{ExtractedView, ViewTarget, VisibleEntities},
Render, RenderApp, RenderSet,
},
};
/// Component that associates a generated shader to a mesh.
#[derive(Component, Default, ExtractComponent, Clone)]
pub struct NodeShaderMesh3d(pub Handle<Shader>);
/// Custom pipeline for meshes with vertex colors
#[derive(Resource)]
pub struct NodeShaderMesh3dPipeline {
/// this pipeline wraps the standard [`MeshPipeline`]
mesh_pipeline: MeshPipeline,
}
#[derive(Clone, Hash, PartialEq, Eq)]
pub struct NodeShaderMesh3dPipelineKey {
/// Handle to the generated shader.
shader: Handle<Shader>,
/// Key for the mesh pipeline
mesh_key: MeshPipelineKey,
}
impl FromWorld for NodeShaderMesh3dPipeline {
fn from_world(world: &mut World) -> Self {
Self {
mesh_pipeline: MeshPipeline::from_world(world),
}
}
}
// We implement `SpecializedPipeline` to customize the default rendering from `MeshPipeline`
impl SpecializedRenderPipeline for NodeShaderMesh3dPipeline {
type Key = NodeShaderMesh3dPipelineKey;
fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
// Customize how to store the meshes' vertex attributes in the vertex buffer
// Our meshes only have positions
let formats = vec![VertexFormat::Float32x3];
let vertex_layout =
VertexBufferLayout::from_vertex_formats(VertexStepMode::Vertex, formats);
let format = match key.mesh_key.contains(MeshPipelineKey::HDR) {
true => ViewTarget::TEXTURE_FORMAT_HDR,
false => TextureFormat::bevy_default(),
};
RenderPipelineDescriptor {
vertex: VertexState {
// Use our custom shader
shader: key.shader.clone(),
entry_point: "vertex".into(),
shader_defs: vec![],
// Use our custom vertex buffer
buffers: vec![vertex_layout],
},
fragment: Some(FragmentState {
// Use our custom shader
shader: key.shader.clone(),
shader_defs: vec![],
entry_point: "fragment".into(),
targets: vec![Some(ColorTargetState {
format,
blend: Some(BlendState::ALPHA_BLENDING),
write_mask: ColorWrites::ALL,
})],
}),
layout: vec![
// Bind group 0 is the view uniform
self.mesh_pipeline
.get_view_layout(MeshPipelineViewLayoutKey::from(key.mesh_key))
.clone(),
// Bind group 1 is the mesh uniform
self.mesh_pipeline.mesh_layouts.model_only.clone(),
],
push_constant_ranges: vec![],
primitive: PrimitiveState {
front_face: FrontFace::Ccw,
cull_mode: Some(Face::Back),
unclipped_depth: false,
polygon_mode: PolygonMode::Fill,
conservative: false,
topology: key.mesh_key.primitive_topology(),
strip_index_format: None,
},
depth_stencil: Some(DepthStencilState {
format: CORE_3D_DEPTH_FORMAT,
depth_write_enabled: true,
depth_compare: CompareFunction::Greater,
stencil: StencilState::default(),
bias: DepthBiasState::default(),
}),
multisample: MultisampleState {
count: key.mesh_key.msaa_samples(),
mask: !0,
alpha_to_coverage_enabled: false,
},
label: Some("node_shader_mesh_pipeline".into()),
}
}
}
// This specifies how to render a node-shader mesh
type DrawColoredMesh = (
// Set the pipeline
SetItemPipeline,
// Set the view uniform as bind group 0
SetMeshViewBindGroup<0>,
// Set the mesh uniform as bind group 1
SetMeshBindGroup<1>,
// Draw the mesh
DrawMesh,
);
/// A render-world system that enqueues the entity with custom rendering into
/// the opaque render phases of each view.
fn queue_node_shader_mesh(
pipeline_cache: Res<PipelineCache>,
ns_mesh_pipeline: Res<NodeShaderMesh3dPipeline>,
msaa: Res<Msaa>,
render_meshes: Res<RenderAssets<RenderMesh>>,
render_mesh_instances: Res<RenderMeshInstances>,
mut opaque_render_phases: ResMut<ViewBinnedRenderPhases<Opaque3d>>,
opaque_draw_functions: Res<DrawFunctions<Opaque3d>>,
mut specialized_render_pipelines: ResMut<SpecializedRenderPipelines<NodeShaderMesh3dPipeline>>,
views: Query<(Entity, &VisibleEntities, &ExtractedView)>,
) {
let draw_custom_phase_item = opaque_draw_functions.read().id::<DrawColoredMesh>();
// Render phases are per-view, so we need to iterate over all views so that
// the entity appears in them.
for (view_entity, view_visible_entities, view) in views.iter() {
let mesh_key = MeshPipelineKey::from_msaa_samples(msaa.samples())
| MeshPipelineKey::from_hdr(view.hdr);
let Some(opaque_phase) = opaque_render_phases.get_mut(&view_entity) else {
continue;
};
// Find all the custom rendered entities that are visible from this view.
for &entity in view_visible_entities.get::<NodeShaderMesh3d>().iter() {
// if has_preview_mesh_marker.get(*visible_entity).is_err() {
// continue;
// }
let Some(mesh_instance) = render_mesh_instances.render_mesh_queue_data(*entity) else {
continue;
};
// Get our specialized pipeline
let mut mesh_key = mesh_key;
if let Some(mesh) = render_meshes.get(mesh_instance.mesh_asset_id) {
mesh_key |= MeshPipelineKey::from_primitive_topology(mesh.primitive_topology());
}
let pipeline_id = specialized_render_pipelines.specialize(
&pipeline_cache,
&ns_mesh_pipeline,
NodeShaderMesh3dPipelineKey {
shader: shader.clone(),
mesh_key,
},
);
// Add the custom render item. We use the
// [`BinnedRenderPhaseType::NonMesh`] type to skip the special
// handling that Bevy has for meshes (preprocessing, indirect
// draws, etc.)
opaque_phase.add(
Opaque3dBinKey {
draw_function: draw_custom_phase_item,
pipeline: pipeline_id,
asset_id: AssetId::<Mesh>::invalid().untyped(),
material_bind_group_id: None,
lightmap_image: None,
},
entity,
BinnedRenderPhaseType::NonMesh,
);
}
}
}
/// Plugin that renders [`NodeShaderMesh3d`]s
pub struct NodeShaderMeshPlugin;
impl Plugin for NodeShaderMeshPlugin {
fn build(&self, app: &mut App) {
app.add_plugins(ExtractComponentPlugin::<NodeShaderMesh3d>::default());
// Register our custom draw function, and add our render systems
let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
return;
};
render_app
.add_render_command::<Transparent3d, DrawColoredMesh>()
.init_resource::<SpecializedRenderPipelines<NodeShaderMesh3dPipeline>>()
.add_systems(Render, queue_node_shader_mesh.in_set(RenderSet::Queue));
}
fn finish(&self, app: &mut App) {
// Register our custom pipeline
let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
return;
};
render_app.init_resource::<NodeShaderMesh3dPipeline>();
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment