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>();
    }
}