use Material abstraction for wireframe and configurable colors

IceSentry · UberLambda · IceSentry · commit f93e37eef86e · 2022-07-13T00:55:03.000-04:00
Co-authored-by: UberLambda &lt;paolo.jovon@gmail.com&gt;
diff --git a/crates/bevy_pbr/src/render/wireframe.wgsl b/crates/bevy_pbr/src/render/wireframe.wgsl
@@ -1,44 +1,11 @@
-#import bevy_pbr::mesh_types
-#import bevy_pbr::mesh_view_bindings
-
-[[group(1), binding(0)]]
-var<uniform> mesh: Mesh;
-
-#ifdef SKINNED
-[[group(1), binding(1)]]
-var<uniform> joint_matrices: SkinnedMesh;
-#import bevy_pbr::skinning
-#endif
-
-// NOTE: Bindings must come before functions that use them!
-#import bevy_pbr::mesh_functions
-
-struct Vertex {
-    [[location(0)]] position: vec3<f32>;
-#ifdef SKINNED
-    [[location(4)]] joint_indexes: vec4<u32>;
-    [[location(5)]] joint_weights: vec4<f32>;
-#endif
+struct WireframeMaterial {
+    color: vec4<f32>;
 };
 
-struct VertexOutput {
-    [[builtin(position)]] clip_position: vec4<f32>;
-};
-
-[[stage(vertex)]]
-fn vertex(vertex: Vertex) -> VertexOutput {
-#ifdef SKINNED
-    let model = skin_model(vertex.joint_indexes, vertex.joint_weights);
-#else
-    let model = mesh.model;
-#endif
-
-    var out: VertexOutput;
-    out.clip_position = mesh_position_local_to_clip(model, vec4<f32>(vertex.position, 1.0));
-    return out;
-}
+[[group(1), binding(0)]]
+var<uniform> material: WireframeMaterial;
 
 [[stage(fragment)]]
 fn fragment() -> [[location(0)]] vec4<f32> {
-    return vec4<f32>(1.0, 1.0, 1.0, 1.0);
+    return material.color;
 }
diff --git a/crates/bevy_pbr/src/wireframe.rs b/crates/bevy_pbr/src/wireframe.rs
@@ -1,174 +1,195 @@
-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_asset::{load_internal_asset, Assets, Handle, HandleUntyped};
+use bevy_ecs::prelude::*;
 use bevy_reflect::{Reflect, TypeUuid};
 use bevy_render::{
     extract_resource::{ExtractResource, ExtractResourcePlugin},
     mesh::{Mesh, MeshVertexBufferLayout},
-    render_asset::RenderAssets,
-    render_phase::{AddRenderCommand, DrawFunctions, RenderPhase, SetItemPipeline},
+    prelude::{Color, 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 =
+pub const WIREFRAME_MATERIAL_SHADER_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 192598014480025766);
 
-#[derive(Debug, Default)]
-pub struct WireframePlugin;
+/// Configuration resource for [`WireframePlugin`].
+#[derive(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.
+    pub global: bool,
+    /// If [`Self::global`] is set, any [`Entity`] that does not have a [`Wireframe`] component attached to it will have
+    /// wireframes in this color. Otherwise, this will be the fallback color for any entity that has a [`Wireframe`],
+    /// but no [`WireframeColor`].
+    pub color: Color,
+}
 
+/// Toggles wireframe rendering for any entity it is attached to.
+///
+/// This requires the [`WireframePlugin`] to be enabled.
+#[derive(Component)]
+pub struct Wireframe;
+
+/// Sets the color of the [`Wireframe`] of the entity it is attached to.
+///
+/// This overrides the [`WireframeConfig::default_color`].
+#[derive(Component)]
+pub struct WireframeColor {
+    pub color: Color,
+}
+
+/// A [`Plugin`] that draws wireframes.
+pub struct WireframePlugin;
 impl Plugin for WireframePlugin {
     fn build(&self, app: &mut bevy_app::App) {
         load_internal_asset!(
             app,
-            WIREFRAME_SHADER_HANDLE,
+            WIREFRAME_MATERIAL_SHADER_HANDLE,
             "render/wireframe.wgsl",
             Shader::from_wgsl
         );
 
+        app.add_plugin(MaterialPlugin::<WireframeMaterial>::default());
+
         app.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);
-        }
+        app.add_system(apply_global)
+            .add_system(apply_material)
+            .add_system(wireframe_color_changed)
+            .add_system(global_color_changed);
     }
 }
 
