Configurable colors for wireframe (#5303)

IceSentry · web-flow · commit 068e42a01fa1 · 2023-10-13T00:06:24.000Z
# Objective - Make the wireframe colors configurable at the global level and the single mesh level - Based on #5314 This video shows what happens when playing with various settings from the example https://github.com/bevyengine/bevy/assets/8348954/1ee9aee0-fab7-4da8-bc5d-8d0562bb34e6 ## Solution - Add a `color` field to the `WireframeMaterial` - Use a `WireframeColor` component to configure the color per entity - Add a `default_color` field to `WireframeConfig` for global wireframes or wireframes with no specified color. ## Notes - Most of the docs and the general idea for `WireframeColor` came from [UberLambda](https://github.com/UberLambda) in #3677 but the code ended up completely different so I created a separate branch. ~~I'm not sure how to correctly credit them on this PR.~~ (I re-created the commit but I added them as co-author in the commit message) ~~Closes #3677 ~~Closes #5301 ~~#5314 should be merged before this PR.~~
diff --git a/crates/bevy_pbr/src/render/wireframe.wgsl b/crates/bevy_pbr/src/render/wireframe.wgsl
@@ -1,6 +1,11 @@
 #import bevy_pbr::mesh_vertex_output MeshVertexOutput
+struct WireframeMaterial {
+    color: vec4<f32>,
+};
 
+@group(1) @binding(0)
+var<uniform> material: WireframeMaterial;
 @fragment
 fn fragment(in: MeshVertexOutput) -> @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
@@ -4,6 +4,7 @@ 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::{
+    color::Color,
     extract_resource::ExtractResource,
     mesh::{Mesh, MeshVertexBufferLayout},
     prelude::Shader,
@@ -25,7 +26,6 @@ pub const WIREFRAME_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(19259
 /// This is a native only feature.
 #[derive(Debug, Default)]
 pub struct WireframePlugin;
-
 impl Plugin for WireframePlugin {
     fn build(&self, app: &mut bevy_app::App) {
         load_internal_asset!(
@@ -43,7 +43,12 @@ impl Plugin for WireframePlugin {
             .add_systems(Startup, setup_global_wireframe_material)
             .add_systems(
                 Update,
-                (apply_global_wireframe_material, apply_wireframe_material),
+                (
+                    global_color_changed.run_if(resource_changed::<WireframeConfig>()),
+                    wireframe_color_changed,
+                    apply_wireframe_material,
+                    apply_global_wireframe_material.run_if(resource_changed::<WireframeConfig>()),
+                ),
             );
     }
 }
@@ -56,6 +61,17 @@ impl Plugin for WireframePlugin {
 #[reflect(Component, Default)]
 pub struct Wireframe;
 
+/// Sets the color of the [`Wireframe`] of the entity it is attached to.
+/// If this component is present but there's no [`Wireframe`] component,
+/// it will still affect the color of the wireframe when [`WireframeConfig::global`] is set to true.
+///
+/// This overrides the [`WireframeConfig::default_color`].
+#[derive(Component, Debug, Clone, Default, Reflect)]
+#[reflect(Component, Default)]
+pub struct WireframeColor {
+    pub color: Color,
+}
+
 /// Disables wireframe rendering for any entity it is attached to.
 /// It will ignore the [`WireframeConfig`] global setting.
 ///
@@ -70,6 +86,10 @@ pub struct WireframeConfig {
     /// Whether to show wireframes for all meshes.
     /// Can be overridden for individual meshes by adding a [`Wireframe`] or [`NoWireframe`] component.
     pub global: bool,
+    /// If [`Self::global`] is set, any [`Entity`] that does not have a [`Wireframe`] component attached to it will have
+    /// wireframes using this color. Otherwise, this will be the fallback color for any entity that has a [`Wireframe`],
+    /// but no [`WireframeColor`].
+    pub default_color: Color,
 }
 
 #[derive(Resource)]
@@ -81,19 +101,53 @@ struct GlobalWireframeMaterial {
 fn setup_global_wireframe_material(
     mut commands: Commands,
     mut materials: ResMut<Assets<WireframeMaterial>>,
+    config: Res<WireframeConfig>,
 ) {
     // Create the handle used for the global material
     commands.insert_resource(GlobalWireframeMaterial {
-        handle: materials.add(WireframeMaterial {}),
+        handle: materials.add(WireframeMaterial {
+            color: config.default_color,
+        }),
     });
 }
 
+/// Updates the wireframe material of all entities without a [`WireframeColor`] or without a [`Wireframe`] component
+fn global_color_changed(
+    config: Res<WireframeConfig>,
+    mut materials: ResMut<Assets<WireframeMaterial>>,
+    global_material: Res<GlobalWireframeMaterial>,
+) {
+    if let Some(global_material) = materials.get_mut(&global_material.handle) {
+        global_material.color = config.default_color;
+    }
+}
+
+/// 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,
+        });
+    }
+}
+
 /// 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>>)>,
+    wireframes: Query<
+        (Entity, Option<&WireframeColor>),
+        (With<Wireframe>, Without<Handle<WireframeMaterial>>),
+    >,
     mut removed_wireframes: RemovedComponents<Wireframe>,
+    global_material: Res<GlobalWireframeMaterial>,
 ) {
     for e in removed_wireframes.read() {
         if let Some(mut commands) = commands.get_entity(e) {
@@ -102,8 +156,16 @@ fn apply_wireframe_material(
     }
 
     let mut wireframes_to_spawn = vec![];
-    for e in &wireframes {
-        wireframes_to_spawn.push((e, materials.add(WireframeMaterial {})));
+    for (e, wireframe_color) in &wireframes {
+        let material = if let Some(wireframe_color) = wireframe_color {
+            materials.add(WireframeMaterial {
+                color: wireframe_color.color,
+            })
+        } else {
+            // If there's no color specified we can use the global material since it's already set to use the default_color
+            global_material.handle.clone()
+        };
+        wireframes_to_spawn.push((e, material));
     }
     commands.insert_or_spawn_batch(wireframes_to_spawn);
 }
@@ -118,10 +180,6 @@ fn apply_global_wireframe_material(
     meshes_with_global_material: Query<Entity, (WireframeFilter, 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 {
@@ -130,7 +188,7 @@ fn apply_global_wireframe_material(
             material_to_spawn.push((e, global_material.handle.clone()));
         }
         commands.insert_or_spawn_batch(material_to_spawn);
-    } else if !config.global {
+    } else {
         for e in &meshes_with_global_material {
             commands.entity(e).remove::<Handle<WireframeMaterial>>();
         }
@@ -139,7 +197,10 @@ fn apply_global_wireframe_material(
 
 #[derive(Default, AsBindGroup, TypeUuid, TypePath, Debug, Clone, Asset)]
 #[uuid = "9e694f70-9963-4418-8bc1-3474c66b13b8"]
-struct WireframeMaterial {}
+pub struct WireframeMaterial {
+    #[uniform(0)]
+    pub color: Color,
+}
 
 impl Material for WireframeMaterial {
     fn fragment_shader() -> ShaderRef {
diff --git a/examples/3d/wireframe.rs b/examples/3d/wireframe.rs
@@ -9,7 +9,7 @@
 //! This is a native only feature.
 
 use bevy::{
-    pbr::wireframe::{NoWireframe, Wireframe, WireframeConfig, WireframePlugin},
+    pbr::wireframe::{NoWireframe, Wireframe, WireframeColor, WireframeConfig, WireframePlugin},
     prelude::*,
     render::{
         render_resource::WgpuFeatures,
@@ -31,12 +31,18 @@ fn main() {
             // You need to add this plugin to enable wireframe rendering
             WireframePlugin,
         ))
-        .insert_resource(WireframeToggleTimer(Timer::from_seconds(
-            1.0,
-            TimerMode::Repeating,
-        )))
+        // Wireframes can be configured with this resource. This can be changed at runtime.
+        .insert_resource(WireframeConfig {
+            // The global wireframe config enables drawing of wireframes on every mesh,
+            // except those with `NoWireframe`. Meshes with `Wireframe` will always have a wireframe,
+            // regardless of the global configuration.
+            global: true,
+            // Controls the default color of all wireframes. Used as the default color for global wireframes.
+            // Can be changed per mesh using the `WireframeColor` component.
+            default_color: Color::WHITE,
+        })
         .add_systems(Startup, setup)
-        .add_systems(Update, toggle_global_wireframe_setting)
+        .add_systems(Update, update_colors)
         .run();
 }
 
@@ -54,14 +60,15 @@ fn setup(
     });
 
     // Red cube: Never renders a wireframe
-    commands
-        .spawn(PbrBundle {
+    commands.spawn((
+        PbrBundle {
             mesh: meshes.add(Mesh::from(shape::Cube { size: 1.0 })),
             material: materials.add(Color::RED.into()),
             transform: Transform::from_xyz(-1.0, 0.5, -1.0),
             ..default()
-        })
-        .insert(NoWireframe);
+        },
+        NoWireframe,
+    ));
     // Orange cube: Follows global wireframe setting
     commands.spawn(PbrBundle {
         mesh: meshes.add(Mesh::from(shape::Cube { size: 1.0 })),
@@ -70,42 +77,89 @@ fn setup(
         ..default()
     });
     // Green cube: Always renders a wireframe
-    commands
-        .spawn(PbrBundle {
+    commands.spawn((
+        PbrBundle {
             mesh: meshes.add(Mesh::from(shape::Cube { size: 1.0 })),
             material: materials.add(Color::GREEN.into()),
             transform: Transform::from_xyz(1.0, 0.5, 1.0),
             ..default()
-        })
-        .insert(Wireframe);
+        },
+        Wireframe,
+        // This lets you configure the wireframe color of this entity.
+        // If not set, this will use the color in `WireframeConfig`
+        WireframeColor {
+            color: Color::GREEN,
+        },
+    ));
 
     // light
     commands.spawn(PointLightBundle {
         transform: Transform::from_xyz(4.0, 8.0, 4.0),
         ..default()
     });
+
     // camera
     commands.spawn(Camera3dBundle {
         transform: Transform::from_xyz(-2.0, 2.5, 5.0).looking_at(Vec3::ZERO, Vec3::Y),
         ..default()
     });
-}
 
-/// This timer is used to periodically toggle the wireframe rendering.
-#[derive(Resource)]
-struct WireframeToggleTimer(Timer);
+    // Text used to show controls
+    commands.spawn(
+        TextBundle::from_section("", TextStyle::default()).with_style(Style {
+            position_type: PositionType::Absolute,
+            top: Val::Px(10.0),
+            left: Val::Px(10.0),
+            ..default()
+        }),
+    );
+}
 
-/// Periodically turns the global wireframe setting on and off, to show the differences between
-/// [`Wireframe`], [`NoWireframe`], and just a mesh.
-fn toggle_global_wireframe_setting(
-    time: Res<Time>,
-    mut timer: ResMut<WireframeToggleTimer>,
-    mut wireframe_config: ResMut<WireframeConfig>,
+/// This system let's you toggle various wireframe settings
+fn update_colors(
+    keyboard_input: Res<Input<KeyCode>>,
+    mut config: ResMut<WireframeConfig>,
+    mut wireframe_colors: Query<&mut WireframeColor>,
+    mut text: Query<&mut Text>,
 ) {
-    if timer.0.tick(time.delta()).just_finished() {
-        // The global wireframe config enables drawing of wireframes on every mesh,
-        // except those with `NoWireframe`. Meshes with `Wireframe` will always have a wireframe,
-        // regardless of the global configuration.
-        wireframe_config.global = !wireframe_config.global;
+    text.single_mut().sections[0].value = format!(
+        "
+Controls
+---------------
+Z - Toggle global
+X - Change global color
+C - Change color of the green cube wireframe
+
+WireframeConfig
+-------------
+Global: {}
+Color: {:?}
+",
+        config.global, config.default_color,
+    );
+
+    // Toggle showing a wireframe on all meshes
+    if keyboard_input.just_pressed(KeyCode::Z) {
+        config.global = !config.global;
+    }
+
+    // Toggle the global wireframe color
+    if keyboard_input.just_pressed(KeyCode::X) {
+        config.default_color = if config.default_color == Color::WHITE {
+            Color::PINK
+        } else {
+            Color::WHITE
+        };
+    }
+
+    // Toggle the color of a wireframe using WireframeColor and not the global color
+    if keyboard_input.just_pressed(KeyCode::C) {
+        for mut color in &mut wireframe_colors {
+            color.color = if color.color == Color::GREEN {
+                Color::RED
+            } else {
+                Color::GREEN
+            };
+        }
     }
 }
