WIP: Update depth buffer with parallax mapping result

crates/bevy_core_pipeline/src/debug_gradient.wgsl

@@ -0,0 +1,41 @@
+#define_import_path bevy_core_pipeline::debug_gradient
+
+// Define a function to map a range of values to a high contrast color gradient
+// optimized for representational clarity.
+// this uses the matplotlib plasma color map.
+fn get_color(value: f32) -> vec3<f32> {
+    let plasma_colors = array(
+        vec3(0.868,0.941,0.011), // #f0f921
+        vec3(0.966,0.386,0.0326), // #fca636
+        vec3(0.753,0.126,0.121), // #e16462
+        vec3(0.444,0.0188,0.282), // #b12a90
+        vec3(0.144,0.0,0.396), // #6a00a8
+        vec3(0.00142,0.000488,0.245), // #0d0887
+    );
+    if value < 1.0 {
+        return plasma_colors[0];
+    } else if value < 2.0 {
+        return plasma_colors[1];
+    } else if value < 3.0 {
+        return plasma_colors[2];
+    } else if value < 4.0 {
+        return plasma_colors[3];
+    } else if value < 5.0 {
+        return plasma_colors[4];
+    } else {
+        return plasma_colors[5];
+    }
+}
+
+// Return color sampled from the plasma color map
+fn debug_gradient(value: f32) -> vec4<f32> {
+    let colors_offset = clamp(value, 0.0, 1.0) * 5.0;
+    let low_ratio = fract(colors_offset);
+
+    let low_color_index = u32(floor(colors_offset));
+    let low_color = get_color(colors_offset);
+    let high_color = get_color(colors_offset + 1.0);
+
+    return vec4(mix(low_color, high_color, low_ratio), 1.0);
+}
+

crates/bevy_core_pipeline/src/lib.rs

@@ -15,6 +15,7 @@ mod taa;
 pub mod tonemapping;
 pub mod upscaling;
 
+use bevy_reflect::TypeUuid;
 pub use skybox::Skybox;
 
 /// Experimental features that are not yet finished. Please report any issues you encounter!
@@ -48,9 +49,12 @@ use crate::{
     upscaling::UpscalingPlugin,
 };
 use bevy_app::{App, Plugin};
-use bevy_asset::load_internal_asset;
+use bevy_asset::{load_internal_asset, HandleUntyped};
 use bevy_render::{extract_resource::ExtractResourcePlugin, prelude::Shader};
 
+pub const DEBUG_GRADIENT_HANDLE: HandleUntyped =
+    HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 17293936987160423658);
+
 #[derive(Default)]
 pub struct CorePipelinePlugin;
 
@@ -62,6 +66,12 @@ impl Plugin for CorePipelinePlugin {
             "fullscreen_vertex_shader/fullscreen.wgsl",
             Shader::from_wgsl
         );
+        load_internal_asset!(
+            app,
+            DEBUG_GRADIENT_HANDLE,
+            "debug_gradient.wgsl",
+            Shader::from_wgsl
+        );
 
         app.register_type::<ClearColor>()
             .register_type::<ClearColorConfig>()

crates/bevy_pbr/src/prepass/mod.rs

@@ -393,7 +393,7 @@ where
             vertex_attributes.push(Mesh::ATTRIBUTE_UV_0.at_shader_location(1));
         }
 
