mod.rs

mod kernel;
use core::cmp::Ordering;
use core::f32;
use core::num::FpCategory;

/// Explicitly stored significand bits in [`prim@f32`]
///
/// This constant is usually used as a shift to access the exponent bits.
pub const EXP_SHIFT: u32 = f32::MANTISSA_DIGITS - 1;

/// Magnitude of `f32`
///
/// Nonzero subnormal numbers are normalized to have an implicit leading bit.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Magnitude {
    /// NaN, see [`FpCategory::Nan`]
    Nan,

    /// Infinity, see [`FpCategory::Infinite`]
    Infinite,

    /// Zero, see [`FpCategory::Zero`]
    ///
    /// Zero cannot be normalized.  A normalized magnitude has an implicit
    /// leading bit.
    Zero,

    /// Normalized magnitude
    ///
    /// The layout of the bits is the same as a normal positive `f32`.  For
    /// subnormal numbers, the stored exponent becomes zero or negative while
    /// the significand is normalized to have an implicit leading bit.
    Normalized(i32),
}

/// Break a `f32` into its sign and magnitude
#[allow(clippy::cast_possible_wrap)]
#[inline]
fn normalize(x: f32) -> (bool, Magnitude) {
    let sign = x.is_sign_negative();
    let magnitude = x.abs().to_bits() as i32;

    match x.classify() {
        FpCategory::Nan => (sign, Magnitude::Nan),
        FpCategory::Infinite => (sign, Magnitude::Infinite),
        FpCategory::Zero => (sign, Magnitude::Zero),
        FpCategory::Normal => (sign, Magnitude::Normalized(magnitude)),
        FpCategory::Subnormal => {
            let shift = magnitude.leading_zeros() as i32 - 8;
            let magnitude = (magnitude << shift) - (shift << EXP_SHIFT);
            (sign, Magnitude::Normalized(magnitude))
        }
    }
}

/// Fast multiply-add
///
/// This function picks the faster way to compute `x * y + a` depending on the
/// target architecture.  The FMA instruction is used if available.  Otherwise,
/// it falls back to `x * y + a` that is faster but gives less accurate results
/// than [`f32::mul_add`]
#[must_use]
#[inline]
pub fn mul_add(x: f32, y: f32, a: f32) -> f32 {
    #[cfg(target_feature = "fma")]
    return x.mul_add(y, a);

    #[cfg(not(target_feature = "fma"))]
    return x * y + a;
}

/// Polynomial evaluation with Horner's method
///
/// This function evaluates a polynomial with coefficients in `p` at `x`.
/// This function calls [`mul_add`] for simplicity.
#[must_use]
#[inline]
pub fn poly(x: f32, p: &[f32]) -> f32 {
    p.iter()
        .copied()
        .rev()
        .reduce(|y, c| mul_add(y, x, c))
        .unwrap_or_default()
}

/// The least number greater than `x`
///
/// This is a less careful version of [`f32::next_up`] regarding subnormal
/// numbers.  This function is useful until `f32::next_up` gets stabilized.
#[must_use]
#[inline]
pub fn next_up(x: f32) -> f32 {
    if x.is_nan() || x == f32::INFINITY {
        x
    } else if x == 0.0 {
        f32::from_bits(1)
    } else if x.is_sign_negative() {
        f32::from_bits(x.to_bits() - 1)
    } else {
        f32::from_bits(x.to_bits() + 1)
    }
}

/// The greatest number less than `x`
///
/// This is a less careful version of [`f32::next_down`] regarding subnormal
/// numbers.  This function is useful until `f32::next_down` gets stabilized.
#[must_use]
#[inline]
pub fn next_down(x: f32) -> f32 {
    if x.is_nan() || x == f32::NEG_INFINITY {
        x
    } else if x == 0.0 {
        f32::from_bits(0x8000_0001)
    } else if x.is_sign_negative() {
        f32::from_bits(x.to_bits() + 1)
    } else {
        f32::from_bits(x.to_bits() - 1)
    }
}

/// Rounds half-way cases away from zero
#[must_use]
#[inline]
pub fn round(x: f32) -> f32 {
    let r = x.abs();
    let i = r.trunc();

    (i + f32::from(r - i >= 0.5)).copysign(x)
}

