alice-det-math 0.5.0

//! `f32x8` versions of the `f32` kernels (feature `simd`).
//!
//! Every function here is lane-for-lane **bit-identical** to its scalar
//! counterpart in the crate root: the same constants, the same operation
//! order, no `mul_add` (which would fuse on FMA targets and not on others),
//! and every branch of the scalar code replaced by a mask + blend over
//! values that were computed the same way. `tests/simd_parity.rs` pins this
//! over dense sweeps and the special values.
//!
//! The kernels are written with `f32` arithmetic, compares + `blend`, lane
//! shifts and `f32 ↔ i32` conversions only. Where the scalar code
//! manipulates bits, the SIMD code uses an *exactly equivalent* float
//! operation: scaling by a power of two instead of masking the exponent,
//! `-0.0 - x` instead of flipping the sign bit, `kf - 4·floor(kf/4)` instead
//! of `k & 3`. (`wide` 0.7 has no NEON path for integer / bitwise lane ops,
//! which run 5–10× slower than the float ops there.)
//!
//! [`sin`], [`cos`], [`sin_cos`], [`exp`], [`ln`], [`round`] and [`sqrt`]
//! are vectorised. The functions whose scalar version evaluates an `f64`
//! kernel ([`atan2`], [`atan`], [`asin`], [`acos`], [`tan`], [`tanh`],
//! [`powf`], [`cbrt`], [`hypot`]) are evaluated per lane with the scalar
//! code; they are provided so callers have one namespace for both shapes.

use bytemuck::cast;
use wide::{f32x8, i32x8, CmpEq, CmpGe, CmpGt, CmpLe, CmpLt};

use crate::single::{
    COS_P, EXP_HI, EXP_LO, EXP_P, FRAC_2_PI, LG1, LG2, LG3, LG4, LN2_HI, LN2_LO, LOG2E, PIO2_1,
    PIO2_2, PIO2_3, PIO2_4, SIN_P, SPLIT_MASK, SQRT2_BITS,
};

