ph-color 0.1.1

//! Feature-gated linear-sRGB <-> Oklab conversion.
//!
//! `a` and `b` are stored as UQ0.16 with a +0.5 offset (`0.5` -> `32768` units
//! via Q4.28 addend [`A_B_BIAS_Q428`]), because [`Color`] channels are unsigned.
//! Cube-root interpolation error is [`CBRT_MAX_ERR_LSB`] LSBs versus the host reference.
//! Tables and matrices are generated from the auditable host formulas in
//! `ph-color-bake`; run `cargo xtask generate` after changing them.

use crate::arith::{div_65535, mul_acc, rem_65535, shift_round_sat_q0_16};
use crate::chromaticities::Chromaticities;
use crate::color::Color;
use crate::encoding::Linear;
use crate::fixed::{Q0_16, Q4_28};
use crate::space::{ColorSpace, Perceptual, Srgb};

#[path = "generated/oklab.rs"]
mod generated;

pub use generated::{
    A_B_BIAS_Q428, CBRT, CBRT_MAX_ERR_LSB, OKLAB_M1, OKLAB_M1_INV, OKLAB_M2, OKLAB_M2_INV,
    OKLAB_M2_INV_OFFSET,
};

/// Ottosson Oklab. Primaries/white are unused identity metadata.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Oklab;

impl ColorSpace for Oklab {
    const PRIMARIES: [Chromaticities; 3] = Srgb::PRIMARIES;
    const WHITE: Chromaticities = Srgb::WHITE;
    const NAME: &'static str = "Oklab";
}

impl Perceptual for Oklab {}

/// One matrix row plus an affine Q4.28 addend.
///
/// The addend is `addend_q428 * 65535` — a Q4.28 coefficient against full
/// scale — which is exactly what [`mul_acc`] computes, so it reuses that
/// kernel instead of spelling out a second `i64 * i64` (which links
/// `__aeabi_lmul` on a target without a widening multiply).
const fn apply_row(coefs: [Q4_28; 3], ch: [Q0_16; 3], addend_q428: Q4_28) -> Q0_16 {
    let [c0, c1, c2] = coefs;
    let [x, y, z] = ch;
    let acc = mul_acc(0, c0, x);
    let acc = mul_acc(acc, c1, y);
    let acc = mul_acc(acc, c2, z);
    let acc = mul_acc(acc, addend_q428, Q0_16::ONE);
    shift_round_sat_q0_16(acc)
}

/// `round_half_up(x^3 / 65535^2)`, entirely in [`u32`].
///
/// The obvious spelling — cube into [`u64`], divide by `65535^2` — links
/// `__aeabi_uldivmod` on a 32-bit target, one of the exact symbols the
/// `no-wide-div` gate step rejects. Dividing by `65535` three times instead
/// keeps every value in [`u32`] and every divide in the `arith` reciprocal
/// helper.
///
/// Writing `p = x*x = q1*65535 + r1` and `r1*x = q2*65535 + r2` gives
/// `x^3 = (q1*x + q2)*65535 + r2`, and one more step gives
/// `x^3 = q3*65535^2 + r3*65535 + r2`. Each intermediate is proven to fit:
/// `x*x <= 65535^2`, `r1*x <= 65534*65535`, and `q1*x + q2 <= 4294901759`.
/// Half-up rounding on `r3*65535 + r2` reduces to `r3 >= 32768`, or
/// `r3 == 32767` with `r2 >= 32768`.
const fn cube_q0_16(x: Q0_16) -> Q0_16 {
    let x = x.to_raw() as u32;
    let p = x.saturating_mul(x);
    let q1 = div_65535(p);
    let r1 = rem_65535(p, q1);
    let n2 = (r1 as u32).saturating_mul(x);
    let q2 = div_65535(n2);
    let r2 = rem_65535(n2, q2);
    let a = q1.saturating_mul(x).saturating_add(q2);
    let q3 = div_65535(a);
    let r3 = rem_65535(a, q3);
    let up = (r3 >= 32768) | (r3 == 32767 && r2 >= 32768);
    let rounded = q3.saturating_add(up as u32);
    if rounded >= u16::MAX as u32 {
        Q0_16::ONE
    } else {
        Q0_16::from_raw(rounded as u16)
    }
}

/// Convert linear sRGB to Oklab. Encoded sRGB is a type error.
///
/// ```compile_fail
/// let encoded = ph_color::Color::<ph_color::Srgb, ph_color::Encoded>::new(
///     ph_color::Q0_16::array_from_raw([0, 0, 0]),
/// );
/// let _ = ph_color::srgb_to_oklab(encoded);
/// ```
#[must_use]
pub const fn srgb_to_oklab(color: Color<Srgb, Linear>) -> Color<Oklab, Linear> {
    let [row0, row1, row2] = OKLAB_M1;
    let lms = [
        apply_row(row0, color.ch, Q4_28::ZERO),
        apply_row(row1, color.ch, Q4_28::ZERO),
        apply_row(row2, color.ch, Q4_28::ZERO),
    ];
    let [l, m, s] = lms;
    let lp = [CBRT.lookup(l), CBRT.lookup(m), CBRT.lookup(s)];
    let [lrow, arow, brow] = OKLAB_M2;
    Color::new([
        apply_row(lrow, lp, Q4_28::ZERO),
        apply_row(arow, lp, A_B_BIAS_Q428),
        apply_row(brow, lp, A_B_BIAS_Q428),
    ])
}