/// The cube root
#[must_use]
#[inline]
pub fn cbrt(x: f32) -> f32 {
    let (sign, Magnitude::Normalized(magnitude)) = normalize(x) else {
        return x;
    };

    // SAFETY: the minimum magnitude is -0x0B00_0000, so the cast integer is
    // always positive
    #[allow(clippy::cast_sign_loss)]
    let magnitude = (0x2A51_2CE3 + magnitude / 3) as u32;

    let iter = |y: f32| mul_add(3.0_f32.recip(), x / (y * y) - y, y);
    iter(iter(iter(f32::from_bits(
        u32::from(sign) << 31 | magnitude,
    ))))
}

/// Hypotenuse of a right-angled triangle with sides `x` and `y`
#[must_use]
#[inline]
pub fn hypot(x: f32, y: f32) -> f32 {
    if x.is_infinite() || y.is_infinite() {
        return f32::INFINITY;
    }

    let x: f64 = x.into();
    let y: f64 = y.into();

    #[allow(clippy::imprecise_flops, clippy::cast_possible_truncation)]
    return (x * x + y * y).sqrt() as f32;
}

/// The exponential function
#[must_use]
#[inline]
pub fn exp(x: f32) -> f32 {
    use core::f32::consts::LN_2;
    use core::f64::consts;

    #[allow(clippy::cast_precision_loss, clippy::cast_possible_wrap)]
    if x < (f32::MIN_EXP - f32::MANTISSA_DIGITS as i32 - 1) as f32 * LN_2 {
        return 0.0;
    }

    #[allow(clippy::cast_precision_loss)]
    if x > f32::MAX_EXP as f32 * LN_2 {
        return f32::INFINITY;
    }

    let x = f64::from(x);
    let n = (x * consts::LOG2_E).round_ties_even();
    let x = crate::f64::mul_add(n, -consts::LN_2, x);
    let y = crate::f64::mul_add(kernel::exp_slope(x), x, 1.0);

    #[allow(clippy::cast_possible_truncation)]
    return kernel::fast_ldexp(y, n as i64) as f32;
}

/// Raise 2 to the power of `x`
#[must_use]
#[inline]
pub fn exp2(x: f32) -> f32 {
    #[allow(clippy::cast_precision_loss, clippy::cast_possible_wrap)]
    if x < (f32::MIN_EXP - f32::MANTISSA_DIGITS as i32 - 1) as f32 {
        return 0.0;
    }

    #[allow(clippy::cast_precision_loss)]
    if x > f32::MAX_EXP as f32 {
        return f32::INFINITY;
    }

    let n = x.round_ties_even();
    let x = poly(
        x - n,
        &[
            1.0,
            6.931_472e-1,
            2.402_265e-1,
            5.550_357e-2,
            9.618_031e-3,
            1.339_086_7e-3,
            1.546_973_5e-4,
        ],
    );

    #[allow(clippy::cast_possible_truncation)]
    return kernel::fast_ldexp(x.into(), n as i64) as f32;
}

/// Raise 10 to the power of `x`
#[must_use]
#[inline]
pub fn exp10(x: f32) -> f32 {
    use core::f32::consts::LOG10_2;

    #[allow(clippy::cast_precision_loss, clippy::cast_possible_wrap)]
    if x < (f32::MIN_EXP - f32::MANTISSA_DIGITS as i32 - 1) as f32 * LOG10_2 {
        return 0.0;
    }

    #[allow(clippy::cast_precision_loss)]
    if x > f32::MAX_EXP as f32 * LOG10_2 {
        return f32::INFINITY;
    }

    let x = f64::from(x) * core::f64::consts::LOG2_10;
    let n = x.round_ties_even();
    let x = crate::f64::poly(
        x - n,
        &[
            9.999_999_999_190_67e-1,
            6.931_472_067_096_466e-1,
            2.402_265_150_505_521e-1,
            5.550_327_215_766_594e-2,
            9.617_994_514_161_479e-3,
            1.340_043_166_700_788_1e-3,
            1.547_802_227_945_78e-4,
        ],
    );

    #[allow(clippy::cast_possible_truncation)]
    return kernel::fast_ldexp(x, n as i64) as f32;
}

