Break up shader into multiple files

Cargo.lock

@@ -1,6 +1,6 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
-version = 3
+version = 4
 
 [[package]]
 name = "ab_glyph"
@@ -1450,9 +1450,9 @@ dependencies = [
 
 [[package]]
 name = "proc-macro2"
-version = "1.0.89"
+version = "1.0.92"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f139b0662de085916d1fb67d2b4169d1addddda1919e696f3252b740b629986e"
+checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0"
 dependencies = [
  "unicode-ident",
 ]
@@ -1501,6 +1501,7 @@ dependencies = [
  "wasm-bindgen-test",
  "web-sys",
  "wgpu",
+ "wgsl-shader-assembler",
  "winit",
 ]
 
@@ -1804,9 +1805,9 @@ checksum = "6637bab7722d379c8b41ba849228d680cc12d0a45ba1fa2b48f2a30577a06731"
 
 [[package]]
 name = "syn"
-version = "2.0.87"
+version = "2.0.90"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "25aa4ce346d03a6dcd68dd8b4010bcb74e54e62c90c573f394c46eae99aba32d"
+checksum = "919d3b74a5dd0ccd15aeb8f93e7006bd9e14c295087c9896a110f490752bcf31"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -2331,6 +2332,17 @@ dependencies = [
  "web-sys",
 ]
 
+[[package]]
+name = "wgsl-shader-assembler"
+version = "0.1.0"
+source = "git+https://git.gordon.earth/matthew/wgsl-shader-assembler.git#2ce14e25d8f1685bd3486257cba4025587fd0fb9"
+dependencies = [
+ "naga",
+ "quote",
+ "regex",
+ "syn",
+]
+
 [[package]]
 name = "winapi-util"
 version = "0.1.9"

Cargo.toml

@@ -11,6 +11,7 @@ bytemuck = { version = "1.19.0", features = ["derive"] }
 glam = "0.29.2"
 log = "0.4.22"
 tiff = "0.9.1"
+wgsl-shader-assembler = { git = "https://git.gordon.earth/matthew/wgsl-shader-assembler.git" }
 
 [target.'cfg(target_arch = "x86_64")'.dependencies]
 winit = { version = "0.30.3", features = ["rwh_06"] }

rust-toolchain.toml

@@ -0,0 +1,2 @@
+[toolchain]
+channel = "nightly"

src/app/dem_renderer/dem_renderer.wgsl

@@ -1,7 +1,7 @@
 struct VertexOutput {
     @location(0) world_space_position: vec4<f32>,
     @builtin(position) position: vec4<f32>,
-};
+}
 
 struct Uniforms {
     transform: mat4x4<f32>,
@@ -20,90 +20,7 @@ var dem_texture: texture_2d<u32>;
 @group(0) @binding(2)
 var dembvh_texture: texture_2d<u32>;
 
-struct BoundingNode {
+//#include ray_intersection.wgsl
-    index: vec2<u32>,
-    level: u32
-}
-
-struct AABB {
-    min_corner: vec3<f32>,
-    max_corner: vec3<f32>
-}
-
-struct Ray { origin: vec3<f32>, direction: vec3<f32> };
-
-fn intersect_ray_with_aabb(ray: Ray, aabb: AABB) -> bool {
-    let t1 = (aabb.min_corner - ray.origin) / ray.direction;
-    let t2 = (aabb.max_corner - ray.origin) / ray.direction;
-    let t_mins = min(t1, t2);
-    let t_min = min(t_mins.x, min(t_mins.y, t_mins.z));
-    let t_maxs = max(t1, t2);
-    let t_max = max(t_maxs.x, max(t_maxs.y, t_maxs.z));
-    return t_min <= t_max;
-}
-
-fn scale_u16(value: u32, range: vec2<f32>) -> f32 {
-    return range.x + (f32(value - 1) / 65534.0) * (range.y - range.x);
-}
-
-fn get_xy_min_for_node(node: BoundingNode) -> vec2<f32> {
-    return uniforms.dem_min_corner.xy + uniforms.dem_cell_size * vec2<f32>(node.index) * exp2(f32(node.level));
-}
-
-fn get_xy_max_for_node(node: BoundingNode) -> vec2<f32> {
-    return uniforms.dem_min_corner.xy + uniforms.dem_cell_size * vec2<f32>(node.index + 1) * exp2(f32(node.level));
-}
-
-struct NodeStack {
-    stack: array<BoundingNode,64>,
-    count: u32
-}
-
-fn push_node_stack(stack: ptr<function,NodeStack>, node: BoundingNode) {
-    (*stack).stack[(*stack).count] = node;
-    (*stack).count++;
-}
-
-fn pop_node_stack(stack: ptr<function,NodeStack>) -> BoundingNode {
-    (*stack).count--;
-    return (*stack).stack[(*stack).count];
-}
-
-fn intersect_ray_with_node(ray: Ray, root_node: BoundingNode) -> bool {
-    var node_stack: NodeStack;
-    node_stack.count = 0u;
-    push_node_stack(&node_stack, root_node);
-    while node_stack.count > 0 {
-        let node = pop_node_stack(&node_stack);
-        let minmax_z = textureLoad(dembvh_texture, node.index, i32(node.level)).rg;
-        if minmax_z.r == 0 {
-            return false;
-        }
-        let min_z = scale_u16(minmax_z.r, uniforms.dem_z_range);
-        let max_z = scale_u16(minmax_z.g, uniforms.dem_z_range);
-        var aabb: AABB ;
-        aabb.min_corner = vec3<f32>(get_xy_min_for_node(node), min_z);
-        aabb.max_corner = vec3<f32>(get_xy_max_for_node(node), min_z);
-        if intersect_ray_with_aabb(ray, aabb) {
-            if node.level == 0 {
-                return true;
-            } else {
-                let next_index = node.index * 2;
-                var next_node: BoundingNode;
-                next_node.index = next_index;
-                next_node.level = node.level - 1;
-                push_node_stack(&node_stack, next_node);
-                next_node.index = next_index + vec2<u32>(1, 0);
-                push_node_stack(&node_stack, next_node);
-                next_node.index = next_index + vec2<u32>(0, 1);
-                push_node_stack(&node_stack, next_node);
-                next_node.index = next_index + vec2<u32>(1, 1);
-                push_node_stack(&node_stack, next_node);
-            }
-        }
-    }
-    return false;
-}
 
 @vertex
 fn vs_main(
@@ -125,7 +42,11 @@ fn fs_solid(vertex: VertexOutput) -> @location(0) vec4<f32> {
     root_node.index = vec2<u32>(0);
     root_node.level = textureNumLevels(dembvh_texture) - 1;
 
-    if intersect_ray_with_node(ray, root_node) {
+    if intersect_ray_with_node(dembvh_texture,
+                               uniforms.dem_min_corner,
+                               uniforms.dem_cell_size,
+                               ray,
+                               root_node) {
         return vec4<f32>(1.0);
     } else {
         discard;

src/app/dem_renderer/mod.rs

@@ -1,3 +1,5 @@
+use wgsl_shader_assembler::wgsl_module;
+
 use super::raster::{Dem, DemBvh};
 use {
     bytemuck::{Pod, Zeroable},
@@ -250,10 +252,7 @@ impl DemRenderer {
             }],
         }];
 
-        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
+        let shader = device.create_shader_module(wgsl_module!("dem_renderer.wgsl"));
-            label: None,
-            source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("dem_renderer.wgsl"))),
-        });
 
         let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
             label: Some("DemRendererPipeline"),

