bevy wireframe2d

lib.rs

//! Extensions to bevy's wireframe renderer.

use bevy::{
    core_pipeline::core_2d::Transparent2d,
    ecs::system::{lifetimeless::SRes, SystemParamItem},
    pbr::wireframe::{Wireframe, WireframeConfig, WireframePlugin},
    prelude::*,
    render::{
        mesh::MeshVertexBufferLayout,
        render_asset::RenderAssets,
        render_phase::{
            AddRenderCommand, DrawFunctions, PhaseItem, RenderCommand, RenderCommandResult,
            RenderPhase, SetItemPipeline, TrackedRenderPass,
        },
        render_resource::{
            BindGroupLayout, PipelineCache, PolygonMode, RenderPipelineDescriptor,
            SpecializedMeshPipeline, SpecializedMeshPipelineError, SpecializedMeshPipelines,
        },
        view::{ExtractedView, VisibleEntities},
        Extract, Render, RenderApp, RenderSet,
    },
    sprite::{
        DrawMesh2d, Material2dPipeline, Mesh2dHandle, Mesh2dPipeline, Mesh2dPipelineKey,
        Mesh2dUniform, RenderMaterials2d, SetMesh2dBindGroup, SetMesh2dViewBindGroup,
    },
    utils::FloatOrd,
};

/// Renders 2D meshes as wireframes if they have the [`Wireframe`] component.
pub struct Wireframe2dPlugin;

impl Plugin for Wireframe2dPlugin {
    fn build(&self, app: &mut App) {
        if !app.is_plugin_added::<WireframePlugin>() {
            app.add_plugins(WireframePlugin);
        }

        let config = Wireframe2dConfig::from_world(&mut app.world);

        app.register_type::<Wireframe2dConfig>()
            .insert_resource(config.clone());

        let render_app = app.sub_app_mut(RenderApp);
        render_app
            .add_render_command::<Transparent2d, DrawWireframes2d>()
            .insert_resource(config)
            .add_systems(ExtractSchedule, extract_wireframes2d)
            .add_systems(Render, queue_wireframes2d.in_set(RenderSet::Queue));
    }

    fn finish(&self, app: &mut App) {
        let render_app = app.sub_app_mut(RenderApp);

        render_app
            .init_resource::<Wireframe2dPipeline>()
            .init_resource::<SpecializedMeshPipelines<Wireframe2dPipeline>>();
    }
}

/// A resource that configures how 2D wireframes should be rendered.
#[derive(Clone, Resource, Reflect)]
pub struct Wireframe2dConfig {
    /// The default material to render 2D wireframes in.
    pub material: Handle<ColorMaterial>,
}

impl FromWorld for Wireframe2dConfig {
    fn from_world(world: &mut World) -> Self {
        let mut mats = world.resource_mut::<Assets<ColorMaterial>>();
        Self {
            material: mats.add(ColorMaterial {
                color: Color::GREEN,
                ..Default::default()
            }),
        }
    }
}

#[derive(Resource)]
struct Wireframe2dPipeline {
    mesh2d_pipeline: Mesh2dPipeline,
    material2d_layout: BindGroupLayout,
    vertex_shader: Option<Handle<Shader>>,
    fragment_shader: Option<Handle<Shader>>,
}

impl SpecializedMeshPipeline for Wireframe2dPipeline {
    type Key = Mesh2dPipelineKey;

    fn specialize(
        &self,
        key: Self::Key,
        layout: &MeshVertexBufferLayout,
    ) -> Result<RenderPipelineDescriptor, SpecializedMeshPipelineError> {
        let mut descriptor = self.mesh2d_pipeline.specialize(key, layout)?;
        if let Some(vertex_shader) = &self.vertex_shader {
            descriptor.vertex.shader = vertex_shader.clone();
        }

        if let Some(fragment_shader) = &self.fragment_shader {
            descriptor.fragment.as_mut().unwrap().shader = fragment_shader.clone();
        }
        descriptor.label = Some("wireframe2d_pipeline".into());
        descriptor.layout = vec![
            self.mesh2d_pipeline.view_layout.clone(),
            self.material2d_layout.clone(),
            self.mesh2d_pipeline.mesh_layout.clone(),
        ];
        descriptor.primitive.polygon_mode = PolygonMode::Line;
        // FIXME: ideally we'd set a line width of 2 or so, since MSAA makes them look very thin

        Ok(descriptor)
    }
}