/// Compute `exp(x) - 1` accurately especially for small `x`
#[must_use]
#[inline]
pub fn exp_m1(x: f32) -> f32 {
    use core::f32::consts::LN_2;
    use core::f64::consts;

    #[allow(clippy::cast_precision_loss)]
    if x < f32::MANTISSA_DIGITS as f32 * -LN_2 {
        return -1.0;
    }

    #[allow(clippy::cast_precision_loss)]
    if x > f32::MAX_EXP as f32 * LN_2 {
        return f32::INFINITY;
    }

    let x = f64::from(x);
    let n = (x * consts::LOG2_E).round_ties_even() + 0.0;
    let x = crate::f64::mul_add(n, -consts::LN_2, x);
    let y = kernel::exp_slope(x);

    if n == 0.0 {
        #[allow(clippy::cast_possible_truncation)]
        return (x * y) as f32;
    }

    #[allow(clippy::cast_possible_truncation)]
    return (kernel::fast_ldexp(crate::f64::mul_add(x, y, 1.0), n as i64) - 1.0) as f32;
}

/// Multiply `x` by 2 raised to the power of `n`
#[must_use]
#[inline]
pub fn ldexp(x: f32, n: i32) -> f32 {
    const MIN_EXP: i32 = f64::MIN_EXP - 1;
    const MAX_EXP: i32 = f64::MAX_EXP;

    let coefficient = match n {
        ..MIN_EXP => 0.5 * f64::MIN_POSITIVE,

        #[allow(clippy::cast_sign_loss)]
        n @ MIN_EXP..MAX_EXP => f64::from_bits(((MAX_EXP - 1 + n) as u64) << crate::f64::EXP_SHIFT),

        MAX_EXP.. => f64::MAX,
    };

    #[allow(clippy::cast_possible_truncation)]
    return (f64::from(x) * coefficient) as f32;
}

/// Decompose into a significand and an exponent
///
/// The absolute value of the significand is in the range of [0.5, 1) for
/// nonzero finite `x` for historical reasons.  This function also explains how
/// [`f32::MAX_EXP`] and [`f32::MIN_EXP`] are defined.
#[must_use]
#[inline]
pub fn frexp(x: f32) -> (f32, i32) {
    let (sign, Magnitude::Normalized(magnitude)) = normalize(x) else {
        return (x, 0);
    };

    let mask = f32::MIN_POSITIVE.to_bits() - 1;

    #[allow(clippy::cast_sign_loss)]
    let significand = magnitude as u32 & mask | 0.5f32.to_bits();

    (
        f32::from_bits(u32::from(sign) << 31 | significand),
        f32::MIN_EXP - 1 + (magnitude >> EXP_SHIFT),
    )
}

/// Natural logarithm
#[must_use]
#[inline]
pub fn ln(x: f32) -> f32 {
    match normalize(x) {
        (false, Magnitude::Infinite) => f32::INFINITY,
        (_, Magnitude::Zero) => f32::NEG_INFINITY,
        (true, _) | (_, Magnitude::Nan) => f32::NAN,

        (false, Magnitude::Normalized(i)) => {
            use core::f32::consts::FRAC_1_SQRT_2;

            #[allow(clippy::cast_possible_wrap)]
            let exponent = (i - FRAC_1_SQRT_2.to_bits() as i32) >> EXP_SHIFT;

            #[allow(clippy::cast_sign_loss)]
            let x = f64::from(f32::from_bits((i - (exponent << EXP_SHIFT)) as u32));

            #[allow(clippy::cast_possible_truncation)]
            return crate::f64::mul_add(
                core::f64::consts::LN_2,
                exponent.into(),
                2.0 * kernel::atanh((x - 1.0) / (x + 1.0)),
            ) as f32;
        }
    }
}