src/app/dem_renderer/pack_unpack.wgsl

@@ -0,0 +1,3 @@
+fn scale_u16(value: u32, range: vec2<f32>) -> f32 {
+    return range.x + (f32(value - 1) / 65534.0) * (range.y - range.x);
+}

src/app/dem_renderer/ray_intersection.wgsl

@@ -0,0 +1,96 @@
+//#include pack_unpack.wgsl
+
+struct BoundingNode {
+    index: vec2<u32>,
+    level: u32
+}
+
+struct AABB {
+    min_corner: vec3<f32>,
+    max_corner: vec3<f32>
+}
+
+struct Ray { origin: vec3<f32>, direction: vec3<f32> };
+
+fn intersect_ray_with_aabb(ray: Ray, aabb: AABB) -> bool {
+    let t1 = (aabb.min_corner - ray.origin) / ray.direction;
+    let t2 = (aabb.max_corner - ray.origin) / ray.direction;
+    let t_mins = min(t1, t2);
+    let t_min = min(t_mins.x, min(t_mins.y, t_mins.z));
+    let t_maxs = max(t1, t2);
+    let t_max = max(t_maxs.x, max(t_maxs.y, t_maxs.z));
+    return t_min <= t_max;
+}
+
+fn get_xy_min_for_node(dem_min_corner: vec2<f32>,
+                       dem_cell_size: vec2<f32>,
+                       node: BoundingNode) -> vec2<f32> {
+    return dem_min_corner.xy + dem_cell_size * vec2<f32>(node.index) * exp2(f32(node.level));
+}
+
+fn get_xy_max_for_node(dem_min_corner: vec2<f32>,
+                       dem_cell_size: vec2<f32>,
+                       node: BoundingNode) -> vec2<f32> {
+    return dem_min_corner.xy + dem_cell_size * vec2<f32>(node.index + 1) * exp2(f32(node.level));
+}
+
+struct NodeStack {
+    stack: array<BoundingNode,64>,
+    count: u32
+}
+
+fn push_node_stack(stack: ptr<function,NodeStack>, node: BoundingNode) {
+    (*stack).stack[(*stack).count] = node;
+    (*stack).count++;
+}
+
+fn pop_node_stack(stack: ptr<function,NodeStack>) -> BoundingNode {
+    (*stack).count--;
+    return (*stack).stack[(*stack).count];
+}
+
+fn intersect_ray_with_node(tree_texture: texture_2d<u32>,
+                           dem_min_corner: vec2<f32>,
+                           dem_cell_size: vec2<f32>,
+                           ray: Ray,
+                           root_node: BoundingNode) -> bool {
+    var node_stack: NodeStack;
+    node_stack.count = 0u;
+    push_node_stack(&node_stack, root_node);
+    while node_stack.count > 0 {
+        let node = pop_node_stack(&node_stack);
+        let minmax_z = textureLoad(tree_texture, node.index, i32(node.level)).rg;
+        if minmax_z.r == 0 {
+            return false;
+        }
+        let min_z = scale_u16(minmax_z.r, uniforms.dem_z_range);
+        let max_z = scale_u16(minmax_z.g, uniforms.dem_z_range);
+        var aabb: AABB ;
+        aabb.min_corner = vec3<f32>(get_xy_min_for_node(dem_min_corner,
+                                                        dem_cell_size,
+                                                        node),
+                                    min_z);
+        aabb.max_corner = vec3<f32>(get_xy_max_for_node(dem_min_corner,
+                                                        dem_cell_size,
+                                                        node),
+                                    max_z);
+        if intersect_ray_with_aabb(ray, aabb) {
+            if node.level == 0 {
+                return true;
+            } else {
+                let next_index = node.index * 2;
+                var next_node: BoundingNode;
+                next_node.index = next_index;
+                next_node.level = node.level - 1;
+                push_node_stack(&node_stack, next_node);
+                next_node.index = next_index + vec2<u32>(1, 0);
+                push_node_stack(&node_stack, next_node);
+                next_node.index = next_index + vec2<u32>(0, 1);
+                push_node_stack(&node_stack, next_node);
+                next_node.index = next_index + vec2<u32>(1, 1);
+                push_node_stack(&node_stack, next_node);
+            }
+        }
+    }
+    return false;
+}