use Material for wireframes

crates/bevy_pbr/src/render/wireframe.wgsl

@@ -1,66 +1,6 @@
-#import bevy_pbr::mesh_bindings    mesh
-#import bevy_pbr::mesh_functions   get_model_matrix, mesh_position_local_to_clip
-#import bevy_pbr::morph
-
-#ifdef SKINNED
-    #import bevy_pbr::skinning
-#endif
-
-struct Vertex {
-    @builtin(instance_index) instance_index: u32,
-    @location(0) position: vec3<f32>,
-#ifdef SKINNED
-    @location(5) joint_indexes: vec4<u32>,
-    @location(6) joint_weights: vec4<f32>,
-#endif
-#ifdef MORPH_TARGETS
-    @builtin(vertex_index) index: u32,
-#endif
-};
-
-struct VertexOutput {
-    @builtin(position) clip_position: vec4<f32>,
-};
-
-
-#ifdef MORPH_TARGETS
-fn morph_vertex(vertex_in: Vertex) -> Vertex {
-    var vertex = vertex_in;
-    let weight_count = bevy_pbr::morph::layer_count();
-    for (var i: u32 = 0u; i < weight_count; i ++) {
-        let weight = bevy_pbr::morph::weight_at(i);
-        if weight == 0.0 {
-            continue;
-        }
-        vertex.position += weight * bevy_pbr::morph::morph(vertex.index, bevy_pbr::morph::position_offset, i);
-    }
-    return vertex;
-}
-#endif
-
-@vertex
-fn vertex(vertex_no_morph: Vertex) -> VertexOutput {
-
-#ifdef MORPH_TARGETS
-    var vertex = morph_vertex(vertex_no_morph);
-#else
-    var vertex = vertex_no_morph;
-#endif
-
-#ifdef SKINNED
-    let model = bevy_pbr::skinning::skin_model(vertex.joint_indexes, vertex.joint_weights);
-#else
-    // Use vertex_no_morph.instance_index instead of vertex.instance_index to work around a wgpu dx12 bug.
-    // See https://github.com/gfx-rs/naga/issues/2416 .
-    let model = get_model_matrix(vertex_no_morph.instance_index);
-#endif
-
-    var out: VertexOutput;
-    out.clip_position = mesh_position_local_to_clip(model, vec4<f32>(vertex.position, 1.0));
-    return out;
-}
+#import bevy_pbr::mesh_vertex_output MeshVertexOutput
 
 @fragment
-fn fragment() -> @location(0) vec4<f32> {
+fn fragment(in: MeshVertexOutput) -> @location(0) vec4<f32> {
     return vec4<f32>(1.0, 1.0, 1.0, 1.0);
-}
+}

crates/bevy_pbr/src/wireframe.rs

@@ -1,30 +1,28 @@
-use crate::{
-    DrawMesh, MeshPipeline, MeshPipelineKey, RenderMeshInstance, RenderMeshInstances,
-    SetMeshBindGroup, SetMeshViewBindGroup,
-};
-use bevy_app::Plugin;
-use bevy_asset::{load_internal_asset, Handle};
-use bevy_core_pipeline::core_3d::Opaque3d;
-use bevy_derive::{Deref, DerefMut};
-use bevy_ecs::{prelude::*, reflect::ReflectComponent};
-use bevy_reflect::{std_traits::ReflectDefault, Reflect};
+use crate::{Material, MaterialPipeline, MaterialPipelineKey, MaterialPlugin};
+use bevy_app::{Plugin, Startup, Update};
+use bevy_asset::{load_internal_asset, Asset, Assets, Handle};
+use bevy_ecs::prelude::*;
+use bevy_reflect::{std_traits::ReflectDefault, Reflect, TypePath, 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, RenderSet,
 };
-use bevy_render::{Extract, ExtractSchedule, Render};
-use bevy_utils::{tracing::error, EntityHashSet};
 
 pub const WIREFRAME_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(192598014480025766);
 
+/// A [`Plugin`] that draws wireframes.
+///
+/// Wireframes currently do not work when using webgl or webgpu.
+/// Supported platforms:
+/// - DX12
+/// - Vulkan
+/// - Metal
+///
+/// This is a native only feature.
 #[derive(Debug, Default)]
 pub struct WireframePlugin;
 