/// Compute `ln(1 + x)` accurately especially for small `x`
#[must_use]
#[inline]
pub fn ln_1p(x: f32) -> f32 {
    match x {
        f32::INFINITY => f32::INFINITY,
        -1.0 => f32::NEG_INFINITY,
        x if x < -1.0 || x.is_nan() => f32::NAN,

        // SAFETY: there is no wrap because `1.0 + x` is positive
        #[allow(clippy::cast_possible_wrap)]
        _ => {
            use core::f64::consts::FRAC_1_SQRT_2;
            let x = f64::from(x);
            let i = (1.0 + x).to_bits() as i64;
            let exponent = (i - FRAC_1_SQRT_2.to_bits() as i64) >> crate::f64::EXP_SHIFT;

            #[allow(clippy::cast_possible_truncation)]
            return (if exponent == 0 {
                2.0 * kernel::atanh(x / (2.0 + x))
            } else {
                #[allow(clippy::cast_sign_loss)]
                let y = f64::from_bits((i - (exponent << crate::f64::EXP_SHIFT)) as u64);

                #[allow(clippy::cast_precision_loss)]
                crate::f64::mul_add(
                    core::f64::consts::LN_2,
                    exponent as f64,
                    2.0 * kernel::atanh((y - 1.0) / (y + 1.0)),
                )
            }) as f32;
        }
    }
}

/// Base 2 logarithm
#[must_use]
#[inline]
pub fn log2(x: f32) -> f32 {
    match normalize(x) {
        (false, Magnitude::Infinite) => f32::INFINITY,
        (_, Magnitude::Zero) => f32::NEG_INFINITY,
        (true, _) | (_, Magnitude::Nan) => f32::NAN,

        (false, Magnitude::Normalized(i)) => {
            use core::f32::consts::FRAC_1_SQRT_2;

            #[allow(clippy::cast_possible_wrap)]
            let exponent = (i - FRAC_1_SQRT_2.to_bits() as i32) >> EXP_SHIFT;

            #[allow(clippy::cast_sign_loss)]
            let x = f64::from(f32::from_bits((i - (exponent << EXP_SHIFT)) as u32));

            #[allow(clippy::cast_possible_truncation)]
            return crate::f64::mul_add(
                2.0 * core::f64::consts::LOG2_E,
                kernel::atanh((x - 1.0) / (x + 1.0)),
                exponent.into(),
            ) as f32;
        }
    }
}

/// Base 10 logarithm
#[must_use]
#[inline]
pub fn log10(x: f32) -> f32 {
    match normalize(x) {
        (false, Magnitude::Infinite) => f32::INFINITY,
        (_, Magnitude::Zero) => f32::NEG_INFINITY,
        (true, _) | (_, Magnitude::Nan) => f32::NAN,

        (false, Magnitude::Normalized(i)) => {
            use core::f32::consts::FRAC_1_SQRT_2;
            use core::f64::consts;

            #[allow(clippy::cast_possible_wrap)]
            let exponent = (i - FRAC_1_SQRT_2.to_bits() as i32) >> EXP_SHIFT;

            #[allow(clippy::cast_sign_loss)]
            let x = f64::from(f32::from_bits((i - (exponent << EXP_SHIFT)) as u32));

            #[allow(clippy::cast_possible_truncation)]
            return crate::f64::mul_add(
                2.0 * consts::LOG10_E,
                kernel::atanh((x - 1.0) / (x + 1.0)),
                consts::LOG10_2 * f64::from(exponent),
            ) as f32;
        }
    }
}

/// Logarithm with arbitrary base
#[must_use]
#[inline]
pub fn log(x: f32, base: f32) -> f32 {
    fn log2(x: f32) -> f64 {
        match (x.is_sign_negative(), x.classify()) {
            (false, FpCategory::Infinite) => f64::INFINITY,
            (_, FpCategory::Zero) => f64::NEG_INFINITY,
            (true, _) | (_, FpCategory::Nan) => f64::NAN,
            _ => kernel::log2(x.into()),
        }
    }
    #[allow(clippy::cast_possible_truncation)]
    return (log2(x) / log2(base)) as f32;
}

