pbr lightning for volumetric sphere

assets/shaders/material.wgsl

@@ -1,6 +1,8 @@
 #import bevy_pbr::{
     forward_io::{VertexOutput, FragmentOutput},
     mesh_view_bindings::view,
+    pbr_functions::{alpha_discard, apply_pbr_lighting, main_pass_post_lighting_processing},
+    pbr_fragment::pbr_input_from_standard_material,
 }
 
 struct VoluMaterial {
@@ -12,33 +14,78 @@ struct VoluMaterial {
 @group(2) @binding(100)
 var<storage, read> volu_material: VoluMaterial;
 
-fn raymarch_hit(position_: vec3<f32>, direction: vec3<f32>) -> bool {
-    var position = position_;
-    for (var i = 0; i < 200; i++) {
-        if sphere_hit(position) {
-          return true;
-        }
-        position += direction * 0.01;
-    }
-    return false;
+fn sphere_sdf(p: vec3<f32>, center: vec3<f32>, radius: f32) -> f32 {
+    return length(p - center) - radius;
 }
 
-fn sphere_hit(p: vec3<f32>) -> bool {
-    return distance(p, volu_material.mesh_translation) < volu_material.sphere_radius;
+fn raymarch(ray_origin: vec3<f32>, ray_direction: vec3<f32>) -> f32 {
+    var t = 0.0;
+    let max_distance = 100.0;
+    let max_steps = 128;
+    let epsilon = 0.001;
+
+    for (var i = 0; i < max_steps; i++) {
+        let current_pos = ray_origin + t * ray_direction;
+        let distance = sphere_sdf(current_pos, volu_material.mesh_translation, volu_material.sphere_radius);
+
+        if distance < epsilon {
+            return t;
+        }
+
+        t += distance;
+
+        if t > max_distance {
+            break;
+        }
+    }
+
+    return -1.0;
+}
+
+fn sphere_normal(p: vec3<f32>) -> vec3<f32> {
+    let epsilon = 0.001;
+    let center = volu_material.mesh_translation;
+    let radius = volu_material.sphere_radius;
+
+    let gradient = vec3<f32>(
+        sphere_sdf(p + vec3<f32>(epsilon, 0.0, 0.0), center, radius) - sphere_sdf(p - vec3<f32>(epsilon, 0.0, 0.0), center, radius),
+        sphere_sdf(p + vec3<f32>(0.0, epsilon, 0.0), center, radius) - sphere_sdf(p - vec3<f32>(0.0, epsilon, 0.0), center, radius),
+        sphere_sdf(p + vec3<f32>(0.0, 0.0, epsilon), center, radius) - sphere_sdf(p - vec3<f32>(0.0, 0.0, epsilon), center, radius)
+    );
+
+    return normalize(gradient);
 }
 
 @fragment
 fn fragment(
-    in: VertexOutput,
+    vertex_output: VertexOutput,
     @builtin(front_facing) is_front: bool,
 ) -> FragmentOutput {
-    let view_direction = normalize(in.world_position.xyz - view.world_position);
+    var in = vertex_output;
+
+    let ray_origin = view.world_position;
+    let ray_direction = normalize(in.world_position.xyz - view.world_position);
+
+    let t = raymarch(ray_origin, ray_direction);
 
     var out: FragmentOutput;
-    if (raymarch_hit(in.world_position.xyz, view_direction)) {
-      out.color = volu_material.color;
+
+    if t > 0.0 {
+        let hit_point = ray_origin + t * ray_direction;
+        let normal = sphere_normal(hit_point);
+
+        in.world_position = vec4<f32>(hit_point, 1.0);
+        in.world_normal = normal;
+
+        var pbr_input = pbr_input_from_standard_material(in, is_front);
+
+        pbr_input.material.base_color = alpha_discard(pbr_input.material, pbr_input.material.base_color);
+
+        out.color = apply_pbr_lighting(pbr_input);
+
+        out.color = main_pass_post_lighting_processing(pbr_input, out.color);
     } else {
-      out.color = vec4(1.0, 1.0, 1.0, 0.2);
+        out.color = vec4(1.0, 1.0, 1.0, 0.1);
     }
 
     return out;

src/main.rs

@@ -1,4 +1,5 @@
 use bevy::{
+    color::palettes::css::RED,
     pbr::{ExtendedMaterial, MaterialExtension},
     prelude::*,
     render::{render_resource::*, storage::ShaderStorageBuffer},
@@ -35,6 +36,7 @@ fn setup(
         Mesh3d(meshes.add(Cuboid::new(2.0, 2.0, 2.0))),
         MeshMaterial3d(materials.add(ExtendedMaterial {
             base: StandardMaterial {
+                base_color: RED.into(),
                 alpha_mode: AlphaMode::Blend,
                 ..default()
             },