@@ -41,23 +39,12 @@ impl Plugin for WireframePlugin {
             .register_type::<NoWireframe>()
             .register_type::<WireframeConfig>()
             .init_resource::<WireframeConfig>()
-            .add_plugins((ExtractResourcePlugin::<WireframeConfig>::default(),));
-
-        if let Ok(render_app) = app.get_sub_app_mut(RenderApp) {
-            render_app
-                .add_render_command::<Opaque3d, DrawWireframes>()
-                .init_resource::<SpecializedMeshPipelines<WireframePipeline>>()
-                .init_resource::<Wireframes>()
-                .init_resource::<NoWireframes>()
-                .add_systems(ExtractSchedule, extract_wireframes)
-                .add_systems(Render, queue_wireframes.in_set(RenderSet::QueueMeshes));
-        }
-    }
-
-    fn finish(&self, app: &mut bevy_app::App) {
-        if let Ok(render_app) = app.get_sub_app_mut(RenderApp) {
-            render_app.init_resource::<WireframePipeline>();
-        }
+            .add_plugins(MaterialPlugin::<WireframeMaterial>::default())
+            .add_systems(Startup, setup_global_wireframe_material)
+            .add_systems(
+                Update,
+                (apply_global_wireframe_material, apply_wireframe_material),
+            );
     }
 }
 
@@ -85,133 +72,103 @@ pub struct WireframeConfig {
     pub global: bool,
 }
 
-#[derive(Resource, Default, Deref, DerefMut)]
-pub struct Wireframes(EntityHashSet<Entity>);
-
-#[derive(Resource, Default, Deref, DerefMut)]
-pub struct NoWireframes(EntityHashSet<Entity>);
+#[derive(Resource)]
+struct GlobalWireframeMaterial {
+    // This handle will be reused when the global config is enabled
+    handle: Handle<WireframeMaterial>,
+}
 
