Add colored wireframes

UberLambda · UberLambda · commit 5e26fcbe8e99 · 2022-05-15T20:54:00.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,19 +1,35 @@
-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_reflect::std_traits::ReflectDefault;
+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_component::{ExtractComponent, ExtractComponentPlugin},
+    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, VisibleEntities},
     RenderApp, RenderStage,
 };
@@ -22,6 +38,7 @@ 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;
 
@@ -36,13 +53,16 @@ impl Plugin for WireframePlugin {
 
         app.init_resource::<WireframeConfig>();
 
+        app.add_plugin(ExtractComponentPlugin::<Wireframe>::extract_visible());
         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);
         }
     }
@@ -54,31 +74,186 @@ 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);
+#[allow(clippy::type_complexity)]
+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.and_then(|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.on_all_meshes {
+        global_query.for_each(add_wireframe_uniform);
+    } else {
+        wireframe_query.for_each(add_wireframe_uniform);
+    }
+
+    wireframe_meta
+        .uniforms
+        .write_buffer(&render_device, &render_queue);
+}
+
+/// Stores the [`BindGroup`] of wireframe data that is used on the GPU side.
+///
+/// Internal [`WireframePlugin`] resource.
+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)]
-#[reflect(Component, Default)]
-pub struct Wireframe;
+/// Controls per-entity wireframe rendering settings (if the [`WireframePlugin`] is enabled)
+#[derive(Component, Debug, Default, Copy, Clone, Reflect)]
+#[reflect(Component)]
+pub struct Wireframe {
+    /// The color of this wireframe.
+    ///
+    /// If `None`, the [`Entity`] wireframe's color will be [`WireframeConfig::default_color`].
+    pub color: Option<Color>,
+}
+
+impl ExtractComponent for Wireframe {
+    type Query = &'static Wireframe;
+
+    type Filter = ();
+
+    #[inline]
+    fn extract_component(item: bevy_ecs::query::QueryItem<Self::Query>) -> Self {
+        *item
+    }
+}
 
-#[derive(Debug, Clone, Default)]
+/// Configuration resource 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.
-    pub global: bool,
+    /// Whether to show wireframes for all meshes. If `false`, only meshes with a [`Wireframe`] component will be rendered.
+    pub on_all_meshes: bool,
+    /// The default color for wireframes.
+    ///
+    /// If [`Self::on_all_meshes`] 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 [`Wireframe`]s that does have a `None`
+    /// [`Wireframe::color`].
+    pub default_color: Color,
 }
 
-pub struct WireframePipeline {
+impl Default for WireframeConfig {
+    fn default() -> Self {
+        Self {
+            on_all_meshes: false,
+            default_color: Color::WHITE,
+        }
+    }
+}
+
+/// Holds the offset of a [`WireframeUniform`] in the [`GlobalWireframeMeta::uniforms`].
+///
+/// Internal [`WireframePlugin`] component.
+#[derive(Component, Copy, Clone, Debug, Default)]
+#[repr(transparent)]
+struct WireframeUniformOffset(u32);
+
+/// [`WireframeUniform`] is the GPU representation of a [`Wireframe`].
+///
+/// Internal [`WireframePlugin`] state.
+#[derive(Debug, AsStd140)]
+struct WireframeUniform {
+    color: Vec4,
+}
+
+/// The data required for rendering [`Wireframe`]s.
+///
+/// Internal [`WireframePlugin`] resource.
+#[derive(Component)]
+struct GlobalWireframeMeta {
+    uniforms: DynamicUniformVec<WireframeUniform>,
+    bind_group_layout: BindGroupLayout,
+}
+
+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,
+        }
+    }
+}
+
+/// [`WireframePipeline`] is the specialized rendering pipeline for wireframes.
+///
+/// Internal [`WireframePlugin`] resource.
+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(),
         }
     }
@@ -95,6 +270,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)
@@ -153,13 +333,8 @@ fn queue_wireframes(
                 }
             };
 
-        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);
+        if wireframe_config.on_all_meshes {
+            material_meshes.p0().iter().for_each(add_render_phase);
         } else {
             let query = material_meshes.p1();
             visible_entities
@@ -171,9 +346,36 @@ fn queue_wireframes(
     }
 }
 
+/// [`SetWireframeBindGroup`]`<bindgroup index>` binds the [`GlobalWireframeBindGroup`] there.
+///
+/// Internal [`WireframePlugin`] render command.
+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
@@ -12,6 +12,14 @@ fn main() {
             ..default()
         })
         .add_plugins(DefaultPlugins)
+        .insert_resource(WireframeConfig {
+            // To draw the wireframe on all entities, set this to 'true'
+            on_all_meshes: false,
+            // You can also change the default color of the wireframes, which controls:
+            // - all wireframes if `WireframeConfig::on_all_meshes` is set to 'true'
+            // - the wireframe of all entities whose `Wireframe::color` is None otherwise
+            default_color: Color::AQUAMARINE,
+        })
         .add_plugin(WireframePlugin)
         .add_startup_system(setup)
         .run();
@@ -20,12 +28,9 @@ fn main() {
 /// set up a simple 3D scene
 fn setup(
     mut commands: Commands,
-    mut wireframe_config: ResMut<WireframeConfig>,
     mut meshes: ResMut<Assets<Mesh>>,
     mut materials: ResMut<Assets<StandardMaterial>>,
 ) {
-    // To draw the wireframe on all entities, set this to 'true'
-    wireframe_config.global = false;
     // plane
     commands.spawn_bundle(PbrBundle {
         mesh: meshes.add(Mesh::from(shape::Plane { size: 5.0 })),
@@ -40,8 +45,10 @@ fn setup(
             transform: Transform::from_xyz(0.0, 0.5, 0.0),
             ..default()
         })
-        // This enables wireframe drawing on this entity
-        .insert(Wireframe);
+        // This enables wireframe drawing on this entity (and it overrides the WireframeConfig's default_color)
+        .insert(Wireframe {
+            color: Some(Color::FUCHSIA),
+        });
     // light
     commands.spawn_bundle(PointLightBundle {
         transform: Transform::from_xyz(4.0, 8.0, 4.0),
