codecraft 0.1.1

// Port of glsl/pathtracing/main.glsl: constants, math utils, local frames,
// sampling, Fresnel, GGX microfacet functions, volumes, spectral helpers.

const PI: f32 = 3.141592653589793;
const PI2: f32 = 6.283185307179586;
const PI_HALF: f32 = 1.5707963267948966;
const RECIPROCAL_PI: f32 = 0.3183098861837907;
const RECIPROCAL_PI2: f32 = 0.15915494309189535;

const HUGE_DIST: f32 = 1.0e20;
// A floor on the subsurface radius, keeping its reciprocal finite. Unlike the
// ray offset this is not a geometric nudge, so it stays a fixed small length
// rather than scaling with the scene.
const RAY_OFFSET: f32 = 1.0e-4;
const DENOM_TOLERANCE: f32 = 1.0e-10;
const RADIANCE_EPSILON: f32 = 1.0e-12;
const TRANSMITTANCE_EPSILON: f32 = 1.0e-4;
const THROUGHPUT_EPSILON: f32 = 1.0e-6;
const PDF_EPSILON: f32 = 1.0e-6;
const IOR_EPSILON: f32 = 1.0e-5;
const FLT_EPSILON: f32 = 1.1920929e-7;
const INFINITY: f32 = 1.0e38;

const MATERIAL_PROPS: i32 = 0;
const MATERIAL_OPENPBR: i32 = 1;
const MATERIAL_GROUND: i32 = 2;

const ambient_ior: f32 = 1.0;

// Per-path state
var<private> rndSeed: u32;
var<private> wavelength_nm: f32;

fn sqr(x: f32) -> f32 { return x * x; }
fn sqr3(x: vec3<f32>) -> vec3<f32> { return x * x; }

fn safe_normalize(N: vec3<f32>) -> vec3<f32> {
    let l = length(N);
    return N / max(l, DENOM_TOLERANCE);
}

fn maxComponent(v: vec3<f32>) -> f32 { return max(v.x, max(v.y, v.z)); }
fn minComponent(v: vec3<f32>) -> f32 { return min(v.x, min(v.y, v.z)); }
fn avgComponent(v: vec3<f32>) -> f32 { return (v.x + v.y + v.z) / 3.0; }

fn cosTheta2(w: vec3<f32>) -> f32 { return w.z * w.z; }
fn cosTheta(w: vec3<f32>) -> f32 { return w.z; }
fn sinTheta2(w: vec3<f32>) -> f32 { return 1.0 - cosTheta2(w); }
fn sinTheta(w: vec3<f32>) -> f32 { return sqrt(max(0.0, sinTheta2(w))); }

// ---------------------------------------------------------------------------
// Local frame
// ---------------------------------------------------------------------------
struct Basis {
    nW: vec3<f32>,
    tW: vec3<f32>,
    bW: vec3<f32>,
    baryCoord: vec3<f32>,
};

fn normalToTangent(N: vec3<f32>) -> vec3<f32> {
    var T: vec3<f32>;
    if (abs(N.z) < abs(N.x)) {
        T = vec3<f32>(N.z, 0.0, -N.x);
    } else {
        T = vec3<f32>(0.0, N.z, -N.y);
    }
    return safe_normalize(T);
}

fn makeBasis(nW: vec3<f32>) -> Basis {
    var basis: Basis;
    basis.nW = safe_normalize(nW);
    basis.tW = normalToTangent(nW);
    basis.bW = cross(basis.nW, basis.tW);
    basis.baryCoord = vec3<f32>(0.0);
    return basis;
}

fn makeBasisT(nW: vec3<f32>, tW: vec3<f32>, baryCoord: vec3<f32>) -> Basis {
    var basis: Basis;
    basis.nW = safe_normalize(nW);
    basis.tW = safe_normalize(tW);
    basis.bW = cross(basis.nW, basis.tW);
    basis.baryCoord = baryCoord;
    return basis;
}

fn worldToLocal(vWorld: vec3<f32>, basis: Basis) -> vec3<f32> {
    return vec3<f32>(dot(vWorld, basis.tW), dot(vWorld, basis.bW), dot(vWorld, basis.nW));
}

fn localToWorld(vLocal: vec3<f32>, basis: Basis) -> vec3<f32> {
    return basis.tW * vLocal.x + basis.bW * vLocal.y + basis.nW * vLocal.z;
}

