colored_mesh_3d.rs

//! This example shows how to manually render 3d items using "mid level render apis" with a custom
//! pipeline for 3d meshes.
//! It doesn't use the [`Material`] abstraction, but changes the vertex buffer to include vertex color.
//!
//! [`Material`]: bevy::pbr::Material

use bevy::{
    core_pipeline::core_3d::{Transparent3d, CORE_3D_DEPTH_FORMAT},
    pbr::{
        DrawMesh, MeshPipeline, MeshPipelineKey, MeshPipelineViewLayoutKey, RenderMeshInstances,
        SetMeshBindGroup, SetMeshViewBindGroup,
    },
    prelude::*,
    render::{
        extract_component::{ExtractComponent, ExtractComponentPlugin},
        mesh::{Indices, PrimitiveTopology, RenderMesh},
        render_asset::{RenderAssetUsages, RenderAssets},
        render_phase::{
            AddRenderCommand, DrawFunctions, PhaseItemExtraIndex, SetItemPipeline,
            ViewSortedRenderPhases,
        },
        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, WithMesh},
        Render, RenderApp, RenderSet,
    },
};

fn main() {
    App::new()
        .add_plugins((DefaultPlugins, ColoredMeshPlugin))
        .add_systems(Startup, setup)
        .run();
}

fn setup(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<StandardMaterial>>,
) {
    // Build a quad
    let half_size = Vec2::splat(2.0);
    let rotation = Quat::from_rotation_arc(Vec3::Y, *Dir3::Y);
    let positions = vec![
        rotation * Vec3::new(half_size.x, 0.0, -half_size.y),
        rotation * Vec3::new(-half_size.x, 0.0, -half_size.y),
        rotation * Vec3::new(-half_size.x, 0.0, half_size.y),
        rotation * Vec3::new(half_size.x, 0.0, half_size.y),
    ];
    let indices = Indices::U32(vec![0, 1, 2, 0, 2, 3]);
    let mesh = Mesh::new(
        PrimitiveTopology::TriangleList,
        RenderAssetUsages::default(),
    )
    .with_inserted_indices(indices)
    .with_inserted_attribute(Mesh::ATTRIBUTE_POSITION, positions);

    commands.spawn((
        // We use a marker component to identify the custom colored meshes
        ColoredMesh3d,
        meshes.add(mesh),
        // This bundle's components are needed for something to be rendered
        SpatialBundle::INHERITED_IDENTITY,
    ));

    // cube
    commands.spawn(PbrBundle {
        mesh: meshes.add(Cuboid::new(1.0, 1.0, 1.0)),
        material: materials.add(Color::srgb_u8(124, 144, 255)),
        transform: Transform::from_xyz(0.0, 0.5, 0.0),
        ..default()
    });
    // light
    commands.spawn(PointLightBundle {
        point_light: PointLight {
            shadows_enabled: true,
            ..default()
        },
        transform: Transform::from_xyz(4.0, 8.0, 4.0),
        ..default()
    });

    // camera
    commands.spawn(Camera3dBundle {
        transform: Transform::from_xyz(-2.5, 4.5, 9.0).looking_at(Vec3::ZERO, Vec3::Y),
        ..default()
    });
}

/// A marker component for colored meshes
#[derive(Component, Default, ExtractComponent, Clone, Copy)]
pub struct ColoredMesh3d;

/// Custom pipeline for meshes with vertex colors
#[derive(Resource)]
pub struct ColoredMesh3dPipeline {
    /// this pipeline wraps the standard [`MeshPipeline`]
    mesh_pipeline: MeshPipeline,
}

impl FromWorld for ColoredMesh3dPipeline {
    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 ColoredMesh3dPipeline {
    type Key = MeshPipelineKey;

    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.contains(MeshPipelineKey::HDR) {
            true => ViewTarget::TEXTURE_FORMAT_HDR,
            false => TextureFormat::bevy_default(),
        };

        RenderPipelineDescriptor {
            vertex: VertexState {
                // Use our custom shader
                shader: COLORED_MESH_SHADER_HANDLE,
                entry_point: "vertex".into(),
                shader_defs: vec![],
                // Use our custom vertex buffer
                buffers: vec![vertex_layout],
            },
            fragment: Some(FragmentState {
                // Use our custom shader
                shader: COLORED_MESH_SHADER_HANDLE,
                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))
                    .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.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.msaa_samples(),
                mask: !0,
                alpha_to_coverage_enabled: false,
            },
            label: Some("colored_mesh_pipeline".into()),
        }
    }
}

