forge-ibl-bake 0.1.0

// SPDX-License-Identifier: MIT OR Apache-2.0
// forge-ibl-bake · deveraux.dev
//! Low-discrepancy sampling, GGX importance sampling, and geometry utilities.
//!
//! Provides Hammersley sequence generation (deterministic, no RNG), GGX
//! importance sampling for Monte-Carlo specular integration, Smith G visibility,
//! TBN basis construction, equirectangular env sampling, and f32↔f16 encoding.
//! All functions are pure with no global state. PROVEN by white-furnace + BRDF
//! boundary discriminators; determinism follows from Hammersley's fixed sequence.

use std::f32::consts::{PI, TAU};

// ── Hammersley low-discrepancy sequence ──────────────────────────────────────

/// 2D Hammersley point: u = i/n (stratum), v = Van der Corput radical inverse.
/// Deterministic by construction — identical inputs always produce identical outputs.
#[inline]
pub fn hammersley(i: u32, n: u32) -> [f32; 2] {
    let mut bits = i;
    bits = bits.rotate_right(16);
    bits = ((bits & 0x5555_5555) << 1) | ((bits & 0xAAAA_AAAA) >> 1);
    bits = ((bits & 0x3333_3333) << 2) | ((bits & 0xCCCC_CCCC) >> 2);
    bits = ((bits & 0x0F0F_0F0F) << 4) | ((bits & 0xF0F0_F0F0) >> 4);
    bits = ((bits & 0x00FF_00FF) << 8) | ((bits & 0xFF00_FF00) >> 8);
    let vdc = (bits as f32) * 2.328_306_4e-10; // = 1 / 2^32
    [i as f32 / n as f32, vdc]
}

// ── GGX importance sampling ──────────────────────────────────────────────────

/// GGX-distributed half-vector in tangent space (Z up = surface normal).
/// roughness is linear (not squared). Result is a unit vector.
#[inline]
pub fn importance_sample_ggx(xi: [f32; 2], roughness: f32) -> [f32; 3] {
    let a = roughness * roughness;
    let phi = TAU * xi[0];
    let denom = 1.0 + (a * a - 1.0) * xi[1];
    let cos_theta = ((1.0 - xi[1]) / denom.max(1e-7)).max(0.0).sqrt();
    let sin_theta = (1.0 - cos_theta * cos_theta).max(0.0).sqrt();
    [phi.cos() * sin_theta, phi.sin() * sin_theta, cos_theta]
}

/// Smith G2 visibility for GGX, roughness IBL remapping: k = α²/2.
#[inline]
pub fn g_smith(n_dot_v: f32, n_dot_l: f32, roughness: f32) -> f32 {
    let k = (roughness * roughness) * 0.5;
    let g1 = |nv: f32| nv / (nv * (1.0 - k) + k).max(1e-7);
    g1(n_dot_v) * g1(n_dot_l)
}

/// Pre-integrate the GGX BRDF split-sum term (Karis 2013, SIGGRAPH).
/// Returns (scale, bias) where F(roughness, NdotV) = F0·scale + bias.
pub fn integrate_brdf(n_dot_v: f32, roughness: f32, samples: u32) -> (f32, f32) {
    let v = [(1.0 - n_dot_v * n_dot_v).max(0.0).sqrt(), 0.0, n_dot_v];
    let mut a = 0.0_f32;
    let mut b = 0.0_f32;

    for i in 0..samples {
        let xi = hammersley(i, samples);
        let h = importance_sample_ggx(xi, roughness);
        let vh = dot3(v, h).max(0.0);
        let l = [
            2.0 * vh * h[0] - v[0],
            2.0 * vh * h[1] - v[1],
            2.0 * vh * h[2] - v[2],
        ];
        let n_dot_l = l[2].max(0.0);
        let n_dot_h = h[2].max(0.0);
        if n_dot_l > 0.0 {
            let g = g_smith(n_dot_v, n_dot_l, roughness);
            let g_vis = (g * vh) / (n_dot_h * n_dot_v).max(1e-7);
            let fc = (1.0 - vh).powi(5);
            a += (1.0 - fc) * g_vis;
            b += fc * g_vis;
        }
    }

    (
        (a / samples as f32).clamp(0.0, 1.0),
        (b / samples as f32).clamp(0.0, 1.0),
    )
}

// ── Cubemap geometry ─────────────────────────────────────────────────────────

/// Cubemap face direction for (u, v) ∈ [−1, 1]² on face 0..=5.
/// Convention: +X, −X, +Y, −Y, +Z, −Z (matches Vulkan / wgpu cubemap).
#[inline]
pub fn cube_face_dir(face: u32, u: f32, v: f32) -> [f32; 3] {
    let raw = match face {
        0 => [1.0, -v, -u],
        1 => [-1.0, -v, u],
        2 => [u, 1.0, v],
        3 => [u, -1.0, -v],
        4 => [u, -v, 1.0],
        5 => [-u, -v, -1.0],
        _ => [0.0, 0.0, 1.0],
    };
    normalize3(raw)
}

