Mesh rendering support for Gizmos

assets/shaders/custom_material_screenspace_texture.wgsl

@@ -1,7 +1,7 @@
+#import bevy_render::view::coords_to_viewport_uv
 #import bevy_pbr::{
     mesh_view_bindings::view,
     forward_io::VertexOutput,
-    utils::coords_to_viewport_uv,
 }
 
 @group(1) @binding(0) var texture: texture_2d<f32>;

crates/bevy_core_pipeline/src/skybox/skybox.wgsl

@@ -1,34 +1,9 @@
-#import bevy_render::view::View
-#import bevy_pbr::utils::coords_to_viewport_uv
+#import bevy_render::view::{View, coords_to_ray_direction}
 
 @group(0) @binding(0) var skybox: texture_cube<f32>;
 @group(0) @binding(1) var skybox_sampler: sampler;
 @group(0) @binding(2) var<uniform> view: View;
 
-fn coords_to_ray_direction(position: vec2<f32>, viewport: vec4<f32>) -> vec3<f32> {
-    // Using world positions of the fragment and camera to calculate a ray direction
-    // breaks down at large translations. This code only needs to know the ray direction.
-    // The ray direction is along the direction from the camera to the fragment position.
-    // In view space, the camera is at the origin, so the view space ray direction is
-    // along the direction of the fragment position - (0,0,0) which is just the
-    // fragment position.
-    // Use the position on the near clipping plane to avoid -inf world position
-    // because the far plane of an infinite reverse projection is at infinity.
-    let view_position_homogeneous = view.inverse_projection * vec4(
-        coords_to_viewport_uv(position, viewport) * vec2(2.0, -2.0) + vec2(-1.0, 1.0),
-        1.0,
-        1.0,
-    );
-    let view_ray_direction = view_position_homogeneous.xyz / view_position_homogeneous.w;
-    // Transforming the view space ray direction by the view matrix, transforms the
-    // direction to world space. Note that the w element is set to 0.0, as this is a
-    // vector direction, not a position, That causes the matrix multiplication to ignore
-    // the translations from the view matrix.
-    let ray_direction = (view.view * vec4(view_ray_direction, 0.0)).xyz;
-
-    return normalize(ray_direction);
-}
-
 struct VertexOutput {
     @builtin(position) position: vec4<f32>,
 };