/// Convert Oklab to linear sRGB.
#[must_use]
pub const fn oklab_to_srgb(color: Color<Oklab, Linear>) -> Color<Srgb, Linear> {
    let [row0, row1, row2] = OKLAB_M2_INV;
    let [off0, off1, off2] = OKLAB_M2_INV_OFFSET;
    let lp = [
        apply_row(row0, color.ch, off0),
        apply_row(row1, color.ch, off1),
        apply_row(row2, color.ch, off2),
    ];
    let [l, m, s] = lp;
    let lms = [cube_q0_16(l), cube_q0_16(m), cube_q0_16(s)];
    let [rrow, grow, brow] = OKLAB_M1_INV;
    Color::new([
        apply_row(rrow, lms, Q4_28::ZERO),
        apply_row(grow, lms, Q4_28::ZERO),
        apply_row(brow, lms, Q4_28::ZERO),
    ])
}

/// Maximum |sRGB - round-trip| in UQ0.16 LSBs on the asserted sample set:
/// the 9-step cube, every axis sample `0..=65535`, `[620, 0, 0]`,
/// `[580, 0, 340]`, and two-channel planes with the third channel 0 at
/// step 64 (including both endpoints). Measured maximum on that set is 327 LSB.
pub const OKLAB_ROUNDTRIP_MAX_LSB: u16 = 327;

/// Maximum |Oklab - quantized Ottosson host reference| in
/// UQ0.16 LSBs on the same asserted sample set as [`OKLAB_ROUNDTRIP_MAX_LSB`].
/// Measured maximum on that set is 6252 LSB. Host-renderer visual parity is
/// not claimed.
pub const OKLAB_FORWARD_MAX_LSB: u16 = 6252;

#[cfg(test)]
mod tests {
    use super::*;

    const GRID: [u16; 9] = [0, 8192, 16384, 24576, 32768, 40960, 49152, 57344, 65535];

    #[test]
    fn cbrt_companion_and_endpoints() {
        assert_eq!(CBRT.max_err_lsb(), CBRT_MAX_ERR_LSB);
        assert_eq!(CBRT.lookup(Q0_16::ZERO), Q0_16::ZERO);
        assert_eq!(CBRT.lookup(Q0_16::ONE), Q0_16::ONE);
        assert_eq!(cube_q0_16(Q0_16::ZERO), Q0_16::ZERO);
        assert_eq!(cube_q0_16(Q0_16::ONE), Q0_16::ONE);
    }

    #[test]
    fn cube_matches_the_u64_reference_on_every_input() {
        // The three-step u32 form replaced a u64 divide; assert it is
        // bit-identical to that reference across the whole domain.
        let den = u64::from(u16::MAX) * u64::from(u16::MAX);
        for x in 0..=u16::MAX {
            let n = u64::from(x) * u64::from(x) * u64::from(x);
            let (quot, rem) = (n / den, n % den);
            let want = if rem * 2 >= den { quot + 1 } else { quot };
            let want = if want >= u64::from(u16::MAX) {
                Q0_16::ONE
            } else {
                Q0_16::from_raw(want as u16)
            };
            assert_eq!(cube_q0_16(Q0_16::from_raw(x)), want, "x={x}");
        }
    }

    #[test]
    fn black_and_white_round_trip() {
        let black = Color::<Srgb, Linear>::new(Q0_16::array_from_raw([0, 0, 0]));
        let white = Color::<Srgb, Linear>::new(Q0_16::array_from_raw([65535, 65535, 65535]));
        assert_eq!(
            oklab_to_srgb(srgb_to_oklab(black)).ch,
            Q0_16::array_from_raw([0, 0, 0])
        );
        let [r, g, b] = oklab_to_srgb(srgb_to_oklab(white)).ch;
        assert!(r.to_raw() > 65000 && g.to_raw() > 65000 && b.to_raw() > 65000);
    }

    #[test]
    fn round_trip_stays_within_asserted_bound() {
        let mut max = 0u16;
        for x in GRID {
            for y in GRID {
                for z in GRID {
                    max = max.max(round_trip_err([x, y, z]));
                }
            }
        }
        for x in 0..=u16::MAX {
            max = max.max(round_trip_err([x, 0, 0]));
            max = max.max(round_trip_err([0, x, 0]));
            max = max.max(round_trip_err([0, 0, x]));
        }
        max = max.max(round_trip_err([620, 0, 0]));
        max = max.max(round_trip_err([580, 0, 340]));
        let mut i = 0u16;
        loop {
            let x = i.saturating_mul(64);
            let mut j = 0u16;
            loop {
                let y = j.saturating_mul(64);
                max = max.max(round_trip_err([x, y, 0]));
                max = max.max(round_trip_err([x, 0, y]));
                max = max.max(round_trip_err([0, x, y]));
                if j == 1024 {
                    break;
                }
                j = j.saturating_add(1);
            }
            if i == 1024 {
                break;
            }
            i = i.saturating_add(1);
        }
        assert!(
            max <= OKLAB_ROUNDTRIP_MAX_LSB,
            "round-trip max {max} > bound {OKLAB_ROUNDTRIP_MAX_LSB}"
        );
        assert!(max > 0, "bound should be tight enough to be non-vacuous");
    }

    fn round_trip_err(ch: [u16; 3]) -> u16 {
        let src = Color::<Srgb, Linear>::new(Q0_16::array_from_raw(ch));
        let back = oklab_to_srgb(srgb_to_oklab(src));
        let [a0, a1, a2] = src.ch;
        let [b0, b1, b2] = back.ch;
        a0.to_raw()
            .abs_diff(b0.to_raw())
            .max(a1.to_raw().abs_diff(b1.to_raw()))
            .max(a2.to_raw().abs_diff(b2.to_raw()))
    }
}