/// Constant vectors as `const` items: they live in rodata and load with one
/// instruction. Building them at run time (`f32x8::splat`, or even
/// `f32x8::new([v; 8])` in an expression) lowers to a `memset_pattern16`
/// *call* per constant on aarch64 with wide 0.7 — that was the whole cost of
/// these kernels (27 calls per `sin_cos`).
const TOINT: f32 = 8_388_608.0; // 2^23
const C_NEG0P0: f32x8 = f32x8::new([-0.0; 8]);
const C_NEG0P5: f32x8 = f32x8::new([-0.5; 8]);
const C_NEG126P0: f32x8 = f32x8::new([-126.0; 8]);
const C_NEG25P0: f32x8 = f32x8::new([-25.0; 8]);
const C_NEG2147483648P0: f32x8 = f32x8::new([-2_147_483_648.0; 8]);
const C_0P0: f32x8 = f32x8::new([0.0; 8]);
const C_0P25: f32x8 = f32x8::new([0.25; 8]);
const C_0P5: f32x8 = f32x8::new([0.5; 8]);
const C_1P0: f32x8 = f32x8::new([1.0; 8]);
const C_127P0: f32x8 = f32x8::new([127.0; 8]);
const C_2P0: f32x8 = f32x8::new([2.0; 8]);
const C_2147483648P0: f32x8 = f32x8::new([2_147_483_648.0; 8]);
const C_3P0: f32x8 = f32x8::new([3.0; 8]);
const C_33554432P0: f32x8 = f32x8::new([33_554_432.0; 8]);
const C_4P0: f32x8 = f32x8::new([4.0; 8]);
const C_COSP0: f32x8 = f32x8::new([COS_P[0]; 8]);
const C_COSP1: f32x8 = f32x8::new([COS_P[1]; 8]);
const C_COSP2: f32x8 = f32x8::new([COS_P[2]; 8]);
const C_EXPHI: f32x8 = f32x8::new([EXP_HI; 8]);
const C_EXPLO: f32x8 = f32x8::new([EXP_LO; 8]);
const C_EXPP0: f32x8 = f32x8::new([EXP_P[0]; 8]);
const C_EXPP1: f32x8 = f32x8::new([EXP_P[1]; 8]);
const C_EXPP2: f32x8 = f32x8::new([EXP_P[2]; 8]);
const C_EXPP3: f32x8 = f32x8::new([EXP_P[3]; 8]);
const C_EXPP4: f32x8 = f32x8::new([EXP_P[4]; 8]);
const C_FRAC2PI: f32x8 = f32x8::new([FRAC_2_PI; 8]);
const C_LG1: f32x8 = f32x8::new([LG1; 8]);
const C_LG2: f32x8 = f32x8::new([LG2; 8]);
const C_LG3: f32x8 = f32x8::new([LG3; 8]);
const C_LG4: f32x8 = f32x8::new([LG4; 8]);
const C_LN2HI: f32x8 = f32x8::new([LN2_HI; 8]);
const C_LN2LO: f32x8 = f32x8::new([LN2_LO; 8]);
const C_LOG2E: f32x8 = f32x8::new([LOG2E; 8]);
const C_PIO21: f32x8 = f32x8::new([PIO2_1; 8]);
const C_PIO22: f32x8 = f32x8::new([PIO2_2; 8]);
const C_PIO23: f32x8 = f32x8::new([PIO2_3; 8]);
const C_PIO24: f32x8 = f32x8::new([PIO2_4; 8]);
const C_SINP0: f32x8 = f32x8::new([SIN_P[0]; 8]);
const C_SINP1: f32x8 = f32x8::new([SIN_P[1]; 8]);
const C_SINP2: f32x8 = f32x8::new([SIN_P[2]; 8]);
const C_TOINT: f32x8 = f32x8::new([TOINT; 8]);
const C_INFINITY: f32x8 = f32x8::new([f32::INFINITY; 8]);
const C_MAX: f32x8 = f32x8::new([f32::MAX; 8]);
const C_MINPOSITIVE: f32x8 = f32x8::new([f32::MIN_POSITIVE; 8]);
const C_NAN: f32x8 = f32x8::new([f32::NAN; 8]);
const C_NEGINFINITY: f32x8 = f32x8::new([f32::NEG_INFINITY; 8]);
const C_SQRT2M: f32x8 = f32x8::new([f32::from_bits(SQRT2_BITS); 8]);

/// Exact negation for every input (`-0.0 - x`): `+0.0 → -0.0`, NaN stays
/// NaN. (`wide`'s `Neg` is `0.0 - x`, which turns `+0.0` into `+0.0`.)
#[inline(always)]
fn neg(x: f32x8) -> f32x8 {
    C_NEG0P0 - x
}

/// `2^k` from a float that holds a small integer `k` (`|k| ≤ 127`): the
/// biased exponent `k + 127` is formed in float (exact) and shifted into place.
#[inline(always)]
fn pow2(k: f32x8) -> f32x8 {
    cast((k + C_127P0).trunc_int() << 23i32)
}

/// `single::nearest_int` per lane: `trunc(t + (t < 0 ? -0.5 : 0.5))`.
#[inline(always)]
fn nearest_int(t: f32x8) -> f32x8 {
    let half = t.cmp_lt(C_0P0).blend(C_NEG0P5, C_0P5);
    f32x8::from_i32x8((t + half).trunc_int())
}

/// `f32::round` per lane — nearest integer, ties away from zero (musl
/// `roundf`, same formula as [`crate::round`]).
#[inline]
#[must_use]
pub fn round(x: f32x8) -> f32x8 {
    let ax = x.abs();
    // e >= 0x7f + 23  ⇔  |x| >= 2^23 (or inf / NaN)
    let integral = !ax.cmp_lt(C_TOINT);
    // e < 0x7f - 1  ⇔  |x| < 0.5
    let small = ax.cmp_lt(C_0P5);
    let neg_x = x.cmp_lt(C_0P0);
    let y0 = ax + C_TOINT - C_TOINT - ax;
    let y_ax = y0 + ax;
    let y_gt = y_ax - C_1P0;
    let y_le = y_ax + C_1P0;
    let y = y0
        .cmp_gt(C_0P5)
        .blend(y_gt, y0.cmp_le(C_NEG0P5).blend(y_le, y_ax));
    // y ≥ 1 here; the scalar `if neg { -y }` is the sign blend
    let y = neg_x.blend(neg(y), y);
    // |x| < 0.5 → ±0 with the sign of x (`bits & 0x8000_0000`): `x * 0.0`
    // keeps the sign of x, including for -0.0 (which is not `< 0`)
    let signed_zero = x * C_0P0;
    // a NaN lane (any payload) → the canonical NaN, as the scalar
    canonical_nan(integral.blend(x, small.blend(signed_zero, y)))
}