-fn extract_wireframes(mut commands: Commands, query: Query<Entity, With<Wireframe>>) {
-    for entity in &query {
-        commands.get_or_spawn(entity).insert(Wireframe);
+/// Apply the wireframe material to any mesh with a `Wireframe` component.
+/// Uses `WireframeConfig::color` as a fallback if no `WireframeColor` component is found
+#[allow(clippy::type_complexity)]
+fn apply_material(
+    mut commands: Commands,
+    config: Res<WireframeConfig>,
+    mut materials: ResMut<Assets<WireframeMaterial>>,
+    wireframes: Query<
+        (Entity, Option<&WireframeColor>),
+        (With<Wireframe>, Without<Handle<WireframeMaterial>>),
+    >,
+) {
+    for (e, color) in &wireframes {
+        commands.entity(e).insert(materials.add(WireframeMaterial {
+            color: if let Some(wireframe_color) = color {
+                wireframe_color.color
+            } else {
+                config.color
+            },
+        }));
     }
 }
 
-/// Controls whether an entity should rendered in wireframe-mode if the [`WireframePlugin`] is enabled
-#[derive(Component, Debug, Clone, Default, Reflect)]
-#[reflect(Component, Default)]
-pub struct Wireframe;
-
-#[derive(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.
-    pub global: bool,
-}
-
-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(),
-        }
+/// Updates the wireframe material when the color in `WireframeColor` changes
+#[allow(clippy::type_complexity)]
+fn wireframe_color_changed(
+    mut materials: ResMut<Assets<WireframeMaterial>>,
+    mut colors_changed: Query<
+        (&mut Handle<WireframeMaterial>, &WireframeColor),
+        (With<Wireframe>, Changed<WireframeColor>),
+    >,
+) {
+    for (mut handle, wireframe_color) in &mut colors_changed {
+        *handle = materials.add(WireframeMaterial {
+            color: wireframe_color.color,
+        });
     }
 }
 
-impl SpecializedMeshPipeline for WireframePipeline {
-    type Key = MeshPipelineKey;
+/// Updates the wireframe material of all entities without a specified color or without a `Wireframe` component
+fn global_color_changed(
+    config: Res<WireframeConfig>,
+    mut materials: ResMut<Assets<WireframeMaterial>>,
+    mut wireframes: Query<&mut Handle<WireframeMaterial>, Without<WireframeColor>>,
+) {
+    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)
+    for mut handle in &mut wireframes {
+        *handle = materials.add(WireframeMaterial {
+            color: config.color,
+        });
     }
 }
 
-#[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>>,
+/// Applies or remove 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>>,
+    mut q1: ParamSet<(
+        Query<
+            Entity,
+            (
+                With<Handle<Mesh>>,
+                Without<Handle<WireframeMaterial>>,
+                Without<Wireframe>,
+            ),
+        >,
+        Query<
+            Entity,
+            (
+                With<Handle<Mesh>>,
+                With<Handle<WireframeMaterial>>,
+                Without<Wireframe>,
+            ),
+        >,
     )>,
-    mut views: Query<(&ExtractedView, &VisibleEntities, &mut RenderPhase<Opaque3d>)>,
+    mut is_global_applied: Local<bool>,
 ) {
-    let draw_custom = opaque_3d_draw_functions
-        .read()
-        .get_id::<DrawWireframes>()
-        .unwrap();
-    let msaa_key = MeshPipelineKey::from_msaa_samples(msaa.samples);
-    for (view, visible_entities, mut opaque_phase) in &mut views {
-        let rangefinder = view.rangefinder3d();
-
-        let add_render_phase =
-            |(entity, mesh_handle, mesh_uniform): (Entity, &Handle<Mesh>, &MeshUniform)| {
-                if let Some(mesh) = render_meshes.get(mesh_handle) {
-                    let key = msaa_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),
-                    });
-                }
-            };
-
-        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);
+    if !config.is_changed() {
+        return;
+    }
+
+    if !*is_global_applied && config.global {
+        let global_material = materials.add(WireframeMaterial {
+            color: config.color,
+        });
+
+        for e in &mut q1.p0() {
+            commands.entity(e).insert(global_material.clone());
         }
+
+        *is_global_applied = true;
+    } else if *is_global_applied && !config.global {
+        for e in &mut q1.p1() {
+            commands.entity(e).remove::<Handle<WireframeMaterial>>();
+        }
+        *is_global_applied = false;
     }
 }
 
-type DrawWireframes = (
-    SetItemPipeline,
-    SetMeshViewBindGroup<0>,
-    SetMeshBindGroup<1>,
-    DrawMesh,
-);
+#[derive(Default, AsBindGroup, TypeUuid, Debug, Clone)]
+#[uuid = "9e694f70-9963-4418-8bc1-3474c66b13b8"]
+struct WireframeMaterial {
+    #[uniform(0)]
+    color: Color,
+}
+
+impl Material for WireframeMaterial {
+    fn fragment_shader() -> ShaderRef {
+        WIREFRAME_MATERIAL_SHADER_HANDLE.typed().into()
+    }
+
+    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(())
+    }
+}
diff --git a/examples/3d/wireframe.rs b/examples/3d/wireframe.rs
