From deeea06e9227149b3563de8787e9676508e337c2 Mon Sep 17 00:00:00 2001
From: Jarrod Doyle <thejarroddoyle@gmail.com>
Date: Sat, 11 May 2024 13:56:57 +0100
Subject: [PATCH] Add basic dda traversal shader

---
 res/basic_dda/traverse.wgsl | 84 +++++++++++++++++++++++++++++++++++++
 src/renderer.rs             | 19 +++++++++
 2 files changed, 103 insertions(+)
 create mode 100644 res/basic_dda/traverse.wgsl

diff --git a/res/basic_dda/traverse.wgsl b/res/basic_dda/traverse.wgsl
new file mode 100644
index 0000000..d8d1427
--- /dev/null
+++ b/res/basic_dda/traverse.wgsl
@@ -0,0 +1,84 @@
+@group(0) @binding(0) var output: texture_storage_2d<rgba8unorm, write>;
+@group(0) @binding(2) var<storage, read_write> rays: array<Ray>;
+
+struct Ray {
+    pos: vec3<f32>,
+    dist: f32,
+    dir: vec3<f32>,
+    hit: u32,
+}
+
+const MAX_RAY_STEPS: i32 = 64;
+
+fn sd_sphere(p: vec3<f32>, d: f32) -> f32 {
+    return length(p) - d;
+}
+
+fn sd_box(p: vec3<f32>, b: vec3<f32>) -> f32 {
+    let d = abs(p) - b;
+    return min(max(d.x, max(d.y, d.z)), 0.0) + length(max(d, vec3<f32>(0.0)));
+}
+
+fn get_voxel(c: vec3<i32>) -> bool {
+    let p = vec3<f32>(c) + vec3<f32>(0.5);
+    let d = max(-sd_sphere(p, 7.5), sd_box(p, vec3<f32>(6.0)));
+    return d < 0.0;
+}
+
+fn get_mask(side_dist: vec3<f32>) -> vec3<f32> {
+    let smallest = min(side_dist.x, min(side_dist.y, side_dist.z));
+    if (smallest == side_dist.x) {
+        return vec3<f32>(1.0, 0.0, 0.0);
+    }
+    else if (smallest == side_dist.y) {
+        return vec3<f32>(0.0, 1.0, 0.0);
+    }
+    return vec3<f32>(0.0, 0.0, 1.0);
+}
+
+fn traverse(ray: Ray) -> Ray {
+    var map_pos = floor(ray.pos);
+    var ray_sign = sign(ray.dir);
+    var delta_dist = 1.0 / ray.dir;
+    var side_dist = delta_dist * ((map_pos - ray.pos) + ray_sign * 0.5 + 0.5);
+    var mask = get_mask(side_dist);
+
+    for (var i: i32 = 0; i < MAX_RAY_STEPS; i++) {
+        if (get_voxel(vec3<i32>(map_pos))) {
+            let mini = delta_dist * ((map_pos - ray.pos) - ray_sign * 0.5 + 0.5);
+            let d = max(mini.x, max(mini.y, mini.z));
+            return Ray(ray.pos, d, ray.dir, 1);
+        }
+
+        mask = get_mask(side_dist);
+        map_pos += mask * ray_sign;
+        side_dist += mask * ray_sign * delta_dist;
+    }
+
+    return ray;
+}
+
+
+@compute @workgroup_size(8, 8, 1)
+fn compute(@builtin(global_invocation_id) global_id: vec3<u32>) {
+    let img_coord = global_id.xy;
+    let img_dims = textureDimensions(output);
+
+    // This discards the extra pixels in cases where the image size isn't perfectly divisible by the kernel.xy
+    if (img_coord.x >= img_dims.x || img_coord.y >= img_dims.y) {
+        return;
+    }
+
+    let idx = img_coord.x + img_coord.y * img_dims.x;
+    rays[idx] = traverse(rays[idx]);
+
+    // Temp
+    var color = vec4<f32>(0.5, 0.6, 0.7, 1.0);
+    let ray = rays[idx];
+    if (ray.hit == 1) {
+        let intersect = ray.pos + ray.dir * (ray.dist + 0.0001);
+        let uv3d = intersect - floor(intersect);
+        color = vec4<f32>(uv3d, 1.0);
+    }
+    textureStore(output, img_coord, color);
+}
\ No newline at end of file
diff --git a/src/renderer.rs b/src/renderer.rs
index d7c80bc..93246c8 100644
--- a/src/renderer.rs
+++ b/src/renderer.rs
@@ -17,6 +17,7 @@ pub struct VoxelRenderer {
     render_pipeline: wgpu::RenderPipeline,
     raygen_bind_group: wgpu::BindGroup,
     raygen_pipeline: wgpu::ComputePipeline,
+    raycast_pipeline: wgpu::ComputePipeline,
 }
 
 impl VoxelRenderer {
@@ -83,11 +84,23 @@ impl VoxelRenderer {
             .with_layout("Raygen PL", &[&raygen_layout], &[])
             .build(context);
 
+        let raycast_shader = context
+            .device
+            .create_shader_module(wgpu::include_wgsl!(concat!(
+                env!("CARGO_MANIFEST_DIR"),
+                "/res/basic_dda/traverse.wgsl"
+            )));
+
+        let raycast_pipeline = crawl::ComputePipelineBuilder::new("Raycast P", &raycast_shader)
+            .with_layout("Raygen PL", &[&raygen_layout], &[])
+            .build(context);
+
         Ok(Self {
             render_texture,
             render_pipeline,
             raygen_pipeline,
             raygen_bind_group,
+            raycast_pipeline,
         })
     }
 }
@@ -101,6 +114,12 @@ impl crawl::Pass for VoxelRenderer {
         compute_pass.dispatch_workgroups(size.width / 8, size.height / 8, 1);
         drop(compute_pass);
 
+        let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
+        compute_pass.set_pipeline(&self.raycast_pipeline);
+        compute_pass.set_bind_group(0, &self.raygen_bind_group, &[]);
+        compute_pass.dispatch_workgroups(size.width / 8, size.height / 8, 1);
+        drop(compute_pass);
+
         let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
             label: Some("Render Pass"),
             color_attachments: &[Some(wgpu::RenderPassColorAttachment {