/// `f32::sqrt` per lane, correctly rounded; canonical NaN for negative / NaN
/// lanes (the hardware default NaN's sign bit differs between `x86` and
/// `AArch64`, so it is replaced here as in [`crate::sqrt`]).
#[inline]
#[must_use]
pub fn sqrt(x: f32x8) -> f32x8 {
    let nan = C_NAN;
    let v = x.cmp_lt(C_0P0).blend(nan, x.sqrt());
    x.cmp_eq(x).blend(v, nan)
}

#[inline(always)]
fn sin_poly(r: f32x8) -> f32x8 {
    let z = r * r;
    let p = ((C_SINP0 * z + C_SINP1) * z + C_SINP2) * z * r;
    r + p
}

#[inline(always)]
fn cos_poly(r: f32x8) -> f32x8 {
    let z = r * r;
    let p = ((C_COSP0 * z + C_COSP1) * z + C_COSP2) * z * z;
    C_1P0 - C_0P5 * z + p
}

/// `single::split_hi` per lane.
#[inline(always)]
fn split_hi(a: f32x8) -> f32x8 {
    let bits: i32x8 = cast(a);
    cast(bits & i32x8::splat(SPLIT_MASK as i32))
}

/// `single::two_prod` per lane, the same operations in the same order.
#[inline(always)]
fn two_prod(a: f32x8, b: f32x8) -> (f32x8, f32x8) {
    let p = a * b;
    let (ah, bh) = (split_hi(a), split_hi(b));
    let (al, bl) = (a - ah, b - bh);
    (p, (((ah * bh - p) + ah * bl) + al * bh) + al * bl)
}

/// `(k mod 4 as a float, r)` with `x = k·π/2 + r`.
///
/// `k mod 4` is `kf - 4·floor(kf/4)` (exact: `kf` is integral, and a
/// multiple of 4 once it exceeds 2^25), with the scalar's `kf as i32`
/// saturation reproduced: `kf ≥ 2^31 → 3`, `kf ≤ -2^31 → 0`, NaN → 0.
#[inline(always)]
fn reduce_pio2(x: f32x8) -> (f32x8, f32x8) {
    let kf = nearest_int(x * C_FRAC2PI);
    let (p3, e3) = two_prod(kf, C_PIO23);
    let r = ((((x - kf * C_PIO21) - kf * C_PIO22) - p3) - e3) - kf * C_PIO24;
    let q = kf - C_4P0 * (kf * C_0P25).floor();
    let q = kf.cmp_ge(C_2147483648P0).blend(C_3P0, q);
    let q = kf.cmp_le(C_NEG2147483648P0).blend(C_0P0, q);
    (q, r)
}

/// Quadrant selection: `k & 1` swaps sin / cos, `k & 2` negates sin,
/// `(k + 1) & 2` negates cos — as blends on `q = k mod 4 ∈ {0, 1, 2, 3}`.
/// Pure selection and exact sign flips: bit-identical to the scalar `match`.
#[inline(always)]
fn quadrant(q: f32x8, s: f32x8, c: f32x8) -> (f32x8, f32x8) {
    let q1 = q.cmp_eq(C_1P0);
    let q2 = q.cmp_eq(C_2P0);
    let q3 = q.cmp_eq(C_3P0);
    let swap = q1 | q3;
    let sv = swap.blend(c, s);
    let cv = swap.blend(s, c);
    let sv = (q2 | q3).blend(neg(sv), sv);
    let cv = (q1 | q2).blend(neg(cv), cv);
    (sv, cv)
}

