chore: Consolidate raytracer shader code

fazil47 · fazil47 · commit 50f864364f07 · 2025-08-27T14:43:01.000+05:30
diff --git a/app/src/rendering/renderer.rs b/app/src/rendering/renderer.rs
@@ -272,6 +272,8 @@ impl Renderer {
             } else {
                 self.frames_till_next_raytracer_compute -= 1;
             }
+        } else {
+            self.frames_till_next_raytracer_compute = 0;
         }
 
         Ok(())
diff --git a/app/src/shaders/probe_lighting/updater.wesl b/app/src/shaders/probe_lighting/updater.wesl
@@ -3,28 +3,10 @@ import super::util::{
     ProbeGridConfig,
     probe_index_to_world_pos,
     probe_index_to_tex_coords,
-    world_to_texture_coords,
 };
 import super::super::raytracer::util::{
     sample_radiance_with_lighting,
-    trace_triangles,
-    get_interpolated_color_from_hit,
-    get_triangle_data,
-    get_interpolated_color,
-    Ray,
-    HitInfo,
-    Material,
-    Vertex,
-    Triangle,
-    FLT_MAX,
-    SUN_INTENSITY,
-    materials,
-    material_stride,
-    vertices,
-    vertex_stride,
-    vertex_normal_offset,
-    vertex_material_offset,
-    indices,
+    MAX_BOUNCES
 };
 
 @group(3) @binding(0) var probe_atlas_l0: texture_storage_3d<rgba32float, write>;
