use Material for wireframes

Author: IceSentry

crates/bevy_pbr/src/render/wireframe.wgsl

@@ -1,11 +1,11 @@
 #import bevy_pbr::mesh_types
 #import bevy_pbr::mesh_view_bindings
 
-@group(1) @binding(0)
+@group(2) @binding(0)
 var<uniform> mesh: Mesh;
 
 #ifdef SKINNED
-@group(1) @binding(1)
+@group(2) @binding(1)
 var<uniform> joint_matrices: SkinnedMesh;
 #import bevy_pbr::skinning
 #endif
@@ -41,4 +41,4 @@ fn vertex(vertex: Vertex) -> VertexOutput {
 @fragment
 fn fragment() -> @location(0) vec4<f32> {
     return vec4<f32>(1.0, 1.0, 1.0, 1.0);
-}
+}

crates/bevy_pbr/src/wireframe.rs

@@ -1,29 +1,21 @@
-use crate::MeshPipeline;
-use crate::{DrawMesh, MeshPipelineKey, MeshUniform, SetMeshBindGroup, SetMeshViewBindGroup};
+use crate::{Material, MaterialPipeline, MaterialPipelineKey, MaterialPlugin};
 use bevy_app::Plugin;
-use bevy_asset::{load_internal_asset, Handle, HandleUntyped};
-use bevy_core_pipeline::core_3d::Opaque3d;
-use bevy_ecs::{prelude::*, reflect::ReflectComponent};
-use bevy_reflect::std_traits::ReflectDefault;
-use bevy_reflect::{Reflect, TypeUuid};
-use bevy_render::Extract;
+use bevy_asset::{load_internal_asset, Assets, Handle, HandleUntyped};
+use bevy_ecs::prelude::*;
+use bevy_reflect::{std_traits::ReflectDefault, Reflect, TypeUuid};
 use bevy_render::{
-    extract_resource::{ExtractResource, ExtractResourcePlugin},
+    extract_resource::ExtractResource,
     mesh::{Mesh, MeshVertexBufferLayout},
-    render_asset::RenderAssets,
-    render_phase::{AddRenderCommand, DrawFunctions, RenderPhase, SetItemPipeline},
+    prelude::Shader,
     render_resource::{
-        PipelineCache, PolygonMode, RenderPipelineDescriptor, Shader, SpecializedMeshPipeline,
-        SpecializedMeshPipelineError, SpecializedMeshPipelines,
+        AsBindGroup, PolygonMode, RenderPipelineDescriptor, ShaderRef, SpecializedMeshPipelineError,
     },
-    view::{ExtractedView, Msaa, VisibleEntities},
-    RenderApp, RenderStage,
 };
-use bevy_utils::tracing::error;
 
 pub const WIREFRAME_SHADER_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 192598014480025766);
 
+/// A [`Plugin`] that draws wireframes.
 #[derive(Debug, Default)]
 pub struct WireframePlugin;
 
@@ -39,137 +31,101 @@ impl Plugin for WireframePlugin {
         app.register_type::<Wireframe>()
             .register_type::<WireframeConfig>()
             .init_resource::<WireframeConfig>()
-            .add_plugin(ExtractResourcePlugin::<WireframeConfig>::default());
-
-        if let Ok(render_app) = app.get_sub_app_mut(RenderApp) {
-            render_app
-                .add_render_command::<Opaque3d, DrawWireframes>()
-                .init_resource::<WireframePipeline>()
-                .init_resource::<SpecializedMeshPipelines<WireframePipeline>>()
-                .add_system_to_stage(RenderStage::Extract, extract_wireframes)
-                .add_system_to_stage(RenderStage::Queue, queue_wireframes);
-        }
-    }
-}
-
-fn extract_wireframes(mut commands: Commands, query: Extract<Query<Entity, With<Wireframe>>>) {
-    for entity in &query {
-        commands.get_or_spawn(entity).insert(Wireframe);
+            .add_plugin(MaterialPlugin::<WireframeMaterial>::default())
+            .add_system(apply_global)
+            .add_system(apply_material);
     }
 }
 
-/// Controls whether an entity should rendered in wireframe-mode if the [`WireframePlugin`] is enabled
+/// Toggles wireframe rendering for any entity it is attached to.
+///
+/// This requires the [`WireframePlugin`] to be enabled.
 #[derive(Component, Debug, Clone, Default, Reflect)]
 #[reflect(Component, Default)]
 pub struct Wireframe;
 
+/// Configuration resource for [`WireframePlugin`].
 #[derive(Resource, Debug, Clone, Default, ExtractResource, Reflect)]
 #[reflect(Resource)]
 pub struct WireframeConfig {
-    /// Whether to show wireframes for all meshes. If `false`, only meshes with a [Wireframe] component will be rendered.
+    /// Whether to show wireframes for all meshes.
+    /// If `false`, only meshes with a [`Wireframe`] component will be rendered.
     pub global: bool,
 }
 