// V is assumed to be in local (+Z) space.
fn orthonormal_basis_ltc(V: vec3<f32>) -> mat3x3<f32> {
    let lenSqr = dot(V.xy, V.xy);
    var X: vec3<f32>;
    if (lenSqr > 0.0) {
        X = vec3<f32>(V.x, V.y, 0.0) * inverseSqrt(lenSqr);
    } else {
        X = vec3<f32>(1.0, 0.0, 0.0);
    }
    let Y = vec3<f32>(-X.y, X.x, 0.0);
    return mat3x3<f32>(X, Y, vec3<f32>(0.0, 0.0, 1.0));
}

// ---------------------------------------------------------------------------
// Sampling
// ---------------------------------------------------------------------------
fn pcg(v: u32) -> u32 {
    let state = v * 747796405u + 2891336453u;
    let word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u;
    return (word >> 22u) ^ word;
}

fn xorshift(seed: u32) -> u32 {
    var s = seed;
    s ^= s << 13u;
    s ^= s >> 17u;
    s ^= s << 5u;
    return s;
}

fn rand() -> f32 {
    rndSeed = pcg(rndSeed);
    let uint_range = 1.0 / 4294967295.0;
    return f32(rndSeed - 1u) * uint_range;
}

fn pdfHemisphereCosineWeighted(wiL: vec3<f32>) -> f32 {
    if (wiL.z <= PDF_EPSILON) { return PDF_EPSILON / PI; }
    return wiL.z / PI;
}

struct DirPdf {
    dir: vec3<f32>,
    pdf: f32,
};

fn sampleHemisphereCosineWeighted() -> DirPdf {
    let r = sqrt(rand());
    let theta = 2.0 * PI * rand();
    let x = r * cos(theta);
    let y = r * sin(theta);
    let z = sqrt(max(0.0, 1.0 - x * x - y * y));
    var out: DirPdf;
    out.pdf = max(PDF_EPSILON, abs(z) / PI);
    out.dir = vec3<f32>(x, y, z);
    return out;
}

fn powerHeuristic(a: f32, b: f32) -> f32 {
    return sqr(a) / max(DENOM_TOLERANCE, sqr(a) + sqr(b));
}

fn sample_triangle_filter(xi: f32) -> f32 {
    if (xi < 0.5) { return sqrt(2.0 * xi) - 1.0; }
    return 1.0 - sqrt(2.0 - 2.0 * xi);
}

// ---------------------------------------------------------------------------
// Fresnel
// ---------------------------------------------------------------------------
fn FresnelDielectricReflectance(mui: f32, eta_ti: f32) -> f32 {
    let c = mui;
    let mut2 = sqr(eta_ti) + sqr(c) - 1.0;
    if (mut2 <= 0.0) { return 1.0; }
    let g = sqrt(mut2);
    return 0.5 * sqr((g - c) / (g + c)) * (1.0 + sqr(((g + c) * c - 1.0) / ((g - c) * c + 1.0)));
}

fn FresnelSchlick(F0: vec3<f32>, mu: f32) -> vec3<f32> {
    return F0 + pow(1.0 - mu, 5.0) * (vec3<f32>(1.0) - F0);
}

// PR #256: clamp to [0,1]
fn FresnelF82Tint(mu: f32, F0: vec3<f32>, F82tint: vec3<f32>) -> vec3<f32> {
    let mu_bar = 1.0 / 7.0;
    let denom = mu_bar * pow(1.0 - mu_bar, 6.0);
    let Fschlick_bar = FresnelSchlick(F0, mu_bar);
    let Fschlick = FresnelSchlick(F0, mu);
    return clamp(Fschlick - mu * pow(1.0 - mu, 6.0) * (vec3<f32>(1.0) - F82tint) * Fschlick_bar / denom,
                 vec3<f32>(0.0), vec3<f32>(1.0));
}

fn E_F(eta: f32) -> f32 {
    return log((10893.0 * eta - 1438.2) / (-774.4 * sqr(eta) + 10212.0 * eta + 1.0));
}

fn DielectricFresnelAvg(eta: f32) -> f32 {
    if (eta > 1.0) { return E_F(eta); }
    else if (eta < 1.0) { return 1.0 - sqr(eta) * (1.0 - E_F(1.0 / eta)); }
    return 0.0;
}

