Add colored wireframes

UberLambda · UberLambda · commit e096e8e0f8c5 · 2022-05-02T12:26:05.000+01:00
diff --git a/crates/bevy_pbr/src/render/wireframe.wgsl b/crates/bevy_pbr/src/render/wireframe.wgsl
@@ -9,8 +9,14 @@ struct Vertex {
 #endif
 };
 
+struct Wireframe {
+    color: vec4<f32>;
+};
+
 [[group(1), binding(0)]]
 var<uniform> mesh: Mesh;
+[[group(2), binding(0)]]
+var<uniform> wireframe: Wireframe;
 
 struct VertexOutput {
     [[builtin(position)]] clip_position: vec4<f32>;
@@ -39,5 +45,5 @@ fn vertex(vertex: Vertex) -> VertexOutput {
 
 [[stage(fragment)]]
 fn fragment() -> [[location(0)]] vec4<f32> {
-    return vec4<f32>(1.0, 1.0, 1.0, 1.0);
+    return wireframe.color;
 }
diff --git a/crates/bevy_pbr/src/wireframe.rs b/crates/bevy_pbr/src/wireframe.rs
@@ -1,18 +1,34 @@
-use crate::MeshPipeline;
-use crate::{DrawMesh, MeshPipelineKey, MeshUniform, SetMeshBindGroup, SetMeshViewBindGroup};
+use crate::{
+    DrawMesh, MeshPipeline, MeshPipelineKey, MeshUniform, SetMeshBindGroup, SetMeshViewBindGroup,
+};
 use bevy_app::Plugin;
 use bevy_asset::{load_internal_asset, Handle, HandleUntyped};
 use bevy_core_pipeline::Opaque3d;
-use bevy_ecs::{prelude::*, reflect::ReflectComponent};
+use bevy_ecs::{
+    prelude::*,
+    reflect::ReflectComponent,
+    system::{lifetimeless::*, SystemParamItem},
+};
+use bevy_math::Vec4;
 use bevy_reflect::{Reflect, TypeUuid};
 use bevy_render::{
+    color::Color,
     mesh::{Mesh, MeshVertexBufferLayout},
     render_asset::RenderAssets,
-    render_phase::{AddRenderCommand, DrawFunctions, RenderPhase, SetItemPipeline},
+    render_phase::{
+        AddRenderCommand, DrawFunctions, EntityRenderCommand, RenderCommandResult, RenderPhase,
+        SetItemPipeline, TrackedRenderPass,
+    },
     render_resource::{
-        PipelineCache, PolygonMode, RenderPipelineDescriptor, Shader, SpecializedMeshPipeline,
-        SpecializedMeshPipelineError, SpecializedMeshPipelines,
+        std140::AsStd140, BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayout,
+        BindGroupLayoutDescriptor, BindGroupLayoutEntry, BindingType, BufferBindingType,
+        BufferSize, DynamicUniformVec,
     },
+    render_resource::{
+        PipelineCache, PolygonMode, RenderPipelineDescriptor, Shader, ShaderStages,
+        SpecializedMeshPipeline, SpecializedMeshPipelineError, SpecializedMeshPipelines,
+    },
+    renderer::{RenderDevice, RenderQueue},
     view::{ExtractedView, Msaa},
     RenderApp, RenderStage,
 };
@@ -38,10 +54,13 @@ impl Plugin for WireframePlugin {
         if let Ok(render_app) = app.get_sub_app_mut(RenderApp) {
             render_app
                 .add_render_command::<Opaque3d, DrawWireframes>()
+                .init_resource::<GlobalWireframeMeta>()
                 .init_resource::<WireframePipeline>()
                 .init_resource::<SpecializedMeshPipelines<WireframePipeline>>()
                 .add_system_to_stage(RenderStage::Extract, extract_wireframes)
                 .add_system_to_stage(RenderStage::Extract, extract_wireframe_config)
+                .add_system_to_stage(RenderStage::Prepare, prepare_wireframes)
+                .add_system_to_stage(RenderStage::Queue, queue_wireframes_bind_group)
                 .add_system_to_stage(RenderStage::Queue, queue_wireframes);
         }
     }
@@ -53,31 +72,179 @@ fn extract_wireframe_config(mut commands: Commands, wireframe_config: Res<Wirefr
     }
 }
 
-fn extract_wireframes(mut commands: Commands, query: Query<Entity, With<Wireframe>>) {
-    for entity in query.iter() {
-        commands.get_or_spawn(entity).insert(Wireframe);
+fn extract_wireframes(mut commands: Commands, query: Query<(Entity, &Wireframe)>) {
+    for (entity, wireframe) in query.iter() {
+        commands.get_or_spawn(entity).insert(wireframe.clone());
+    }
+}
+
+fn prepare_wireframes(
+    mut commands: Commands,
+    config: Res<WireframeConfig>,
+    render_device: Res<RenderDevice>,
+    render_queue: Res<RenderQueue>,
+    mut wireframe_meta: ResMut<GlobalWireframeMeta>,
+    global_query: Query<(Entity, Option<&Wireframe>), (With<Handle<Mesh>>, With<MeshUniform>)>,
+    wireframe_query: Query<
+        (Entity, Option<&Wireframe>),
+        (With<Handle<Mesh>>, With<MeshUniform>, With<Wireframe>),
+    >,
+) {
+    wireframe_meta.uniforms.clear();
+    wireframe_meta.uniforms.push(WireframeUniform {
+        color: config.default_color.as_linear_rgba_f32().into(),
+    });
+
+    let add_wireframe_uniform = |(entity, wireframe): (Entity, Option<&Wireframe>)| {
+        let custom_color = wireframe.map(|wireframe| wireframe.color);
+        let uniform_offset = WireframeUniformOffset(if let Some(custom_color) = custom_color {
+            wireframe_meta.uniforms.push(WireframeUniform {
+                color: custom_color.as_linear_rgba_f32().into(),
+            })
+        } else {
+            0
+        });
+        commands.entity(entity).insert(uniform_offset);
+    };
+
+    if config.global {
+        global_query.for_each(add_wireframe_uniform);
+    } else {
+        wireframe_query.for_each(add_wireframe_uniform);
+    }
+
+    wireframe_meta
+        .uniforms
+        .write_buffer(&render_device, &render_queue);
+}
+
+/// Internal [`WireframePlugin`] resource.  
+/// [`GlobalWireframeBindGroup`] stores the [`BindGroup`] of wireframe data that is used on the GPU side.
+struct GlobalWireframeBindGroup {
+    bind_group: BindGroup,
+}
+
+fn queue_wireframes_bind_group(
+    mut commands: Commands,
+    render_device: Res<RenderDevice>,
+    meta: Res<GlobalWireframeMeta>,
+    bind_group: Option<Res<GlobalWireframeBindGroup>>,
+) {
+    if bind_group.is_none() {
+        commands.insert_resource(GlobalWireframeBindGroup {
+            bind_group: render_device.create_bind_group(&BindGroupDescriptor {
+                entries: &[BindGroupEntry {
+                    binding: 0,
+                    resource: meta.uniforms.binding().unwrap(),
+                }],
+                label: Some("wireframe_bind_group"),
+                layout: &meta.bind_group_layout,
+            }),
+        });
     }
 }
 
-/// Controls whether an entity should rendered in wireframe-mode if the [`WireframePlugin`] is enabled
-#[derive(Component, Debug, Clone, Default, Reflect)]
+/// Controls per-entity wireframe rendering settings (if the [`WireframePlugin`] is enabled)
+#[derive(Component, Debug, Clone, Reflect)]
 #[reflect(Component)]
-pub struct Wireframe;
+pub struct Wireframe {
+    /// The color of this wireframe.  
+    /// For the [`Entity`] that has this [`Wireframe`] component, this color overrides any that was set in [`WireframeConfig::default_color`].
+    pub color: Color,
+}
 
-#[derive(Debug, Clone, Default)]
+impl Default for Wireframe {
+    fn default() -> Self {
+        Self {
+            color: Color::WHITE,
+        }
+    }
+}
+
+/// Configuration for [`WireframePlugin`].
+#[derive(Debug, Clone)]
 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,
+    /// The default color for wireframes. If [`Self::global`] is set, any [`Entity`] that does not have a [`Wireframe`]
+    /// component attached to it will have wireframes in this color.
+    pub default_color: Color,
+}
+
+impl Default for WireframeConfig {
+    fn default() -> Self {
+        Self {
+            global: false,
+            default_color: Color::WHITE,
+        }
+    }
+}
+
+/// Internal [`WireframePlugin`] component.  
+/// [`WireframeUniformOffset`] holds the offset of a [`WireframeUniform`] in the [`GlobalWireframeMeta::uniforms`].
+#[derive(Component, Copy, Clone, Debug, Default)]
+#[repr(transparent)]
+struct WireframeUniformOffset(u32);
+
+/// Internal [`WireframePlugin`] state.  
+/// [`WireframeUniform`] is the GPU representation of a [`Wireframe`].
+#[derive(Debug, AsStd140)]
+struct WireframeUniform {
+    color: Vec4,
+}
+
+/// Internal [`WireframePlugin`] resource.  
+/// This is the data required for rendering [`Wireframe`]s.
+#[derive(Component)]
+struct GlobalWireframeMeta {
+    uniforms: DynamicUniformVec<WireframeUniform>,
+    bind_group_layout: BindGroupLayout,
 }
 
-pub struct WireframePipeline {
+impl FromWorld for GlobalWireframeMeta {
+    fn from_world(world: &mut World) -> Self {
+        let render_device = world.get_resource::<RenderDevice>().unwrap();
+
+        let bind_group_layout =
+            render_device.create_bind_group_layout(&BindGroupLayoutDescriptor {
+                entries: &[BindGroupLayoutEntry {
+                    binding: 0,
+                    visibility: ShaderStages::FRAGMENT,
+                    ty: BindingType::Buffer {
+                        ty: BufferBindingType::Uniform,
+                        has_dynamic_offset: true,
+                        min_binding_size: BufferSize::new(
+                            WireframeUniform::std140_size_static() as u64
+                        ),
+                    },
+                    count: None,
+                }],
+                label: Some("wireframe_bind_group_layout"),
+            });
+
+        Self {
+            uniforms: Default::default(),
+            bind_group_layout,
+        }
+    }
+}
+
+/// Internal [`WireframePlugin`] resource.  
+/// [`WireframePipeline`] is the specialized rendering pipeline for wireframes.
+struct WireframePipeline {
     mesh_pipeline: MeshPipeline,
+    wireframe_bind_group_layout: BindGroupLayout,
     shader: Handle<Shader>,
 }
 impl FromWorld for WireframePipeline {
     fn from_world(render_world: &mut World) -> Self {
         WireframePipeline {
             mesh_pipeline: render_world.resource::<MeshPipeline>().clone(),
+            wireframe_bind_group_layout: render_world
+                .get_resource::<GlobalWireframeMeta>()
+                .unwrap()
+                .bind_group_layout
+                .clone(),
             shader: WIREFRAME_SHADER_HANDLE.typed(),
         }
     }
@@ -94,6 +261,11 @@ impl SpecializedMeshPipeline for WireframePipeline {
         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
+            .layout
+            .as_mut()
+            .unwrap()
+            .push(self.wireframe_bind_group_layout.clone());
         descriptor.primitive.polygon_mode = PolygonMode::Line;
         descriptor.depth_stencil.as_mut().unwrap().bias.slope_scale = 1.0;
         Ok(descriptor)
@@ -160,9 +332,35 @@ fn queue_wireframes(
     }
 }
 
+/// Internal [`WireframePlugin`] render command.  
+/// [`SetWireframeBindGroup`]`<bindgroup index>` binds the [`GlobalWireframeBindGroup`] there.
+struct SetWireframeBindGroup<const I: usize>;
+impl<const I: usize> EntityRenderCommand for SetWireframeBindGroup<I> {
+    type Param = (
+        SRes<GlobalWireframeBindGroup>,
+        SQuery<Read<WireframeUniformOffset>, With<Handle<Mesh>>>,
+    );
+    #[inline]
+    fn render<'w>(
+        _view: Entity,
+        item: Entity,
+        (global_wireframe_bind_group, view_query): SystemParamItem<'w, '_, Self::Param>,
+        pass: &mut TrackedRenderPass<'w>,
+    ) -> RenderCommandResult {
+        let wireframe_uniform_offset = view_query.get(item).unwrap();
+        pass.set_bind_group(
+            I,
+            &global_wireframe_bind_group.into_inner().bind_group,
+            &[wireframe_uniform_offset.0],
+        );
+        RenderCommandResult::Success
+    }
+}
+
 type DrawWireframes = (
     SetItemPipeline,
     SetMeshViewBindGroup<0>,
     SetMeshBindGroup<1>,
+    SetWireframeBindGroup<2>,
     DrawMesh,
 );
diff --git a/examples/3d/wireframe.rs b/examples/3d/wireframe.rs
@@ -26,6 +26,10 @@ fn setup(
 ) {
     // To draw the wireframe on all entities, set this to 'true'
     wireframe_config.global = false;
+    // You can also change the default color of the wireframes, which is used only if `global` is set.
+    // This is the fallback wireframe color, which is used for any entities that do not have a Wireframe component.
+    wireframe_config.default_color = Color::AQUAMARINE;
+
     // plane
     commands.spawn_bundle(PbrBundle {
         mesh: meshes.add(Mesh::from(shape::Plane { size: 5.0 })),
@@ -41,7 +45,9 @@ fn setup(
             ..default()
         })
         // This enables wireframe drawing on this entity
-        .insert(Wireframe);
+        .insert(Wireframe {
+            color: Color::FUCHSIA,
+        });
     // light
     commands.spawn_bundle(PointLightBundle {
         transform: Transform::from_xyz(4.0, 8.0, 4.0),
