cyborg/src/shader.wgsl

struct CameraUniform {
  eye: vec4<f32>;
  vp: mat4x4<f32>;
};

struct MeshInstance {
  transform: mat4x4<f32>;
};

struct MeshData {
  instances: array<MeshInstance>;
};

struct PointLight {
  center: vec4<f32>;
  intensity: vec4<f32>;
};

struct PointLightData {
  num: i32;
  lights: array<PointLight>;
};

struct VertexInput {
  [[location(0)]] position: vec3<f32>;
  [[location(1)]] normal: vec3<f32>;
  [[location(2)]] tex_coords: vec2<f32>;
};

struct VertexOutput {
  [[builtin(position)]] clip_position: vec4<f32>;
  [[location(0)]] position: vec3<f32>;
  [[location(1)]] normal: vec3<f32>;
  [[location(2)]] tex_coords: vec2<f32>;
};

[[group(0), binding(0)]]
var<uniform> camera: CameraUniform;

[[group(0), binding(1)]]
var<storage,read> point_lights: PointLightData;

[[group(1), binding(0)]]
var<storage,read> meshes: MeshData;

[[group(2), binding(0)]] var m_sampler: sampler;
[[group(2), binding(1)]] var m_albedo: texture_2d<f32>;
[[group(2), binding(2)]] var m_metallic_roughness: texture_2d<f32>;

let PI: f32 = 3.141592;

fn D_GGX(NoH: f32, roughness: f32) -> f32 {
  let a = roughness * roughness;
  let a2 = a * a;
  let NoH2 = NoH * NoH;
  let f = NoH * (a2 - 1.0) + 1.0;
  return a2 / (PI * f * f);
}

fn g1(NoV: f32, roughness: f32, k: f32) -> f32 {
  let denom = NoV * (1.0 - k) + k;
  return NoV / denom;
}

fn G_SmithGGXCorrelated(NoV: f32, NoL: f32, roughness: f32) -> f32 {
  let r = roughness + 1.0;
  let k = (r * r) / 8.0;

  let g1l = g1(NoV, roughness, k);
  let g1v = g1(NoL, roughness, k);

  return g1l * g1v;
}

fn F_Schlick(u: f32, f0: vec3<f32>) -> vec3<f32> {
  let f = pow(1.0 - u, 5.0);
  return f + f0 * (1.0 - f);
}

fn BRDF(
  l: vec3<f32>,      // normalized light direction
  n: vec3<f32>,      // normalized surface normal
  v: vec3<f32>,      // normalized view direction
  albedo: vec3<f32>, // surface albedo
  metallic: f32,     // surface metallic
  roughness: f32,    // surface roughness
) -> vec3<f32> {
  let h = normalize(v + l);
  let NoL = max(dot(n, l), 0.0);
  let NoV = max(dot(n, v), 0.0);
  let NoH = max(dot(n, h), 0.0);
  let LoH = max(dot(l, h), 0.0);

  // calculate reflectance at surface incidence
  let f0 = mix(vec3<f32>(0.04), albedo, metallic);

  // specular BRDF
  let D = D_GGX(NoH, roughness);
  let G = G_SmithGGXCorrelated(NoV, NoL, roughness);
  let F = F_Schlick(LoH, f0);

  let numerator = (D * G) * F;
  let denominator = 4.0 * NoV * NoL;

  let Fr = numerator / max(denominator, 0.01);

  // diffuse BRDF
  let diffuse_fresnel =
    (vec3<f32>(1.0) - F_Schlick(NoL, f0)) *
    (vec3<f32>(1.0) - F_Schlick(NoV, f0));

  let lambertian = albedo / PI;
  let Fd = diffuse_fresnel * lambertian;

  // TODO multiple scattering
  return (Fr + Fd) * NoL;
}

[[stage(vertex)]]
fn vs_main(
  [[builtin(instance_index)]] mesh_idx: u32,
  vertex: VertexInput,
) -> VertexOutput {
  let transform = meshes.instances[mesh_idx].transform;
  let world_pos = transform * vec4<f32>(vertex.position, 1.0);
  let world_normal = transform * vec4<f32>(vertex.normal, 0.0);

  var out: VertexOutput;
  out.clip_position = camera.vp * world_pos;
  out.position = world_pos.xyz;
  out.normal = world_normal.xyz;
  out.tex_coords = vertex.tex_coords;
  return out;
}

[[stage(fragment)]]
fn fs_main(
  frag: VertexOutput,
) -> [[location(0)]] vec4<f32> {
  let albedo = textureSample(m_albedo, m_sampler, frag.tex_coords).rgb;
  let normal = normalize(frag.normal);
  let view = normalize(camera.eye.xyz - frag.position);

  let metallic_roughness = textureSample(m_metallic_roughness, m_sampler, frag.tex_coords).bg;
  let metallic = metallic_roughness.x;
  let roughness = metallic_roughness.y;

  var lum = vec3<f32>(0.0);
  for(var i = 0; i < 4; i = i + 1) {
    let light = point_lights.lights[i];
    let light_position = light.center.xyz - frag.position;
    let light_intensity = light.intensity.rgb;
    let light_direction = normalize(light_position);

    let radiance = light_intensity / dot(light_position, light_position);

    let reflected = BRDF(
      light_direction, normal, view,
      albedo, metallic, roughness
    );

    lum = lum + (radiance * reflected);
  }

  return vec4<f32>(lum, 1.0);
}