@@ -87,7 +69,7 @@ fn project_to_sh(probe_pos: vec3f, probe_index: u32) -> SHCoefficients {
 
         // Use probe position and sample index for unique bounce seeds
         let bounce_seed = f32(probe_index) * 0.1 + fi * 0.01;
-        let radiance = sample_radiance_with_lighting(probe_pos, direction, bounce_seed);
+        let radiance = sample_radiance_with_lighting(probe_pos, direction, MAX_BOUNCES, bounce_seed).rgb;
 
         // L0 basis function (constant)
         let y00 = SH_C0;
@@ -109,7 +91,7 @@ fn project_to_sh(probe_pos: vec3f, probe_index: u32) -> SHCoefficients {
     }
 
     // Apply a small normalization factor to prevent excessive brightness
-    let normalization_factor = 0.5;
+    let normalization_factor = 1.0;
     coeffs.l0 *= normalization_factor;
     coeffs.l1x *= normalization_factor;
     coeffs.l1y *= normalization_factor;
diff --git a/app/src/shaders/raytracer/compute.wesl b/app/src/shaders/raytracer/compute.wesl
@@ -3,34 +3,11 @@ import super::super::util::{
     decodeRGBA,
 };
 import super::util::{
-    create_ray,
-    get_triangle_intersection_mt,
-    get_interpolated_color,
-    get_sky_color,
-    random_in_hemisphere,
-    get_triangle_data,
-    get_vertex,
-    get_material,
-    Ray,
-    HitInfo,
-    SceneParams,
-    Material,
-    Vertex,
-    Triangle,
+    sample_radiance_with_lighting,
     MAX_BOUNCES,
     MAX_LOW_FRAME_BOUNCES,
     LOW_BOUNCE_FRAMES,
     CLEAR_COLOR,
-    FLT_MAX,
-    SUN_INTENSITY,
-    materials,
-    material_stride,
-    vertices,
-    vertex_stride,
-    vertex_normal_offset,
-    vertex_material_offset,
-    indices,
-    sun_direction,
 };
 
 @group(3) @binding(0) @if(wasm) var result: texture_storage_2d<r32float, read_write>;
@@ -41,39 +18,16 @@ import super::util::{
 
 
 
-fn trace_triangles(ray: Ray) -> HitInfo {
-    let num_triangles = u32(arrayLength(&indices) / 3u);
 
-    var nearest_hit_info: HitInfo;
-    nearest_hit_info.did_hit = false;
-    nearest_hit_info.t = FLT_MAX;
 
-    for (var i: u32 = 0u; i < num_triangles; i = i + 1u) {
-        let triangle: Triangle = get_triangle_data(i);
-        let hit_info: HitInfo = get_triangle_intersection_mt(triangle, ray);
-
-        if hit_info.did_hit {
-            if hit_info.t < nearest_hit_info.t {
-                nearest_hit_info = hit_info;
-                nearest_hit_info.i = i;
-            }
-        }
-    }
-
-    return nearest_hit_info;
+fn get_camera_position(uv: vec2f) -> vec3f {
+    return (camera_to_world * vec4f(0.0, 0.0, 0.0, 1.0)).xyz;
 }
 
-fn get_interpolated_color_from_hit(hit_info: HitInfo) -> vec4f {
-    let triangle: Triangle = get_triangle_data(hit_info.i);
-    return vec4f(get_interpolated_color(hit_info, triangle), 1.0);
-}
-
-fn create_camera_ray(uv: vec2f) -> Ray {
-    let origin: vec3f = (camera_to_world * vec4f(0.0, 0.0, 0.0, 1.0)).xyz;
+fn get_camera_direction(uv: vec2f) -> vec3f {
     var direction = (camera_inverse_projection * vec4f(uv, 0.0, 1.0)).xyz;
     direction = (camera_to_world * vec4f(direction, 0.0)).xyz;
-    direction = normalize(direction);
-    return create_ray(origin, direction);
+    return normalize(direction);
 }
 
 @compute
@@ -103,9 +57,6 @@ fn main(@builtin(global_invocation_id) id: vec3u) {
     // Reference: https://github.com/gfx-rs/wgpu?tab=readme-ov-file#coordinate-systems
     let uv = ((vec2f(id.xy) + 0.5) / dims * 2.0 - 1.0) * vec2f(1.0, -1.0);
 
-    // Get a ray for the UVs
-    var ray: Ray = create_camera_ray(uv);
-
     // Load the previous frame's color
     var prev_color: vec4f;
     @if(wasm) {
@@ -114,55 +65,15 @@ fn main(@builtin(global_invocation_id) id: vec3u) {
         prev_color = textureLoad(result, coords);
     }
 
-    let scene_params = SceneParams(
-        material_stride,
-        vertex_stride,
-        vertex_normal_offset,
-        vertex_material_offset,
-        sun_direction,
-    );
-
-    // Trace the ray against the triangles
-    var ray_color: vec4f = vec4f(0.0);
-    var ray_throughput: vec4f = vec4f(1.0);
-
+    
     var max_bounces = MAX_BOUNCES;
     if frame_count < LOW_BOUNCE_FRAMES {
         max_bounces = MAX_LOW_FRAME_BOUNCES;
     }
 
-    for (var bounce = 0u; bounce <= max_bounces; bounce += 1u) {
-        if bounce == max_bounces {
-            // Set ray color as black if the ray has bounced too many times
-            ray_color = vec4f(0.0);
-            break;
-        }
-
-        let hit_info: HitInfo = trace_triangles(ray);
-        if hit_info.did_hit {
-            let tri_color = get_interpolated_color_from_hit(hit_info);
-            ray_color += tri_color * ray_throughput;
-            ray_throughput *= tri_color;
-
-            // Use frame number in random seed for temporal variation
-            ray = create_ray(
-                hit_info.p + hit_info.normal * 0.001, // Move the origin slightly to avoid self-intersection
-                random_in_hemisphere(hit_info.normal, hit_info.p + vec3f(f32(frame_count) * 0.1))
-            );
-        } else {
-            if bounce == 0u {
-                // If ray misses all triangles, return the sky color
-                ray_color = get_sky_color(ray, scene_params);
-            } else {
-                let ray_sun_intensity = max(0.0, dot(scene_params.sun_direction, ray.direction)) * SUN_INTENSITY;
-                ray_color *= ray_sun_intensity * ray_throughput;
-            }
-
-            break;
-        }
-    }
-
-    ray_color = clamp(ray_color, vec4f(0.0), vec4f(1.0));
+    var camera_position = get_camera_position(uv);
+    var camera_direction = get_camera_direction(uv);
+    var ray_color: vec4f = sample_radiance_with_lighting(camera_position, camera_direction, max_bounces, f32(frame_count) * 0.1);
 
     // Blend with previous frame
     let blend_factor = 1.0 / pow(f32(frame_count + 1), 1.05);
diff --git a/app/src/shaders/raytracer/util.wesl b/app/src/shaders/raytracer/util.wesl
@@ -1,6 +1,6 @@
 const MAX_BOUNCES: u32 = 8u;
 const MAX_LOW_FRAME_BOUNCES: u32 = 2u;
-const LOW_BOUNCE_FRAMES: u32 = 8u;
+const LOW_BOUNCE_FRAMES: u32 = 4u;
 const CLEAR_COLOR: vec4f = vec4f(0.0, 0.0, 0.0, 1.0);
 const FLT_MAX: f32 = 1e12;
 const K_EPSILON: f32 = 1e-6;
@@ -213,7 +213,7 @@ fn get_interpolated_color_from_hit(hit_info: HitInfo) -> vec4f {
     return vec4f(get_interpolated_color(hit_info, triangle), 1.0);
 }
 
-fn sample_radiance_with_lighting(origin: vec3f, direction: vec3f, bounce_seed: f32) -> vec3f {
+fn sample_radiance_with_lighting(origin: vec3f, direction: vec3f, max_bounces: u32, bounce_seed: f32) -> vec4f {
     let scene_params = SceneParams(
         material_stride,
         vertex_stride,
@@ -226,32 +226,35 @@ fn sample_radiance_with_lighting(origin: vec3f, direction: vec3f, bounce_seed: f
     var ray_color = vec4f(0.0);
     var ray_throughput = vec4f(1.0);
 
-    for (var bounce = 0u; bounce < 3u; bounce++) {
-        let hit_info = trace_triangles(ray);
+    for (var bounce = 0u; bounce <= max_bounces; bounce += 1u) {
+        if bounce == max_bounces {
+            // Set ray color as black if the ray has bounced too many times
+            ray_color = vec4f(0.0);
+            break;
+        }
 
+        let hit_info: HitInfo = trace_triangles(ray);
         if hit_info.did_hit {
-            let surface_color = get_interpolated_color_from_hit(hit_info);
+            let tri_color = get_interpolated_color_from_hit(hit_info);
+            ray_color += tri_color * ray_throughput;
+            ray_throughput *= tri_color;
 
-            ray_color += surface_color * ray_throughput;
-            ray_throughput *= surface_color;
-
-            // Generate next bounce ray
-            let seed = origin + vec3f(bounce_seed, f32(bounce), 0.0);
             ray = create_ray(
-                hit_info.p + hit_info.normal * 0.001,
-                random_in_hemisphere(hit_info.normal, seed)
+                hit_info.p + hit_info.normal * 0.001, // Move the origin slightly to avoid self-intersection
+                random_in_hemisphere(hit_info.normal, hit_info.p + bounce_seed)
             );
         } else {
-            // Hit sky - add sun and ambient lighting
             if bounce == 0u {
+                // If ray misses all triangles, return the sky color
                 ray_color = get_sky_color(ray, scene_params);
             } else {
                 let ray_sun_intensity = max(0.0, dot(scene_params.sun_direction, ray.direction)) * SUN_INTENSITY;
                 ray_color *= ray_sun_intensity * ray_throughput;
             }
+
             break;
         }
     }
 
-    return clamp(ray_color.rgb, vec3f(0.0), vec3f(1.0));
+    return clamp(ray_color, vec4f(0.0), vec4f(1.0));
 }

Original file line number	Diff line number	Diff line change
`@@ -272,6 +272,8 @@ impl Renderer {`
`272`	`272`	`} else {`
`273`	`273`	`self.frames_till_next_raytracer_compute -= 1;`
`274`	`274`	`}`
	`275`	`+ } else {`
	`276`	`+ self.frames_till_next_raytracer_compute = 0;`
`275`	`277`	`}`
`276`	`278`
`277`	`279`	`Ok(())`