/// Scalar `!x.is_finite() → NaN` plus `single::canon` (a NaN produced by the
/// arithmetic on a huge argument gets the platform's default NaN sign).
#[inline(always)]
fn canon(x: f32x8, v: f32x8) -> f32x8 {
    let nan = C_NAN;
    let finite = x.abs().cmp_le(C_MAX);
    let v = v.cmp_eq(v).blend(v, nan);
    finite.blend(v, nan)
}

/// Scalar `if x == 0.0 { return x }` of the odd functions: `±0` lanes return
/// themselves (the polynomial's `r + p` turns `−0` into `+0`).
#[inline(always)]
fn keep_zero(x: f32x8, v: f32x8) -> f32x8 {
    x.cmp_eq(C_0P0).blend(x, v)
}

/// Deterministic `sin` per lane (`NaN` for non-finite lanes).
#[inline]
#[must_use]
pub fn sin(x: f32x8) -> f32x8 {
    let (q, r) = reduce_pio2(x);
    let (sv, _) = quadrant(q, sin_poly(r), cos_poly(r));
    keep_zero(x, canon(x, sv))
}

/// Deterministic `cos` per lane (`NaN` for non-finite lanes).
#[inline]
#[must_use]
pub fn cos(x: f32x8) -> f32x8 {
    let (q, r) = reduce_pio2(x);
    let (_, cv) = quadrant(q, sin_poly(r), cos_poly(r));
    canon(x, cv)
}

/// Deterministic `(sin, cos)` per lane from one range reduction.
#[inline]
#[must_use]
pub fn sin_cos(x: f32x8) -> (f32x8, f32x8) {
    let (q, r) = reduce_pio2(x);
    let (sv, cv) = quadrant(q, sin_poly(r), cos_poly(r));
    (keep_zero(x, canon(x, sv)), canon(x, cv))
}

/// Deterministic `exp` per lane (`+inf` above ≈ 88.72, `0` below ≈ −103.97,
/// gradual subnormal tail, `NaN` stays `NaN`).
#[inline]
#[must_use]
pub fn exp(x: f32x8) -> f32x8 {
    let over = x.cmp_gt(C_EXPHI);
    let under = x.cmp_lt(C_EXPLO);
    // Clamp so the lanes that are blended out below never feed an
    // out-of-range exponent into `pow2`; in-range lanes are unchanged.
    let xc = over.blend(C_EXPHI, under.blend(C_EXPLO, x));
    let kf = nearest_int(xc * C_LOG2E);
    let r = (xc - kf * C_LN2HI) - kf * C_LN2LO;
    let p = (((((C_EXPP0 * r + C_EXPP1) * r + C_EXPP2) * r + C_EXPP3) * r + C_EXPP4) * r + C_0P5)
        * r
        * r
        + r
        + C_1P0;
    // Scalar: k > 127 → p·2^127·2^(k−127); k < −126 → p·2^−126·2^(k+126);
    // else p·2^k. Multiplying by 2^0 = 1.0 is exact, so one formula covers all.
    let k_hi = kf.max(C_NEG126P0).min(C_127P0);
    let k_lo = kf - k_hi;
    let v = p * pow2(k_hi) * pow2(k_lo);
    let v = under.blend(C_0P0, v);
    let v = over.blend(C_INFINITY, v);
    x.cmp_eq(x).blend(v, C_NAN)
}