-        if key.mesh_key.contains(MeshPipelineKey::NORMAL_PREPASS) {
+        if true || key.mesh_key.contains(MeshPipelineKey::NORMAL_PREPASS) {
             vertex_attributes.push(Mesh::ATTRIBUTE_NORMAL.at_shader_location(2));
             shader_defs.push("NORMAL_PREPASS".into());
 
@@ -410,13 +410,6 @@ where
             shader_defs.push("MOTION_VECTOR_PREPASS".into());
         }
 
-        if key
-            .mesh_key
-            .intersects(MeshPipelineKey::NORMAL_PREPASS | MeshPipelineKey::MOTION_VECTOR_PREPASS)
-        {
-            shader_defs.push("PREPASS_FRAGMENT".into());
-        }
-
         if layout.contains(Mesh::ATTRIBUTE_JOINT_INDEX)
             && layout.contains(Mesh::ATTRIBUTE_JOINT_WEIGHT)
         {
@@ -431,39 +424,38 @@ where
         let vertex_buffer_layout = layout.get_layout(&vertex_attributes)?;
 
         // Setup prepass fragment targets - normals in slot 0 (or None if not needed), motion vectors in slot 1
+        let has_normal = key.mesh_key.contains(MeshPipelineKey::NORMAL_PREPASS);
+        let has_motion_vectors = key
+            .mesh_key
+            .contains(MeshPipelineKey::MOTION_VECTOR_PREPASS);
+        let has_prepass_fragment = true || has_normal || has_motion_vectors;
         let mut targets = vec![];
-        targets.push(
-            key.mesh_key
-                .contains(MeshPipelineKey::NORMAL_PREPASS)
-                .then_some(ColorTargetState {
-                    format: NORMAL_PREPASS_FORMAT,
-                    blend: Some(BlendState::REPLACE),
-                    write_mask: ColorWrites::ALL,
-                }),
-        );
-        targets.push(
-            key.mesh_key
-                .contains(MeshPipelineKey::MOTION_VECTOR_PREPASS)
-                .then_some(ColorTargetState {
-                    format: MOTION_VECTOR_PREPASS_FORMAT,
-                    blend: Some(BlendState::REPLACE),
-                    write_mask: ColorWrites::ALL,
-                }),
-        );
-        if targets.iter().all(Option::is_none) {
-            // if no targets are required then clear the list, so that no fragment shader is required
-            // (though one may still be used for discarding depth buffer writes)
+        targets.push(has_normal.then_some(ColorTargetState {
+            format: NORMAL_PREPASS_FORMAT,
+            blend: Some(BlendState::REPLACE),
+            write_mask: ColorWrites::ALL,
+        }));
+        targets.push(has_motion_vectors.then_some(ColorTargetState {
+            format: MOTION_VECTOR_PREPASS_FORMAT,
+            blend: Some(BlendState::REPLACE),
+            write_mask: ColorWrites::ALL,
+        }));
+        if targets.iter().all(|opt| opt.is_none()) {
             targets.clear();
         }
 
         // The fragment shader is only used when the normal prepass or motion vectors prepass
         // is enabled or the material uses alpha cutoff values and doesn't rely on the standard
         // prepass shader
-        let fragment_required = !targets.is_empty()
-            || ((key.mesh_key.contains(MeshPipelineKey::ALPHA_MASK)
-                || blend_key == MeshPipelineKey::BLEND_PREMULTIPLIED_ALPHA
-                || blend_key == MeshPipelineKey::BLEND_ALPHA)
-                && self.material_fragment_shader.is_some());
+        let has_alpha = key.mesh_key.contains(MeshPipelineKey::ALPHA_MASK)
+            || blend_key == MeshPipelineKey::BLEND_PREMULTIPLIED_ALPHA
+            || blend_key == MeshPipelineKey::BLEND_ALPHA;
+        let has_prepass_material = has_alpha && self.material_fragment_shader.is_some();
+        let fragment_required = has_prepass_fragment || has_prepass_material;
+
+        if fragment_required {
+            shader_defs.push("PREPASS_FRAGMENT".into());
+        }
 
         let fragment = fragment_required.then(|| {
             // Use the fragment shader from the material

crates/bevy_pbr/src/prepass/prepass.wgsl

@@ -1,5 +1,6 @@
 #import bevy_pbr::prepass_bindings
 #import bevy_pbr::mesh_functions
+#import bevy_pbr::pbr_bindings
 
 // Most of these attributes are not used in the default prepass fragment shader, but they are still needed so we can
 // pass them to custom prepass shaders like pbr_prepass.wgsl.
@@ -25,6 +26,7 @@ struct Vertex {
 
 struct VertexOutput {
     @builtin(position) clip_position: vec4<f32>,
+    @location(3) world_position: vec4<f32>,
 
 #ifdef VERTEX_UVS
     @location(0) uv: vec2<f32>,
@@ -38,7 +40,6 @@ struct VertexOutput {
 #endif // NORMAL_PREPASS
 
 #ifdef MOTION_VECTOR_PREPASS
-    @location(3) world_position: vec4<f32>,
     @location(4) previous_world_position: vec4<f32>,
 #endif // MOTION_VECTOR_PREPASS
 }
@@ -53,6 +54,7 @@ fn vertex(vertex: Vertex) -> VertexOutput {
     var model = mesh.model;
 #endif // SKINNED
 
+    out.world_position = mesh_position_local_to_world(model, vec4<f32>(vertex.position, 1.0));
     out.clip_position = mesh_position_local_to_clip(model, vec4(vertex.position, 1.0));
 #ifdef DEPTH_CLAMP_ORTHO
         out.clip_position.z = min(out.clip_position.z, 1.0);
@@ -75,7 +77,6 @@ fn vertex(vertex: Vertex) -> VertexOutput {
 #endif // NORMAL_PREPASS
 
 #ifdef MOTION_VECTOR_PREPASS
-    out.world_position = mesh_position_local_to_world(model, vec4<f32>(vertex.position, 1.0));
     out.previous_world_position = mesh_position_local_to_world(mesh.previous_model, vec4<f32>(vertex.position, 1.0));
 #endif // MOTION_VECTOR_PREPASS
 
@@ -84,8 +85,12 @@ fn vertex(vertex: Vertex) -> VertexOutput {
 
 #ifdef PREPASS_FRAGMENT
 struct FragmentInput {
+#ifdef VERTEX_UVS
+    @location(0) uv: vec2<f32>,
+#endif // VERTEX_UVS
 #ifdef NORMAL_PREPASS
     @location(1) world_normal: vec3<f32>,
+    @location(2) world_tangent: vec4<f32>,
 #endif // NORMAL_PREPASS
 
 #ifdef MOTION_VECTOR_PREPASS
@@ -95,6 +100,7 @@ struct FragmentInput {
 }
 
 struct FragmentOutput {
+    @builtin(frag_depth) depth: f32,
 #ifdef NORMAL_PREPASS
     @location(0) normal: vec4<f32>,
 #endif // NORMAL_PREPASS

crates/bevy_pbr/src/render/parallax_mapping.wgsl

@@ -24,9 +24,9 @@ fn parallaxed_uv(
     uv: vec2<f32>,
     // The vector from the camera to the fragment at the surface in tangent space
     Vt: vec3<f32>,
-) -> vec2<f32> {
+) -> vec3<f32> {
     if max_layer_count < 1.0 {
-        return uv;
+        return vec3(uv, 0.0);
     }
     var uv = uv;
 
@@ -110,7 +110,12 @@ fn parallaxed_uv(
     current_layer_depth += mix(next_depth, previous_depth, weight);
 #endif
 
-    // Note: `current_layer_depth` is not returned, but may be useful
-    // for light computation later on in future improvements of the pbr shader.
-    return uv;
+    return vec3<f32>(uv, sample_depth_map(uv) * depth_scale);
+}
+
+/// Additional depth in tangent space resulting from parallaxing.
+///
+/// This is a simple pythagorean theorem.
+fn additional_depth(old_uv: vec2<f32>, new_uv: vec2<f32>, depth: f32) -> f32 {
+    return length(vec3(new_uv - old_uv, depth));
 }

crates/bevy_pbr/src/render/pbr.wgsl

@@ -10,19 +10,22 @@
 #import bevy_pbr::fog
 #import bevy_pbr::pbr_functions
 #import bevy_pbr::parallax_mapping
+#import bevy_core_pipeline::debug_gradient
 
 #import bevy_pbr::prepass_utils
 
 struct FragmentInput {
     @builtin(front_facing) is_front: bool,
     @builtin(position) frag_coord: vec4<f32>,
+    @builtin(sample_index) sample_index: u32,
     #import bevy_pbr::mesh_vertex_output
 };
 
 @fragment
 fn fragment(in: FragmentInput) -> @location(0) vec4<f32> {
     let is_orthographic = view.projection[3].w == 1.0;
     let V = calculate_view(in.world_position, is_orthographic);
+    var my_depth = 100000.0;
 #ifdef VERTEX_UVS
     var uv = in.uv;
 #ifdef VERTEX_TANGENTS
@@ -31,8 +34,9 @@ fn fragment(in: FragmentInput) -> @location(0) vec4<f32> {
         let T = in.world_tangent.xyz;
         let B = in.world_tangent.w * cross(N, T);
         // Transform V from fragment to camera in world space to tangent space.
-        let Vt = vec3(dot(V, T), dot(V, B), dot(V, N));
-        uv = parallaxed_uv(
+        let TBN = mat3x3(T, B, N);
+        let Vt = V * TBN; //vec3(dot(V, T), dot(V, B), dot(V, N));
+        let parallaxed = parallaxed_uv(
             material.parallax_depth_scale,
             material.max_parallax_layer_count,
             material.max_relief_mapping_search_steps,
@@ -42,6 +46,15 @@ fn fragment(in: FragmentInput) -> @location(0) vec4<f32> {
             // about.
             -Vt,
         );
+        uv = parallaxed.xy;
+        if parallaxed.z > 0.0001 {
+            let NBT = transpose(mat3x3(T, B, N));
+            let tangent_extra_depth = Vt * additional_depth(in.uv, uv, parallaxed.z);
+            let extra_depth = tangent_extra_depth * NBT;
+            let parallaxed_view_position = view.view_proj * (in.world_position - vec4(extra_depth, 0.0));
+            // let parallaxed_view_position = view.view_proj * in.world_position;
+            my_depth = parallaxed_view_position.z / parallaxed_view_position.w;
+        }
     }
 #endif
 #endif
@@ -151,5 +164,14 @@ fn fragment(in: FragmentInput) -> @location(0) vec4<f32> {
 #ifdef PREMULTIPLY_ALPHA
     output_color = premultiply_alpha(material.flags, output_color);
 #endif
+    let depth = prepass_depth(in.frag_coord, in.sample_index);
+    // if depth > my_depth {
+    //     discard;
+    // }
+    // var depth = my_depth;
+    // if depth == 0.0 {
+    //     depth = in.frag_coord.z;
+    // }
+    // output_color = debug_gradient(1.0 - depth * 14.0);
     return output_color;
 }