-#[derive(Resource, Clone)]
-pub struct WireframePipeline {
-    mesh_pipeline: MeshPipeline,
-    shader: Handle<Shader>,
-}
-impl FromWorld for WireframePipeline {
-    fn from_world(render_world: &mut World) -> Self {
-        WireframePipeline {
-            mesh_pipeline: render_world.resource::<MeshPipeline>().clone(),
-            shader: WIREFRAME_SHADER_HANDLE.typed(),
-        }
+/// Applies the wireframe material to any mesh with a [`Wireframe`] component.
+fn apply_material(
+    mut commands: Commands,
+    mut materials: ResMut<Assets<WireframeMaterial>>,
+    wireframes: Query<Entity, (With<Wireframe>, Without<Handle<WireframeMaterial>>)>,
+) {
+    for e in &wireframes {
+        commands
+            .entity(e)
+            .insert(materials.add(WireframeMaterial {}));
     }
 }
 
-impl SpecializedMeshPipeline for WireframePipeline {
-    type Key = MeshPipelineKey;
+/// Applies or removes a wireframe material on any mesh without a [`Wireframe`] component.
+#[allow(clippy::type_complexity)]
+fn apply_global(
+    mut commands: Commands,
+    config: Res<WireframeConfig>,
+    mut materials: ResMut<Assets<WireframeMaterial>>,
+    meshes_without_material: Query<
+        Entity,
+        (
+            With<Handle<Mesh>>,
+            Without<Wireframe>,
+            Without<Handle<WireframeMaterial>>,
+        ),
+    >,
+    meshes_with_material: Query<
+        Entity,
+        (
+            With<Handle<Mesh>>,
+            Without<Wireframe>,
+            With<Handle<WireframeMaterial>>,
+        ),
+    >,
+) {
+    if !config.is_changed() {
+        return;
+    }
 
-    fn specialize(
-        &self,
-        key: Self::Key,
-        layout: &MeshVertexBufferLayout,
-    ) -> Result<RenderPipelineDescriptor, SpecializedMeshPipelineError> {
-        let mut descriptor = self.mesh_pipeline.specialize(key, layout)?;
-        descriptor.vertex.shader = self.shader.clone_weak();
-        descriptor.fragment.as_mut().unwrap().shader = self.shader.clone_weak();
-        descriptor.primitive.polygon_mode = PolygonMode::Line;
-        descriptor.depth_stencil.as_mut().unwrap().bias.slope_scale = 1.0;
-        Ok(descriptor)
+    if config.global {
+        let global_material = materials.add(WireframeMaterial {});
+        for e in &meshes_without_material {
+            commands.entity(e).insert(global_material.clone());
+        }
+    } else if !config.global {
+        for e in &meshes_with_material {
+            commands.entity(e).remove::<Handle<WireframeMaterial>>();
+        }
     }
 }
 
-#[allow(clippy::too_many_arguments)]
-fn queue_wireframes(
-    opaque_3d_draw_functions: Res<DrawFunctions<Opaque3d>>,
-    render_meshes: Res<RenderAssets<Mesh>>,
-    wireframe_config: Res<WireframeConfig>,
-    wireframe_pipeline: Res<WireframePipeline>,
-    mut pipelines: ResMut<SpecializedMeshPipelines<WireframePipeline>>,
-    mut pipeline_cache: ResMut<PipelineCache>,
-    msaa: Res<Msaa>,
-    mut material_meshes: ParamSet<(
-        Query<(Entity, &Handle<Mesh>, &MeshUniform)>,
-        Query<(Entity, &Handle<Mesh>, &MeshUniform), With<Wireframe>>,
-    )>,
-    mut views: Query<(&ExtractedView, &VisibleEntities, &mut RenderPhase<Opaque3d>)>,
-) {
-    let draw_custom = opaque_3d_draw_functions.read().id::<DrawWireframes>();
-    let msaa_key = MeshPipelineKey::from_msaa_samples(msaa.samples);
-    for (view, visible_entities, mut opaque_phase) in &mut views {
-        let rangefinder = view.rangefinder3d();
+#[derive(Default, AsBindGroup, TypeUuid, Debug, Clone)]
+#[uuid = "9e694f70-9963-4418-8bc1-3474c66b13b8"]
+struct WireframeMaterial {}
 
-        let view_key = msaa_key | MeshPipelineKey::from_hdr(view.hdr);
-        let add_render_phase =
-            |(entity, mesh_handle, mesh_uniform): (Entity, &Handle<Mesh>, &MeshUniform)| {
-                if let Some(mesh) = render_meshes.get(mesh_handle) {
-                    let key = view_key
-                        | MeshPipelineKey::from_primitive_topology(mesh.primitive_topology);
-                    let pipeline_id = pipelines.specialize(
-                        &mut pipeline_cache,
-                        &wireframe_pipeline,
-                        key,
-                        &mesh.layout,
-                    );
-                    let pipeline_id = match pipeline_id {
-                        Ok(id) => id,
-                        Err(err) => {
-                            error!("{}", err);
-                            return;
-                        }
-                    };
-                    opaque_phase.add(Opaque3d {
-                        entity,
-                        pipeline: pipeline_id,
-                        draw_function: draw_custom,
-                        distance: rangefinder.distance(&mesh_uniform.transform),
-                    });
-                }
-            };
+impl Material for WireframeMaterial {
+    fn vertex_shader() -> ShaderRef {
+        WIREFRAME_SHADER_HANDLE.typed().into()
+    }
 
-        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);
-        }
+    fn fragment_shader() -> ShaderRef {
+        WIREFRAME_SHADER_HANDLE.typed().into()
     }
-}
 
-type DrawWireframes = (
-    SetItemPipeline,
-    SetMeshViewBindGroup<0>,
-    SetMeshBindGroup<1>,
-    DrawMesh,
-);
+    fn specialize(
+        _pipeline: &MaterialPipeline<Self>,
+        descriptor: &mut RenderPipelineDescriptor,
+        _layout: &MeshVertexBufferLayout,
+        _key: MaterialPipelineKey<Self>,
+    ) -> Result<(), SpecializedMeshPipelineError> {
+        descriptor.primitive.polygon_mode = PolygonMode::Line;
+        descriptor.depth_stencil.as_mut().unwrap().bias.slope_scale = 1.0;
+        Ok(())
+    }
+}