/// Deterministic natural logarithm per lane (`ln(0) = -inf`,
/// `ln(x < 0) = NaN`, `ln(inf) = inf`).
#[inline]
#[must_use]
pub fn ln(x: f32x8) -> f32x8 {
    // Normalise subnormals so the exponent extraction below is exact
    // (scalar: `bits < 0x0080_0000 → x·2^25, k = -25`).
    let sub = x.cmp_lt(C_MINPOSITIVE);
    let xs = sub.blend(x * C_33554432P0, x);
    let k = sub.blend(C_NEG25P0, C_0P0);
    // biased exponent of a positive normal number: `bits >> 23`
    let bits: i32x8 = cast(xs);
    let e = f32x8::from_i32x8(bits >> 23i32) - C_127P0;
    let k = k + e;
    // m = mantissa in [1, 2): scaling by 2^-e is exact, the same bits as
    // `(bits & 0x007f_ffff) | 0x3f80_0000`. Done as `2^(1-e)` then `·0.5`
    // because `2^-127` is not a normal number (`e = 127` for x ≥ 2^127).
    let m = (xs * pow2(C_1P0 - e)) * C_0P5;
    // Reduce mantissa to [sqrt(2)/2, sqrt(2)).
    let ge = m.cmp_ge(C_SQRT2M);
    let m = ge.blend(m * C_0P5, m);
    let kf = ge.blend(k + C_1P0, k);
    let f = m - C_1P0;
    let s = f / (C_2P0 + f);
    let z = s * s;
    let w = z * z;
    let t1 = w * (C_LG2 + w * C_LG4);
    let t2 = z * (C_LG1 + w * C_LG3);
    let r = t2 + t1;
    let hfsq = C_0P5 * f * f;
    let v = kf * C_LN2HI - ((hfsq - (s * (hfsq + r) + kf * C_LN2LO)) - f);
    let v = x.cmp_eq(C_INFINITY).blend(C_INFINITY, v);
    let v = x.cmp_eq(C_0P0).blend(C_NEGINFINITY, v);
    let v = x.cmp_lt(C_0P0).blend(C_NAN, v);
    x.cmp_eq(x).blend(v, C_NAN)
}

#[inline(always)]
fn map(x: f32x8, f: impl Fn(f32) -> f32) -> f32x8 {
    let a = x.to_array();
    f32x8::new([
        f(a[0]),
        f(a[1]),
        f(a[2]),
        f(a[3]),
        f(a[4]),
        f(a[5]),
        f(a[6]),
        f(a[7]),
    ])
}

#[inline(always)]
fn map2(x: f32x8, y: f32x8, f: impl Fn(f32, f32) -> f32) -> f32x8 {
    let a = x.to_array();
    let b = y.to_array();
    f32x8::new([
        f(a[0], b[0]),
        f(a[1], b[1]),
        f(a[2], b[2]),
        f(a[3], b[3]),
        f(a[4], b[4]),
        f(a[5], b[5]),
        f(a[6], b[6]),
        f(a[7], b[7]),
    ])
}

/// [`crate::atan2`] per lane (scalar `f64` kernel).
#[inline]
#[must_use]
pub fn atan2(y: f32x8, x: f32x8) -> f32x8 {
    map2(y, x, crate::atan2)
}

/// [`crate::atan`] per lane (scalar `f64` kernel).
#[inline]
#[must_use]
pub fn atan(x: f32x8) -> f32x8 {
    map(x, crate::atan)
}

/// [`crate::asin`] per lane (scalar `f64` kernel).
#[inline]
#[must_use]
pub fn asin(x: f32x8) -> f32x8 {
    map(x, crate::asin)
}

/// [`crate::acos`] per lane (scalar `f64` kernel).
#[inline]
#[must_use]
pub fn acos(x: f32x8) -> f32x8 {
    map(x, crate::acos)
}

/// [`crate::tan`] per lane (scalar `f64` kernel).
#[inline]
#[must_use]
pub fn tan(x: f32x8) -> f32x8 {
    map(x, crate::tan)
}

/// [`crate::tanh`] per lane (scalar `f64` kernel).
#[inline]
#[must_use]
pub fn tanh(x: f32x8) -> f32x8 {
    map(x, crate::tanh)
}

/// [`crate::powf`] per lane (scalar `f64` kernel).
#[inline]
#[must_use]
pub fn powf(x: f32x8, y: f32x8) -> f32x8 {
    map2(x, y, crate::powf)
}

/// [`crate::cbrt`] per lane.
#[inline]
#[must_use]
pub fn cbrt(x: f32x8) -> f32x8 {
    map(x, crate::cbrt)
}

/// [`crate::hypot`] per lane.
#[inline]
#[must_use]
pub fn hypot(x: f32x8, y: f32x8) -> f32x8 {
    map2(x, y, crate::hypot)
}