@@ -59,7 +34,7 @@ fn skybox_vertex(@builtin(vertex_index) vertex_index: u32) -> VertexOutput {
 
 @fragment
 fn skybox_fragment(in: VertexOutput) -> @location(0) vec4<f32> {
-    let ray_direction = coords_to_ray_direction(in.position.xy, view.viewport);
+    let ray_direction = coords_to_ray_direction(in.position.xy, view);
 
     // Cube maps are left-handed so we negate the z coordinate.
     return textureSample(skybox, skybox_sampler, ray_direction * vec3(1.0, 1.0, -1.0));

crates/bevy_gizmos/Cargo.toml

@@ -25,3 +25,7 @@ bevy_core = { path = "../bevy_core", version = "0.12.0" }
 bevy_reflect = { path = "../bevy_reflect", version = "0.12.0" }
 bevy_core_pipeline = { path = "../bevy_core_pipeline", version = "0.12.0" }
 bevy_transform = { path = "../bevy_transform", version = "0.12.0" }
+bevy_derive = { path = "../bevy_derive", version = "0.12.0" }
+bevy_log = { path = "../bevy_log", version = "0.12.0" }
+# Other
+thread_local = "1.0"

crates/bevy_gizmos/src/gizmos.rs

@@ -2,13 +2,14 @@
 
 use std::{f32::consts::TAU, iter};
 
+use bevy_asset::{AssetId, Handle};
 use bevy_ecs::{
     system::{Deferred, Resource, SystemBuffer, SystemMeta, SystemParam},
     world::World,
 };
 use bevy_math::{Mat2, Quat, Vec2, Vec3};
-use bevy_render::color::Color;
-use bevy_transform::TransformPoint;
+use bevy_render::{color::Color, mesh::Mesh};
+use bevy_transform::{components::Transform, TransformPoint};
 
 type PositionItem = [f32; 3];
 type ColorItem = [f32; 4];
@@ -21,6 +22,7 @@ pub(crate) struct GizmoStorage {
     pub list_colors: Vec<ColorItem>,
     pub strip_positions: Vec<PositionItem>,
     pub strip_colors: Vec<ColorItem>,
+    pub meshes: Vec<(AssetId<Mesh>, Transform, Color)>,
 }
 
 /// A [`SystemParam`] for drawing gizmos.
@@ -39,6 +41,7 @@ struct GizmoBuffer {
     list_colors: Vec<ColorItem>,
     strip_positions: Vec<PositionItem>,
     strip_colors: Vec<ColorItem>,
+    meshes: Vec<(AssetId<Mesh>, Transform, Color)>,
 }
 
 impl SystemBuffer for GizmoBuffer {
@@ -48,6 +51,7 @@ impl SystemBuffer for GizmoBuffer {
         storage.list_colors.append(&mut self.list_colors);
         storage.strip_positions.append(&mut self.strip_positions);
         storage.strip_colors.append(&mut self.strip_colors);
+        storage.meshes.append(&mut self.meshes);
     }
 }
 
@@ -571,6 +575,31 @@ impl<'s> Gizmos<'s> {
         self.linestrip_2d([tl, tr, br, bl, tl], color);
     }
 
+    /// Draw a mesh.
+    ///
+    /// # Example
+    /// ```
+    /// # use bevy_gizmos::prelude::*;
+    /// # use bevy_render::prelude::*;
+    /// # use bevy_math::prelude::*;
+    /// # use bevy_asset::prelude::*;
+    /// fn system(
+    ///     mut gizmos: Gizmos,
+    ///     asset_server: Res<AssetServer>,
+    /// ) {
+    ///     gizmos.mesh(
+    ///         asset_server.load("my_mesh.gltf#Mesh0/Primitive0"),
+    ///         Transform::IDENTITY,
+    ///         Color::GREEN,
+    ///     );
+    /// }
+    /// # bevy_ecs::system::assert_is_system(system);
+    /// ```
+    #[inline]
+    pub fn mesh(&mut self, mesh: &Handle<Mesh>, transform: Transform, color: Color) {
+        self.buffer.meshes.push((mesh.id(), transform, color))
+    }
+
     #[inline]
     fn extend_list_positions(&mut self, positions: impl IntoIterator<Item = Vec3>) {
         self.buffer

crates/bevy_gizmos/src/lib.rs

@@ -19,6 +19,7 @@
 mod arrows;
 pub mod gizmos;
 
+mod mesh_pipeline;
 #[cfg(feature = "bevy_sprite")]
 mod pipeline_2d;
 #[cfg(feature = "bevy_pbr")]
@@ -27,7 +28,11 @@ mod pipeline_3d;
 /// The `bevy_gizmos` prelude.
 pub mod prelude {
     #[doc(hidden)]
-    pub use crate::{gizmos::Gizmos, AabbGizmo, AabbGizmoConfig, GizmoConfig};
+    pub use crate::{
+        gizmos::Gizmos,
+        mesh_pipeline::{GizmoStyle, GizmoMeshBundle},
+        AabbGizmo, AabbGizmoConfig, GizmoConfig,
+    };
 }
 
 use bevy_app::{Last, Plugin, PostUpdate};
@@ -69,20 +74,27 @@ use bevy_transform::{
 use gizmos::{GizmoStorage, Gizmos};
 use std::mem;
 
+use crate::mesh_pipeline::GizmoMeshPlugin;
+
 const LINE_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(7414812689238026784);
+const UTILS_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(0866298078381252566);
 
-/// A [`Plugin`] that provides an immediate mode drawing api for visual debugging.
+/// A [`Plugin`] that provides drawing API's for visual debugging.
 pub struct GizmoPlugin;
 
 impl Plugin for GizmoPlugin {
     fn build(&self, app: &mut bevy_app::App) {
         load_internal_asset!(app, LINE_SHADER_HANDLE, "lines.wgsl", Shader::from_wgsl);
+        load_internal_asset!(app, UTILS_SHADER_HANDLE, "utils.wgsl", Shader::from_wgsl);
 
         app.register_type::<GizmoConfig>()
             .register_type::<AabbGizmoConfig>()
-            .add_plugins(UniformComponentPlugin::<LineGizmoUniform>::default())
+            .add_plugins((
+                UniformComponentPlugin::<LineGizmoUniform>::default(),
+                GizmoMeshPlugin,
+                RenderAssetPlugin::<LineGizmo>::default(),
+            ))
             .init_asset::<LineGizmo>()
-            .add_plugins(RenderAssetPlugin::<LineGizmo>::default())
             .init_resource::<LineGizmoHandles>()
             .init_resource::<GizmoConfig>()
             .init_resource::<GizmoStorage>()

crates/bevy_gizmos/src/lines.wgsl

@@ -1,5 +1,6 @@
 // TODO use common view binding
 #import bevy_render::view::View
+#import bevy_gizmos::utils::{calculate_depth, EPSILON}
 
 @group(0) @binding(0) var<uniform> view: View;
 
@@ -28,8 +29,6 @@ struct VertexOutput {
     @location(0) color: vec4<f32>,
 };
 
-const EPSILON: f32 = 4.88e-04;
-
 @vertex
 fn vertex(vertex: VertexInput) -> VertexOutput {
     var positions = array<vec3<f32>, 6>(
@@ -77,19 +76,7 @@ fn vertex(vertex: VertexInput) -> VertexOutput {
     let offset = line_width * (position.x * x_basis + position.y * y_basis);
     let screen = mix(screen_a, screen_b, position.z) + offset;
 
-    var depth: f32;
-    if line_gizmo.depth_bias >= 0. {
-        depth = clip.z * (1. - line_gizmo.depth_bias);
-    } else {
-        // depth * (clip.w / depth)^-depth_bias. So that when -depth_bias is 1.0, this is equal to clip.w
-        // and when equal to 0.0, it is exactly equal to depth.
-        // the epsilon is here to prevent the depth from exceeding clip.w when -depth_bias = 1.0
-        // clip.w represents the near plane in homogeneous clip space in bevy, having a depth
-        // of this value means nothing can be in front of this
-        // The reason this uses an exponential function is that it makes it much easier for the
-        // user to chose a value that is convenient for them
-        depth = clip.z * exp2(-line_gizmo.depth_bias * log2(clip.w / clip.z - EPSILON));
-    }
+    let depth = calculate_depth(line_gizmo.depth_bias, clip);
 
     var clip_position = vec4(clip.w * ((2. * screen) / resolution - 1.), depth, clip.w);
 

crates/bevy_gizmos/src/mesh_pipeline/gizmo_mesh.wgsl

@@ -0,0 +1,112 @@
+#import bevy_render::{maths::{affine_to_square, mat2x4_f32_to_mat3x3_unpack}, instance_index::get_instance_index, view::{View, coords_to_ray_direction}}
+#import bevy_core_pipeline::tonemapping::{screen_space_dither, powsafe, tone_mapping}
+// #import bevy_gizmos::utils::calculate_depth
+
+@group(0) @binding(0) var<uniform> view: View;
+
+struct Gizmo {
+    // Affine 4x3 matrices transposed to 3x4
+    // Use bevy_render::maths::affine_to_square to unpack
+    transform: mat3x4<f32>,
+    // 3x3 matrix packed in mat2x4 and f32 as:
+    // [0].xyz, [1].x,
+    // [1].yz, [2].xy
+    // [2].z
+    // Use bevy_render::maths::mat2x4_f32_to_mat3x3_unpack to unpack
+    inverse_transpose_transform_a: mat2x4<f32>,
+    inverse_transpose_transform_b: f32,
+    color: vec4<f32>
+};
+
+#ifdef PER_OBJECT_BUFFER_BATCH_SIZE
+@group(1) @binding(0) var<uniform> gizmos: array<Gizmo, #{PER_OBJECT_BUFFER_BATCH_SIZE}u>;
+#else
+@group(1) @binding(0) var<storage> gizmos: array<Gizmo>;
+#endif
+
+struct VertexInput {
+    @builtin(instance_index) instance_index: u32,
+    @location(0) position: vec3<f32>,
+#ifdef VERTEX_NORMALS
+    @location(1) normal: vec3<f32>,
+#endif
+#ifdef VERTEX_COLORS
+    @location(5) color: vec4<f32>,
+#endif
+}
+
+struct VertexOutput {
+    @builtin(position) position: vec4<f32>,
+    @location(0) color: vec4<f32>,
+    @location(1) world_normal: vec3<f32>,
+}
+
+@vertex
+fn vertex(in: VertexInput) -> VertexOutput {
+    var out: VertexOutput;
+
+    let gizmo = gizmos[in.instance_index];
+
+    let transform = affine_to_square(gizmo.transform);
+    out.position = view.view_proj * transform * vec4(in.position, 1.0);
+
+#ifdef VERTEX_COLORS
+    out.color = in.color * gizmo.color;
+#elseif
+    out.color = gizmo.color;
+#endif
+
+#ifdef VERTEX_NORMALS
+    let inverse_transform = mat2x4_f32_to_mat3x3_unpack(
+        gizmo.inverse_transpose_transform_a,
+        gizmo.inverse_transpose_transform_b,
+    );
+    out.world_normal = normalize(inverse_transform * in.normal);
+#endif
+
+    return out;
+}
+
+struct FragmentOutput {
+    @location(0) color: vec4<f32>,
+}
+
+@fragment
+fn fragment(in: VertexOutput) -> FragmentOutput {
+    var out: FragmentOutput;
+
+    var color = in.color;
+
+#ifdef VERTEX_NORMALS
+    #ifdef 3D 
+        let view_direction = coords_to_ray_direction(in.position.xy, view);
+    #elseif
+        let view_direction = vec3(0., 0., -1.);
+    #endif
+    // Fake lighting
+    let d = dot(view_direction, in.world_normal);
+    color = mix(color, vec4(vec3(0.), 1.), d);
+#endif
+
+// TODO: We don't need tonemapping in this shader, but screen_space_dither would be nice to have.
+// The fake lighting above also doesn't work without tonemapping for some reason.
+
+#ifdef TONEMAP_IN_SHADER
+    color = tone_mapping(color, view.color_grading);
+#ifdef DEBAND_DITHER
+    var color_rgb = color.rgb;
+
+    // Convert to sRGB
+    color_rgb = powsafe(color_rgb, 1.0 / 2.2);
+    color_rgb += screen_space_dither(in.position.xy);
+    // Convert back to Linear sRGB
+    color_rgb = powsafe(color_rgb, 2.2);
+
+    color = vec4(color_rgb, color.a);
+#endif
+#endif
+
+    out.color = color;
+
+    return out;
+}