// ---------------------------------------------------------------------------
// Microfacet multiple-scattering energy compensation (albedo scaling)
// ---------------------------------------------------------------------------
fn energy_table_lookup(table: u32, mu: f32) -> f32 {
    let x = clamp(mu, 0.0, 1.0) * f32(ENERGY_BINS) - 0.5;
    let i0 = u32(clamp(floor(x), 0.0, f32(ENERGY_BINS - 1u)));
    let i1 = min(i0 + 1u, ENERGY_BINS - 1u);
    let t = clamp(x - f32(i0), 0.0, 1.0);
    // `energy_base` is ours: the path tracer shaded one material and read
    // the tables from zero, and a table per material in one buffer needs to
    // know where its own set starts. The only edit to this file.
    let base = energy_base + table * ENERGY_BINS;
    return mix(energy_tables[base + i0], energy_tables[base + i1], t);
}

const MAX_ENERGY_SCALE: f32 = 8.0;

// Scale for the metal lobe: 1 + F0 (1/E_ss - 1)  (Turquin 2019)
fn metal_energy_scale(winputL: vec3<f32>) -> vec3<f32> {
    if (!energy_compensation) { return vec3<f32>(1.0); }
    let E = energy_table_lookup(0u, abs(winputL.z));
    let F0 = clamp(base_weight * base_color, vec3<f32>(0.0), vec3<f32>(1.0));
    return vec3<f32>(1.0) + F0 * min(1.0 / max(E, 1e-3) - 1.0, MAX_ENERGY_SCALE - 1.0);
}

// Scale for the dielectric reflection + transmission lobes: 1 / E_ss(R+T)
fn dielectric_energy_scale(winputL: vec3<f32>) -> f32 {
    if (!energy_compensation) { return 1.0; }
    let E = energy_table_lookup(select(2u, 1u, winputL.z > 0.0), abs(winputL.z));
    return min(1.0 / max(E, 1e-3), MAX_ENERGY_SCALE);
}

// ---------------------------------------------------------------------------
// GGX
// ---------------------------------------------------------------------------
fn ggx_ndf_eval(m: vec3<f32>, alpha_x: f32, alpha_y: f32) -> f32 {
    let ax = max(alpha_x, DENOM_TOLERANCE);
    let ay = max(alpha_y, DENOM_TOLERANCE);
    let Ddenom = PI * ax * ay * sqr(sqr(m.x / ax) + sqr(m.y / ay) + sqr(m.z));
    return 1.0 / max(Ddenom, DENOM_TOLERANCE);
}

// Dupuy et al. 2023, "Sampling Visible GGX Normals with Spherical Caps"
fn ggx_ndf_sample(wiL: vec3<f32>, alpha_x: f32, alpha_y: f32) -> vec3<f32> {
    let Xi = vec2<f32>(rand(), rand());
    let alpha = vec2<f32>(alpha_x, alpha_y);
    let V = normalize(vec3<f32>(wiL.xy * alpha, wiL.z));
    let phi = 2.0 * PI * Xi.x;
    let z = (1.0 - Xi.y) * (1.0 + V.z) - V.z;
    let sinT = sqrt(clamp(1.0 - z * z, 0.0, 1.0));
    let x = sinT * cos(phi);
    let y = sinT * sin(phi);
    let c = vec3<f32>(x, y, z);
    var H = c + V;
    H = normalize(vec3<f32>(H.xy * alpha, H.z));
    return H;
}

fn ggx_lambda(w: vec3<f32>, alpha_x: f32, alpha_y: f32) -> f32 {
    if (abs(w.z) < FLT_EPSILON) { return 0.0; }
    return (-1.0 + sqrt(1.0 + (sqr(alpha_x * w.x) + sqr(alpha_y * w.y)) / sqr(w.z))) / 2.0;
}

fn ggx_G1(w: vec3<f32>, alpha_x: f32, alpha_y: f32) -> f32 {
    return 1.0 / (1.0 + ggx_lambda(w, alpha_x, alpha_y));
}

fn ggx_G2(woL: vec3<f32>, wiL: vec3<f32>, alpha_x: f32, alpha_y: f32) -> f32 {
    return 1.0 / (1.0 + ggx_lambda(woL, alpha_x, alpha_y) + ggx_lambda(wiL, alpha_x, alpha_y));
}

// ---------------------------------------------------------------------------
// Volumetrics
// ---------------------------------------------------------------------------
struct Volume {
    extinction: vec3<f32>,
    albedo: vec3<f32>,
    anisotropy: f32,
};