// This specifies how to render a colored 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,
);

// The custom shader can be inline like here, included from another file at build time
// using `include_str!()`, or loaded like any other asset with `asset_server.load()`.
const COLORED_MESH_SHADER: &str = r"
#import bevy_pbr::{
    mesh_functions,
    view_transformations::position_world_to_clip
}

struct Vertex {
    @builtin(instance_index) instance_index: u32,
    @location(0) position: vec3<f32>,
};
struct VertexOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) world_position: vec4<f32>,
};
@vertex
fn vertex(vertex: Vertex) -> VertexOutput {
    var out: VertexOutput;
    var world_from_local = mesh_functions::get_world_from_local(vertex.instance_index);
    out.world_position = mesh_functions::mesh_position_local_to_world(world_from_local, vec4(vertex.position, 1.0));
    out.position = position_world_to_clip(out.world_position.xyz);
    return out;
}

@fragment
fn fragment() -> @location(0) vec4<f32> {
    return vec4(1.0, 0.0, 0.0, 1.0);
}
";

/// Plugin that renders [`ColoredMesh3d`]s
pub struct ColoredMeshPlugin;

/// Handle to the custom shader with a unique random ID
pub const COLORED_MESH_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(13828845428412094821);

impl Plugin for ColoredMeshPlugin {
    fn build(&self, app: &mut App) {
        // Load our custom shader
        let mut shaders = app.world_mut().resource_mut::<Assets<Shader>>();
        shaders.insert(
            &COLORED_MESH_SHADER_HANDLE,
            Shader::from_wgsl(COLORED_MESH_SHADER, file!()),
        );

        app.add_plugins(ExtractComponentPlugin::<ColoredMesh3d>::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<ColoredMesh3dPipeline>>()
            .add_systems(Render, queue_colored_mesh.in_set(RenderSet::QueueMeshes));
    }

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

/// Queue the meshes marked with [`ColoredMesh3d`] using our custom pipeline and draw function
#[allow(clippy::too_many_arguments)]
pub fn queue_colored_mesh(
    transparent_draw_functions: Res<DrawFunctions<Transparent3d>>,
    colored_mesh_pipeline: Res<ColoredMesh3dPipeline>,
    mut colored_mesh_pipelines: ResMut<SpecializedRenderPipelines<ColoredMesh3dPipeline>>,
    pipeline_cache: Res<PipelineCache>,
    msaa: Res<Msaa>,
    render_meshes: Res<RenderAssets<RenderMesh>>,
    render_mesh_instances: Res<RenderMeshInstances>,
    mut transparent_render_phases: ResMut<ViewSortedRenderPhases<Transparent3d>>,
    mut views: Query<(Entity, &VisibleEntities, &ExtractedView)>,
    has_colored_mesh_marker: Query<(), With<ColoredMesh3d>>,
) {
    // Iterate each view (a camera is a view)
    for (view_entity, visible_entities, view) in &mut views {
        let Some(transparent_phase) = transparent_render_phases.get_mut(&view_entity) else {
            continue;
        };

        let draw_colored_mesh = transparent_draw_functions.read().id::<DrawColoredMesh>();

        let mesh_key = MeshPipelineKey::from_msaa_samples(msaa.samples())
            | MeshPipelineKey::from_hdr(view.hdr);

        let rangefinder = view.rangefinder3d();
        // Queue all entities visible to that view
        for visible_entity in visible_entities.iter::<WithMesh>() {
            // TODO figure out why we can't check this on VisibleEntities
            if has_colored_mesh_marker.get(*visible_entity).is_err() {
                continue;
            }

            let Some(mesh_instance) = render_mesh_instances.render_mesh_queue_data(*visible_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 = colored_mesh_pipelines.specialize(
                &pipeline_cache,
                &colored_mesh_pipeline,
                mesh_key,
            );

            transparent_phase.add(Transparent3d {
                entity: *visible_entity,
                draw_function: draw_colored_mesh,
                pipeline: pipeline_id,
                // This material is not batched
                batch_range: 0..1,
                extra_index: PhaseItemExtraIndex::NONE,
                distance: rangefinder.distance_translation(&mesh_instance.translation),
            });
        }
    }
}