/// Raise to a floating-point power
#[must_use]
#[inline]
pub fn powf(x: f32, y: f32) -> f32 {
    fn magnitude(x: f32, y: f32) -> f32 {
        match x.classify() {
            FpCategory::Nan => f32::NAN,
            FpCategory::Infinite => match y.partial_cmp(&0.0) {
                Some(Ordering::Greater) => f32::INFINITY,
                Some(Ordering::Less) => 0.0,
                Some(Ordering::Equal) => 1.0,
                None => f32::NAN,
            },
            FpCategory::Zero => match y.partial_cmp(&0.0) {
                Some(Ordering::Greater) => 0.0,
                Some(Ordering::Less) => f32::INFINITY,
                Some(Ordering::Equal) => 1.0,
                None => f32::NAN,
            },
            _ => match x {
                1.0 => 1.0,
                x if x.is_sign_negative() => f32::NAN,

                #[allow(clippy::cast_possible_truncation)]
                _ => kernel::exp2(f64::from(y) * kernel::log2(x.into())) as f32,
            },
        }
    }

    fn is_integer(x: f32) -> bool {
        x.trunc().eq(&x)
    }

    if y == 0.0 {
        return 1.0;
    }

    if x.is_sign_negative() && is_integer(y) {
        let sign = if is_integer(0.5 * y) { 1.0 } else { -1.0 };
        return sign * magnitude(-x, y);
    }

    magnitude(x, y)
}

/// Inverse hyperbolic tangent
#[must_use]
#[inline]
pub fn atanh(x: f32) -> f32 {
    match x.abs().partial_cmp(&1.0) {
        #[allow(clippy::cast_possible_wrap)]
        Some(core::cmp::Ordering::Less) => {
            use crate::f64::EXP_SHIFT;
            use core::f64::consts;

            let x: f64 = x.into();
            let i = ((1.0 + x) / (1.0 - x)).to_bits() as i64;
            let exponent = (i - consts::FRAC_1_SQRT_2.to_bits() as i64) >> EXP_SHIFT;

            #[allow(clippy::cast_possible_truncation)]
            if exponent == 0 {
                return kernel::atanh(x) as f32;
            }

            #[allow(clippy::cast_sign_loss)]
            let x = f64::from_bits((i - (exponent << EXP_SHIFT)) as u64);

            #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
            return crate::f64::mul_add(
                0.5 * consts::LN_2,
                exponent as f64,
                kernel::atanh((x - 1.0) / (x + 1.0)),
            ) as f32;
        }
        Some(core::cmp::Ordering::Equal) => f32::INFINITY.copysign(x),
        _ => f32::NAN,
    }
}

/// Inverse hyperbolic sine
#[must_use]
#[inline]
pub fn asinh(x: f32) -> f32 {
    #[allow(clippy::cast_possible_wrap)]
    fn magnitude(s: f64) -> f64 {
        use crate::f64::EXP_SHIFT;
        use core::f64::consts;

        let c = crate::f64::mul_add(s, s, 1.0).sqrt();
        let i = (c + s).to_bits() as i64;
        let exponent = (i - consts::FRAC_1_SQRT_2.to_bits() as i64) >> EXP_SHIFT;

        if exponent == 0 {
            return 2.0 * kernel::atanh(s / (c + 1.0));
        }

        #[allow(clippy::cast_sign_loss)]
        let x = f64::from_bits((i - (exponent << EXP_SHIFT)) as u64);

        #[allow(clippy::cast_precision_loss)]
        crate::f64::mul_add(
            consts::LN_2,
            exponent as f64,
            2.0 * kernel::atanh((x - 1.0) / (x + 1.0)),
        )
    }

    if !x.is_finite() {
        return x;
    };

    #[allow(clippy::cast_possible_truncation)]
    (magnitude(x.abs().into()) as f32).copysign(x)
}

/// Inverse hyperbolic cosine
#[must_use]
#[inline]
pub fn acosh(x: f32) -> f32 {
    match x {
        f32::INFINITY => f32::INFINITY,

        #[allow(clippy::cast_possible_wrap)]
        (1.0..) => {
            use crate::f64::EXP_SHIFT;
            use core::f64::consts;

            let c: f64 = x.into();
            let s = crate::f64::mul_add(c, c, -1.0).sqrt();
            let i = (c + s).to_bits() as i64;
            let exponent = (i - consts::FRAC_1_SQRT_2.to_bits() as i64) >> EXP_SHIFT;

            #[allow(clippy::cast_sign_loss)]
            let x = f64::from_bits((i - (exponent << EXP_SHIFT)) as u64);

            #[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
            return crate::f64::mul_add(
                consts::LN_2,
                exponent as f64,
                2.0 * kernel::atanh((x - 1.0) / (x + 1.0)),
            ) as f32;
        }

        _ => f32::NAN,
    }
}