// The Henyey-Greenstein inversion, kept identical to
// `openpbr_shader::math::sample_phase_function` -- see the note there for what
// the quotient and the square are, and for how the previous form failed. The
// engine-parity test only binds if the two draw the same sequence.
fn samplePhaseFunction(dW: vec3<f32>, anisotropy: f32) -> vec3<f32> {
    let U = rand();
    let V = rand();
    let g = anisotropy;
    var costheta: f32;
    if (abs(g) < 1.0e-3) {
        // `2U - 1`, matching the limit of the branch below; see the note on
        // `hg_cos_theta`.
        costheta = 2.0 * U - 1.0;
    } else {
        let s = (1.0 - g * g) / (1.0 - g + 2.0 * g * U);
        costheta = (1.0 + g * g - s * s) / (2.0 * g);
    }
    let sintheta = sqrt(max(0.0, 1.0 - costheta * costheta));
    let phi = 2.0 * PI * V;
    let basis = makeBasis(dW);
    return costheta * dW + sintheta * (cos(phi) * basis.tW + sin(phi) * basis.bW);
}

// Wavelength-dependent IOR according to Cauchy formula
fn specular_ior_dispersive() -> f32 {
    let lambda_C = 656.3;
    let lambda_d = 587.6;
    let lambda_F = 486.1;
    let lambda_FC2 = 1.0 / (1.0 / (lambda_F * lambda_F) - 1.0 / (lambda_C * lambda_C));
    let Vd = transmission_dispersion_abbe_number / max(DENOM_TOLERANCE, transmission_dispersion_scale);
    let nd = specular_ior;
    let B = (nd - 1.0) * lambda_FC2 / max(DENOM_TOLERANCE, Vd);
    let A = nd - B / sqr(lambda_d);
    return A + B / sqr(wavelength_nm);
}

// ---------------------------------------------------------------------------
// Color / spectral
// ---------------------------------------------------------------------------
fn luminance_srgb(C: vec3<f32>) -> f32 {
    return 0.2126 * C.r + 0.7152 * C.g + 0.0722 * C.b;
}

fn xFit_1931(w: f32) -> f32 {
    let t1 = (w - 442.0) * select(0.0374, 0.0624, w < 442.0);
    let t2 = (w - 599.8) * select(0.0323, 0.0264, w < 599.8);
    let t3 = (w - 501.1) * select(0.0382, 0.0490, w < 501.1);
    return 0.362 * exp(-0.5 * t1 * t1) + 1.056 * exp(-0.5 * t2 * t2) - 0.065 * exp(-0.5 * t3 * t3);
}
fn yFit_1931(w: f32) -> f32 {
    let t1 = (w - 568.8) * select(0.0247, 0.0213, w < 568.8);
    let t2 = (w - 530.9) * select(0.0322, 0.0613, w < 530.9);
    return 0.821 * exp(-0.5 * t1 * t1) + 0.286 * exp(-0.5 * t2 * t2);
}
fn zFit_1931(w: f32) -> f32 {
    let t1 = (w - 437.0) * select(0.0278, 0.0845, w < 437.0);
    let t2 = (w - 459.0) * select(0.0725, 0.0385, w < 459.0);
    return 1.217 * exp(-0.5 * t1 * t1) + 0.681 * exp(-0.5 * t2 * t2);
}
fn xyzFit_1931(w: f32) -> vec3<f32> {
    return vec3<f32>(xFit_1931(w), yFit_1931(w), zFit_1931(w));
}

fn xyzToRgb(XYZ: vec3<f32>) -> vec3<f32> {
    return vec3<f32>(
        dot(XYZ, vec3<f32>(3.240479, -1.537150, -0.498535)),
        dot(XYZ, vec3<f32>(-0.969256, 1.875991, 0.041556)),
        dot(XYZ, vec3<f32>(0.055648, -0.204043, 1.057311)));
}

// Spectral normalization: chosen so that E[xyzToRgb(xyzFit_1931(λ)) * SPECTRAL_NORM] = (1,1,1)
// for uniform λ in [360, 700] nm (1 / mean of the fit; the reference viewer uses the rougher
// (2.7, 3.3, 3.45), which leaves a slight pink cast).
const SPECTRAL_NORM: vec3<f32> = vec3<f32>(2.65507, 3.34676, 3.50145);