/// Build an orthonormal TBN basis from a surface normal.
/// Returns (tangent, bitangent, normal). Safe for any unit normal.
#[inline]
pub fn orthonormal_basis(n: [f32; 3]) -> ([f32; 3], [f32; 3], [f32; 3]) {
    let up = if n[1].abs() < 0.999 {
        [0.0f32, 1.0, 0.0]
    } else {
        [1.0, 0.0, 0.0]
    };
    let t = normalize3(cross3(up, n));
    let b = cross3(n, t);
    (t, b, n)
}

/// Rotate a vector from tangent space to world space using TBN.
#[inline]
pub fn tbn_to_world(v: [f32; 3], t: [f32; 3], b: [f32; 3], n: [f32; 3]) -> [f32; 3] {
    [
        v[0] * t[0] + v[1] * b[0] + v[2] * n[0],
        v[0] * t[1] + v[1] * b[1] + v[2] * n[1],
        v[0] * t[2] + v[1] * b[2] + v[2] * n[2],
    ]
}

// ── Equirectangular sampling ─────────────────────────────────────────────────

/// Bilinear sample of a row-major equirectangular image at world direction `dir`.
/// `pixels` is width×height with rows top-to-bottom (standard image order).
pub fn sample_equirect(pixels: &[[f32; 3]], width: u32, height: u32, dir: [f32; 3]) -> [f32; 3] {
    if width == 0 || height == 0 {
        return [0.0, 0.0, 0.0];
    }
    let d = normalize3(dir);
    // Spherical coords: longitude φ ∈ [−π,π], latitude θ ∈ [−π/2,π/2].
    let phi = d[2].atan2(d[0]); // atan2(z, x)
    let theta = d[1].clamp(-1.0, 1.0).asin();
    // Map to [0,W) and [0,H).
    let uf = (phi / TAU + 0.5).rem_euclid(1.0) * width as f32;
    let vf = (0.5 - theta / PI).clamp(0.0, 1.0) * height as f32;

    let x0 = (uf.floor() as u32).min(width - 1);
    let y0 = (vf.floor() as u32).min(height - 1);
    let x1 = (x0 + 1) % width; // wrap at seam, not clamp
    let y1 = (y0 + 1).min(height - 1);
    let fx = uf - uf.floor();
    let fy = vf - vf.floor();

    let s = |x: u32, y: u32| pixels[(y * width + x) as usize];
    let lerp3 = |a: [f32; 3], b: [f32; 3], t: f32| -> [f32; 3] {
        [
            a[0] * (1.0 - t) + b[0] * t,
            a[1] * (1.0 - t) + b[1] * t,
            a[2] * (1.0 - t) + b[2] * t,
        ]
    };
    lerp3(lerp3(s(x0, y0), s(x1, y0), fx), lerp3(s(x0, y1), s(x1, y1), fx), fy)
}

// ── f32 ↔ f16 ────────────────────────────────────────────────────────────────

/// IEEE 754 half-precision encoder (round-toward-zero, subnormals flush to ±0).
#[inline]
pub fn f32_to_f16_bits(f: f32) -> u16 {
    let bits = f.to_bits();
    let sign = ((bits >> 31) & 0x1) as u16;
    let exp = ((bits >> 23) & 0xFF) as i32;
    let mant = bits & 0x007F_FFFF;

    if exp == 0xFF {
        // Inf or NaN passthrough.
        return (sign << 15) | 0x7C00 | if mant == 0 { 0 } else { 0x0200 };
    }
    let new_exp = exp - 127 + 15;
    if new_exp >= 31 {
        return (sign << 15) | 0x7C00; // overflow → Inf
    }
    if new_exp <= 0 {
        return sign << 15; // underflow → ±0
    }
    let new_mant = (mant >> 13) as u16;
    (sign << 15) | ((new_exp as u16) << 10) | new_mant
}

/// Decode a half-precision bit pattern back to f32 (used in tests).
#[inline]
pub fn f16_bits_to_f32(bits: u16) -> f32 {
    let sign = ((bits >> 15) & 0x1) as u32;
    let exp = ((bits >> 10) & 0x1F) as i32;
    let mant = (bits & 0x03FF) as u32;
    if exp == 0 {
        if mant == 0 {
            return if sign != 0 { -0.0 } else { 0.0 };
        }
        // Subnormal half: (-1)^sign × mant × 2^(-24).
        let val = mant as f32 * (2.0f32.powi(-24));
        return if sign != 0 { -val } else { val };
    }
    if exp == 0x1F {
        return if mant != 0 {
            f32::NAN
        } else if sign != 0 {
            f32::NEG_INFINITY
        } else {
            f32::INFINITY
        };
    }
    let new_exp = (exp - 15 + 127) as u32;
    f32::from_bits((sign << 31) | (new_exp << 23) | (mant << 13))
}

// ── Geometry primitives ──────────────────────────────────────────────────────

#[inline]
pub fn normalize3(v: [f32; 3]) -> [f32; 3] {
    let m = (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt().max(1e-7);
    [v[0] / m, v[1] / m, v[2] / m]
}

#[inline]
pub fn dot3(a: [f32; 3], b: [f32; 3]) -> f32 {
    a[0] * b[0] + a[1] * b[1] + a[2] * b[2]
}

#[inline]
pub fn cross3(a: [f32; 3], b: [f32; 3]) -> [f32; 3] {
    [
        a[1] * b[2] - a[2] * b[1],
        a[2] * b[0] - a[0] * b[2],
        a[0] * b[1] - a[1] * b[0],
    ]
}