impl FromWorld for Wireframe2dPipeline {
    fn from_world(world: &mut World) -> Self {
        let p = &world.resource::<Material2dPipeline<ColorMaterial>>();
        Self {
            mesh2d_pipeline: p.mesh2d_pipeline.clone(),
            material2d_layout: p.material2d_layout.clone(),
            vertex_shader: p.vertex_shader.clone(),
            fragment_shader: p.fragment_shader.clone(),
        }
    }
}

fn extract_wireframes2d(
    config_in: Extract<Res<Wireframe2dConfig>>,
    mut config_out: ResMut<Wireframe2dConfig>,
) {
    // We rely on the wireframe plugin to do most of the work.
    config_out.clone_from(&config_in);
}

fn queue_wireframes2d(
    transparent_2d_draw_functions: Res<DrawFunctions<Transparent2d>>,
    render_meshes: Res<RenderAssets<Mesh>>,
    wireframe_config: Res<WireframeConfig>,
    wireframe_pipeline: Res<Wireframe2dPipeline>,
    mut pipelines: ResMut<SpecializedMeshPipelines<Wireframe2dPipeline>>,
    mut pipeline_cache: ResMut<PipelineCache>,
    msaa: Res<Msaa>,
    mut material_meshes: ParamSet<(
        Query<(Entity, &Mesh2dHandle, &Mesh2dUniform)>,
        Query<(Entity, &Mesh2dHandle, &Mesh2dUniform), With<Wireframe>>,
    )>,
    mut views: Query<(
        &ExtractedView,
        &VisibleEntities,
        &mut RenderPhase<Transparent2d>,
    )>,
) {
    let draw_custom = transparent_2d_draw_functions
        .read()
        .get_id::<DrawWireframes2d>()
        .unwrap();
    let msaa_key = Mesh2dPipelineKey::from_msaa_samples(msaa.samples());
    for (view, visible_entities, mut opaque_phase) in &mut views {
        let view_key = msaa_key | Mesh2dPipelineKey::from_hdr(view.hdr);
        let add_render_phase =
            |(entity, mesh2d_handle, mesh2d_uniform): (Entity, &Mesh2dHandle, &Mesh2dUniform)| {
                if let Some(mesh) = render_meshes.get(&mesh2d_handle.0) {
                    let mesh_key = view_key
                        | Mesh2dPipelineKey::from_primitive_topology(mesh.primitive_topology);
                    let pipeline_id = pipelines.specialize(
                        &mut pipeline_cache,
                        &wireframe_pipeline,
                        mesh_key,
                        &mesh.layout,
                    );
                    let pipeline_id = match pipeline_id {
                        Ok(id) => id,
                        Err(err) => {
                            error!("{}", err);
                            return;
                        }
                    };

                    let mesh_z = mesh2d_uniform.transform.w_axis.z;
                    opaque_phase.add(Transparent2d {
                        // Ensure to draw the wireframe closer to the camera than the base mesh.
                        sort_key: FloatOrd(mesh_z + 1.0),
                        entity,
                        pipeline: pipeline_id,
                        draw_function: draw_custom,
                        batch_range: None,
                    });
                }
            };

        if wireframe_config.global {
            let query = material_meshes.p0();
            visible_entities
                .entities
                .iter()
                .filter_map(|visible_entity| query.get(*visible_entity).ok())
                .for_each(add_render_phase);
        } else {
            let query = material_meshes.p1();
            visible_entities
                .entities
                .iter()
                .filter_map(|visible_entity| query.get(*visible_entity).ok())
                .for_each(add_render_phase);
        }
    }
}

struct SetWireframe2dBindGroup<const I: usize>;
impl<P: PhaseItem, const I: usize> RenderCommand<P> for SetWireframe2dBindGroup<I> {
    type Param = (
        SRes<RenderMaterials2d<ColorMaterial>>,
        SRes<Wireframe2dConfig>,
    );
    type ItemWorldQuery = ();
    type ViewWorldQuery = ();

    fn render<'w>(
        _item: &P,
        _view: (),
        _entity: (),
        (materials, config): SystemParamItem<'w, '_, Self::Param>,
        pass: &mut TrackedRenderPass<'w>,
    ) -> RenderCommandResult {
        // TODO: add a component for overriding the material
        let material2d = materials.into_inner().get(&config.material).unwrap();
        pass.set_bind_group(I, &material2d.bind_group, &[]);
        RenderCommandResult::Success
    }
}

type DrawWireframes2d = (
    SetItemPipeline,
    SetMesh2dViewBindGroup<0>,
    SetWireframe2dBindGroup<1>,
    SetMesh2dBindGroup<2>,
    DrawMesh2d,
);