-fn extract_wireframes(
-    mut wireframes: ResMut<Wireframes>,
-    mut no_wireframes: ResMut<NoWireframes>,
-    wireframe_query: Extract<Query<Entity, With<Wireframe>>>,
-    no_wireframe_query: Extract<Query<Entity, With<NoWireframe>>>,
+fn setup_global_wireframe_material(
+    mut commands: Commands,
+    mut materials: ResMut<Assets<WireframeMaterial>>,
 ) {
-    wireframes.clear();
-    wireframes.extend(wireframe_query.iter());
-    no_wireframes.clear();
-    no_wireframes.extend(no_wireframe_query.iter());
+    // Create the handle used for the global material
+    commands.insert_resource(GlobalWireframeMaterial {
+        handle: materials.add(WireframeMaterial {}),
+    });
 }
 
-#[derive(Resource, Clone)]
-pub struct WireframePipeline {
-    mesh_pipeline: MeshPipeline,
-    shader: Handle<Shader>,
+/// Applies or remove the wireframe material to any mesh with a [`Wireframe`] component.
+fn apply_wireframe_material(
+    mut commands: Commands,
+    mut materials: ResMut<Assets<WireframeMaterial>>,
+    wireframes: Query<Entity, (With<Wireframe>, Without<Handle<WireframeMaterial>>)>,
+    mut removed_wireframes: RemovedComponents<Wireframe>,
+) {
+    for e in removed_wireframes.read() {
+        if let Some(mut commands) = commands.get_entity(e) {
+            commands.remove::<Handle<WireframeMaterial>>();
+        }
+    }
+
+    let mut wireframes_to_spawn = vec![];
+    for e in &wireframes {
+        wireframes_to_spawn.push((e, materials.add(WireframeMaterial {})));
+    }
+    commands.insert_or_spawn_batch(wireframes_to_spawn);
 }
-impl FromWorld for WireframePipeline {
-    fn from_world(render_world: &mut World) -> Self {
-        WireframePipeline {
-            mesh_pipeline: render_world.resource::<MeshPipeline>().clone(),
-            shader: WIREFRAME_SHADER_HANDLE,
+
+/// Applies or removes a wireframe material on any mesh without a [`Wireframe`] component.
+#[allow(clippy::type_complexity)]
+fn apply_global_wireframe_material(
+    mut commands: Commands,
+    config: Res<WireframeConfig>,
+    meshes_without_material: Query<
+        Entity,
+        (
+            With<Handle<Mesh>>,
+            Without<Wireframe>,
+            Without<NoWireframe>,
+            Without<Handle<WireframeMaterial>>,
+        ),
+    >,
+    meshes_with_global_material: Query<
+        Entity,
+        (
+            With<Handle<Mesh>>,
+            Without<Wireframe>,
+            Without<NoWireframe>,
+            With<Handle<WireframeMaterial>>,
+        ),
+    >,
+    global_material: Res<GlobalWireframeMaterial>,
+) {
+    if !config.is_changed() {
+        return;
+    }
+
+    if config.global {
+        let mut material_to_spawn = vec![];
+        for e in &meshes_without_material {
+            // We only add the material handle but not the Wireframe component
+            // This makes it easy to detect which mesh is using the global material and which ones are user specified
+            material_to_spawn.push((e, global_material.handle.clone()));
+        }
+        commands.insert_or_spawn_batch(material_to_spawn);
+    } else if !config.global {
+        for e in &meshes_with_global_material {
+            commands.entity(e).remove::<Handle<WireframeMaterial>>();
         }
     }
 }
 
-impl SpecializedMeshPipeline for WireframePipeline {
-    type Key = MeshPipelineKey;
+#[derive(Default, AsBindGroup, TypeUuid, TypePath, Debug, Clone, Asset)]
+#[uuid = "9e694f70-9963-4418-8bc1-3474c66b13b8"]
+struct WireframeMaterial {}
+
+impl Material for WireframeMaterial {
+    fn fragment_shader() -> ShaderRef {
+        WIREFRAME_SHADER_HANDLE.into()
+    }
 
     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
-            .vertex
-            .shader_defs
-            .push("MESH_BINDGROUP_1".into());
-        descriptor.fragment.as_mut().unwrap().shader = self.shader.clone_weak();
+        _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(descriptor)
-    }
-}
-
-#[allow(clippy::too_many_arguments)]
-fn queue_wireframes(
-    opaque_3d_draw_functions: Res<DrawFunctions<Opaque3d>>,
-    render_meshes: Res<RenderAssets<Mesh>>,
-    render_mesh_instances: Res<RenderMeshInstances>,
-    wireframes: Res<Wireframes>,
-    no_wireframes: Res<NoWireframes>,
-    wireframe_config: Res<WireframeConfig>,
-    wireframe_pipeline: Res<WireframePipeline>,
-    mut pipelines: ResMut<SpecializedMeshPipelines<WireframePipeline>>,
-    pipeline_cache: Res<PipelineCache>,
-    msaa: Res<Msaa>,
-    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();
-
-        let view_key = msaa_key | MeshPipelineKey::from_hdr(view.hdr);
-        let add_render_phase = |phase_item: (Entity, &RenderMeshInstance)| {
-            let (entity, mesh_instance) = phase_item;
-
-            let Some(mesh) = render_meshes.get(mesh_instance.mesh_asset_id) else {
-                return;
-            };
-            let mut key =
-                view_key | MeshPipelineKey::from_primitive_topology(mesh.primitive_topology);
-            if mesh.morph_targets.is_some() {
-                key |= MeshPipelineKey::MORPH_TARGETS;
-            }
-            let pipeline_id =
-                pipelines.specialize(&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_translation(&mesh_instance.transforms.transform.translation),
-                batch_range: 0..1,
-                dynamic_offset: None,
-            });
-        };
-
-        visible_entities
-            .entities
-            .iter()
-            .filter_map(|visible_entity| {
-                if no_wireframes.get(visible_entity).is_some() {
-                    return None;
-                }
-
-                if wireframe_config.global || wireframes.get(visible_entity).is_some() {
-                    render_mesh_instances
-                        .get(visible_entity)
-                        .map(|mesh_instance| (*visible_entity, mesh_instance))
-                } else {
-                    None
-                }
-            })
-            .for_each(add_render_phase);
+        Ok(())
     }
 }
-
-type DrawWireframes = (
-    SetItemPipeline,
-    SetMeshViewBindGroup<0>,
-    SetMeshBindGroup<1>,
-    DrawMesh,
-);