/// Lane-wise counterpart of `metric::canonical`: a NaN raised by arithmetic
/// has a target-dependent sign, so it is replaced by the canonical one.
#[inline]
fn canonical_nan(v: f32x8) -> f32x8 {
    v.cmp_eq(v).blend(v, C_NAN)
}

/// [`crate::metric::lerp`] per lane — `(1 − t)·a + t·b`, exact at both ends.
#[inline]
#[must_use]
pub fn lerp(a: f32x8, b: f32x8, t: f32x8) -> f32x8 {
    canonical_nan((C_1P0 - t) * a + t * b)
}

/// [`crate::metric::clamp`] per lane.
///
/// A NaN in any argument gives the canonical NaN, as the scalar early return
/// does (the comparisons alone are false against NaN and would keep a NaN
/// `x` but ignore a NaN bound).
#[inline]
#[must_use]
pub fn clamp(x: f32x8, low: f32x8, high: f32x8) -> f32x8 {
    let y = high.cmp_lt(x).blend(high, x);
    let y = y.cmp_lt(low).blend(low, y);
    any_nan3(x, low, high).blend(C_NAN, y)
}

/// Lanes where any of the three arguments is NaN.
#[inline(always)]
fn any_nan3(a: f32x8, b: f32x8, c: f32x8) -> f32x8 {
    !(a.cmp_eq(a) & b.cmp_eq(b) & c.cmp_eq(c))
}

/// [`crate::metric::smoothstep`] per lane — `3t² − 2t³` over the clamped
/// interval, a step where the interval is degenerate, a NaN in any argument
/// → the canonical NaN.
#[inline]
#[must_use]
pub fn smoothstep(edge0: f32x8, edge1: f32x8, x: f32x8) -> f32x8 {
    let degenerate = !edge1.cmp_gt(edge0);
    let step = x.cmp_lt(edge0).blend(C_0P0, C_1P0);
    // the division is garbage in degenerate lanes and blended away there
    let t = clamp((x - edge0) / (edge1 - edge0), C_0P0, C_1P0);
    let cubic = canonical_nan(t * t * (C_3P0 - C_2P0 * t));
    let r = degenerate.blend(step, cubic);
    any_nan3(edge0, edge1, x).blend(C_NAN, r)
}

/// [`crate::metric::MetricWeights`] broadcast to lanes.
///
/// Splatting is a call rather than an instruction on `aarch64` with `wide`
/// 0.7, so a hot loop builds this once outside the loop instead of passing
/// the scalar weights per sample.
#[derive(Debug, Clone, Copy)]
pub struct MetricWeightsX8 {
    l1: f32x8,
    l2: f32x8,
    linf: f32x8,
}

impl MetricWeightsX8 {
    /// Broadcasts the three weights to all lanes.
    #[inline]
    #[must_use]
    pub fn splat(w: crate::metric::MetricWeights) -> Self {
        let (l1, l2, linf) = w.weights();
        Self {
            l1: f32x8::splat(l1),
            l2: f32x8::splat(l2),
            linf: f32x8::splat(linf),
        }
    }
}

/// [`crate::metric::MetricWeights::norm`] per lane, same operation order.
///
/// `‖·‖∞` is a mask + blend rather than `f32x8::max`, whose NaN behaviour
/// differs between the SSE and NEON backends; the blend reproduces the
/// scalar `if a > b { a } else { b }` on every target.
#[inline]
#[must_use]
pub fn metric_norm(w: MetricWeightsX8, x: f32x8, y: f32x8, z: f32x8) -> f32x8 {
    let ax = x.abs();
    let ay = y.abs();
    let az = z.abs();
    let l1 = ax + ay + az;
    let l2 = sqrt(x * x + y * y + z * z);
    let m = ax.cmp_gt(ay).blend(ax, ay);
    let linf = m.cmp_gt(az).blend(m, az);
    // the blends drop a NaN in the first two positions, but a NaN component
    // also makes `l1` and `l2` NaN, and `0 · NaN` is NaN: the sum is NaN for
    // every metric, as the scalar `norm` (whose `norm_linf` rejects it outright)
    canonical_nan(w.l1 * l1 + w.l2 * l2 + w.linf * linf)
}