wgsl include preprocessor; refactor pbr

Cargo.lock

@@ -1556,6 +1556,7 @@ dependencies = [
  "image 0.24.5",
  "log",
  "pollster",
+ "regex",
  "tobj",
  "wgpu",
  "wgpu-types",

Cargo.toml

@@ -18,6 +18,7 @@ cgmath = "0.18"
 tobj = { version = "3.2.1", features = ["async"] }
 gltf = "1.0.0"
 wgpu-types = "0.14.1"
+regex = "1.7.1"
 
 [build-dependencies]
 anyhow = "1.0.68"

src/core/state.rs

@@ -10,6 +10,7 @@ use super::light::{DrawLight, LightUniform};
 use super::model::{DrawModel, Model, ModelVertex, Vertex};
 use super::resources;
 use super::texture::Texture;
+use crate::shaders::preprocessor::preprocess_wgsl;
 
 const NUM_INSTANCES_PER_ROW: u32 = 1;
 
@@ -275,7 +276,7 @@ impl State {
             });
             let shader = wgpu::ShaderModuleDescriptor {
                 label: Some("Normal Shader"),
-                source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/test.wgsl").into()),
+                source: preprocess_wgsl("test.wgsl"),
             };
             create_render_pipeline(
                 &device,
@@ -295,7 +296,7 @@ impl State {
             });
             let shader = wgpu::ShaderModuleDescriptor {
                 label: Some("Light Shader"),
-                source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/light.wgsl").into()),
+                source: preprocess_wgsl("light.wgsl"),
             };
             create_render_pipeline(
                 &device,

src/main.rs

@@ -1,4 +1,5 @@
 mod core;
+mod shaders;
 
 fn main() {
     env_logger::init();

src/shaders/mod.rs

@@ -0,0 +1 @@
+pub mod preprocessor;

src/shaders/preprocessor.rs

@@ -0,0 +1,24 @@
+use wgpu;
+use regex::Regex;
+use std::fs::read_to_string;
+
+pub fn preprocess_wgsl(filename: &str) -> wgpu::ShaderSource {
+    let source_path = env!("CARGO_MANIFEST_DIR").to_owned() + "/src/shaders/wgsl/" + filename;
+    println!("preprocess_wgsl: loading source {}", source_path);
+    let mut source =
+        read_to_string(&source_path)
+            .unwrap();
+
+    let re = Regex::new(r"#include (.*?)\n").unwrap();
+    for cap in re.captures_iter(&source.clone()) {
+        let whole_match = &cap[0];
+        let mut full_path: String = source_path.to_owned();
+        full_path = full_path.replace(filename, &cap[1]);
+
+        println!("preprocess_wgsl: replacing {} with file {}", whole_match, full_path);
+        let nested_source = read_to_string(full_path).unwrap();
+        source = source.replace(whole_match, &nested_source);
+    }
+
+    return wgpu::ShaderSource::Wgsl(source.into());
+}

src/shaders/wgsl/brdf.wgsl

@@ -0,0 +1,63 @@
+// normal distribution function (Trowbridge-Reitz GGX)
+
+fn distribution_ggx(n: vec3<f32>, h: vec3<f32>, a: f32) -> f32 {
+    let a2 = a * a;
+    let n_dot_h = max(dot(n, h), 0.0);
+    let n_dot_h2 = n_dot_h * n_dot_h;
+
+    var denom = (n_dot_h2 * (a2 - 1.0) + 1.0);
+    denom = PI * denom * denom;
+
+    return a2 / denom;
+}
+
+// geometry function (Smith's Schlick-GGX)
+
+fn geometry_schlick_ggx(nom: f32, k: f32) -> f32 {
+    let denom = nom * (1.0 - k) + k;
+    return nom / denom;
+}
+
+fn geometry_smith(n: vec3<f32>, v: vec3<f32>, l: vec3<f32>, k: f32) -> f32 {
+    let n_dot_v = max(dot(n, v), 0.0);
+    let n_dot_l = max(dot(n, l), 0.0);
+    let ggx1 = geometry_schlick_ggx(n_dot_v, k);
+    let ggx2 = geometry_schlick_ggx(n_dot_l, k);
+    return ggx1 * ggx2;
+}
+
+// fresnel function (Fresnel-Schlick approximation)
+
+fn fresnel_schlick(cos_theta: f32, f: vec3<f32>) -> vec3<f32> {
+    return f + (1.0 - f) * pow(1.0 - cos_theta, 5.0);
+}
+
+fn brdf(
+    normal_dir: vec3<f32>,
+    light_dir: vec3<f32>,
+    view_dir: vec3<f32>,
+    half_dir: vec3<f32>,
+    albedo: vec3<f32>,
+    roughness: f32,
+    metalness: f32    
+) -> vec3<f32> {
+    // fresnel
+    var dialect = vec3(0.04);
+    dialect = mix(dialect, albedo, metalness);
+    let fresnel = fresnel_schlick(max(dot(half_dir, view_dir), 0.0), dialect);
+
+    // distribution
+    let ndf = distribution_ggx(normal_dir, half_dir, roughness);
+
+    // geometry
+    let geo = geometry_smith(normal_dir, view_dir, light_dir, roughness);
+
+    // specular
+    let nom = ndf * geo * fresnel;
+    let denom = 4.0 * max(dot(normal_dir, view_dir), 0.0) * max(dot(normal_dir, light_dir), 0.0) + 0.0001;
+    let specular = nom / denom;
+
+    let k_d = (vec3(1.0) - fresnel) * (1.0 - metalness);
+    let n_dot_l = max(dot(normal_dir, light_dir), 0.0);
+    return (k_d * albedo / PI + specular) * n_dot_l;
+}

src/shaders/wgsl/constants.wgsl

@@ -0,0 +1 @@
+let PI = 3.14159;

src/shaders/wgsl/test.wgsl

@@ -1,4 +1,5 @@
-let PI = 3.14159;
+#include constants.wgsl
+#include brdf.wgsl
 
 // Vertex shader
 
@@ -86,40 +87,6 @@ fn vs_main(
 
 // Fragment shader
 
-// normal distribution function (Trowbridge-Reitz GGX)
-
-fn distribution_ggx(n: vec3<f32>, h: vec3<f32>, a: f32) -> f32 {
-    let a2 = a * a;
-    let n_dot_h = max(dot(n, h), 0.0);
-    let n_dot_h2 = n_dot_h * n_dot_h;
-
-    var denom = (n_dot_h2 * (a2 - 1.0) + 1.0);
-    denom = PI * denom * denom;
-
-    return a2 / denom;
-}
-
-// geometry function (Smith's Schlick-GGX)
-
-fn geometry_schlick_ggx(nom: f32, k: f32) -> f32 {
-    let denom = nom * (1.0 - k) + k;
-    return nom / denom;
-}
-
-fn geometry_smith(n: vec3<f32>, v: vec3<f32>, l: vec3<f32>, k: f32) -> f32 {
-    let n_dot_v = max(dot(n, v), 0.0);
-    let n_dot_l = max(dot(n, l), 0.0);
-    let ggx1 = geometry_schlick_ggx(n_dot_v, k);
-    let ggx2 = geometry_schlick_ggx(n_dot_l, k);
-    return ggx1 * ggx2;
-}
-
-// fresnel function (Fresnel-Schlick approximation)
-
-fn fresnel_schlick(cos_theta: f32, f: vec3<f32>) -> vec3<f32> {
-    return f + (1.0 - f) * pow(1.0 - cos_theta, 5.0);
-}
-
 @group(0) @binding(0)
 var t_diffuse: texture_2d<f32>;
 @group(0)@binding(1)
@@ -131,26 +98,26 @@ var t_normal: texture_2d<f32>;
 var s_normal: sampler;
 
 @group(0)@binding(4)
-var t_metallic_roughness: texture_2d<f32>;
+var t_roughness_metalness: texture_2d<f32>;
 @group(0) @binding(5)
-var s_metallic_roughness: sampler;
+var s_roughness_metalness: sampler;
 
 @fragment
 fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
     // textures
     let object_color: vec4<f32> = textureSample(t_diffuse, s_diffuse, in.tex_coords);
     let object_normal: vec4<f32> = textureSample(t_normal, s_normal, in.tex_coords);
-    let object_metallic_roughness: vec4<f32> = textureSample(
-        t_metallic_roughness, s_metallic_roughness, in.tex_coords);
+    let object_roughness_metalness: vec4<f32> = textureSample(
+        t_roughness_metalness, s_roughness_metalness, in.tex_coords);
     // TODO: AO
 
     let albedo = object_color.xyz;
     // TODO: pass factors to shader
-    let roughness = object_metallic_roughness.y * 1.0;
-    let metallic = object_metallic_roughness.z * 1.0;
+    let roughness = object_roughness_metalness.y * 1.0;
+    let metalness = object_roughness_metalness.z * 1.0;
     
     // lighting vecs
-    let tangent_normal = object_normal.xyz * 2.0 - 1.0;
+    let normal_dir = object_normal.xyz * 2.0 - 1.0;
     var light_dir = normalize(in.tangent_light_position - in.tangent_position);
     let view_dir = normalize(in.tangent_view_position - in.tangent_position);
     let half_dir = normalize(view_dir + light_dir);
@@ -162,28 +129,19 @@ fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
     let light_attenuation = 1.0 / (1.0 + coef_a * light_dist + coef_b * light_dist * light_dist);
 
     // radiance
-    let radiance_strength = max(dot(tangent_normal, light_dir), 0.0);
+    let radiance_strength = max(dot(normal_dir, light_dir), 0.0);
     let radiance = radiance_strength * light.color.xyz * light.color.w * light_attenuation;
 
-    // fresnel
-    var f = vec3(0.04);
-    f = mix(f, albedo, metallic);
-    let fresnel = fresnel_schlick(max(dot(half_dir, view_dir), 0.0), f);
-
-    // distribution
-    let ndf = distribution_ggx(tangent_normal, half_dir, roughness);
-
-    // geometry
-    let geo = geometry_smith(tangent_normal, view_dir, light_dir, roughness);
-
-    // brdf
-    let nom = ndf * geo * fresnel;
-    let denom = 4.0 * max(dot(tangent_normal, view_dir), 0.0) * max(dot(tangent_normal, light_dir), 0.0) + 0.0001;
-    let specular = nom / denom;
-
-    let k_d = (vec3(1.0) - fresnel) * (1.0 - metallic);
-    let n_dot_l = max(dot(tangent_normal, light_dir), 0.0);
-    let total_radiance = (k_d * albedo / PI + specular) * radiance * n_dot_l;
+    // brdf shading
+    let total_radiance = radiance * brdf(
+        normal_dir,
+        light_dir,
+        view_dir,
+        half_dir,
+        albedo,
+        roughness,
+        metalness
+    );
 
     // ambient
     let ambient_light_color = vec3(1.0);