path: root/src/eval.rs

use num::bigint::Sign;
use num::{BigInt, FromPrimitive, Integer, Signed, ToPrimitive, Zero};
use std::fmt::{Display, Formatter};

fn int_to_float(i: &BigInt) -> f64 {
    match i.sign() {
        Sign::Minus => return -int_to_float(&-i),
        Sign::NoSign => return 0.,
        _ => (),
    }
    let mut exponent = i.bits() - 1;
    if exponent > 1023 {
        return f64::INFINITY;
    }
    let mut fraction;
    if exponent <= 52 {
        // The number can be represented exactly without rounding.
        fraction = i.to_u64().unwrap() << (52 - exponent);
    } else {
        // We need to round.
        let power = BigInt::from(1) << (exponent - 52);
        let (rounded_fraction, rem) = i.div_rem(&power);
        fraction = rounded_fraction.to_u64().unwrap();
        let half = power >> 1;
        // Round up?
        if rem > half || rem == half && fraction % 2 == 1 {
            fraction += 1;
            // Check for overflow.
            if fraction == 1 << 53 {
                exponent += 1;
                fraction >>= 1;
                if exponent > 1023 {
                    return f64::INFINITY;
                }
            }
        }
    }

    // Build the float from exponent and fraction
    // https://en.wikipedia.org/wiki/IEEE_754
    f64::from_bits(((exponent + 1023) << 52) + (fraction & !(1 << 52)))
}

#[derive(Debug, PartialEq)]
pub enum Num {
    Int(BigInt),
    Float(f64),
}

fn int(i: BigInt) -> Num {
    if i.bits() > 1 << 30 {
        return Num::Float(int_to_float(&i));
    }
    Num::Int(i)
}

fn float(f: f64) -> Num {
    Num::Float(f)
}

impl Num {
    fn as_float(&self) -> f64 {
        match self {
            Num::Int(i) => int_to_float(i),
            &Num::Float(f) => f,
        }
    }

    fn square(self) -> Num {
        match self {
            Num::Int(i) => int(i.pow(2)),
            Num::Float(f) => float(f * f),
        }
    }

    fn powi(self, p: BigInt) -> Num {
        let mut n = self;
        let mut p = p;
        let mut acc = int(BigInt::from(1));
        loop {
            if p.is_odd() {
                acc = acc.mul(&n);
            }
            p >>= 1;
            if p.is_zero() {
                return acc;
            }
            n = n.square();
        }
    }

    pub fn pow(self, other: Num) -> Num {
        match other {
            Num::Int(i) => {
                if !i.is_negative() {
                    return self.powi(i);
                }
                if let Ok(i_32) = i32::try_from(&i) {
                    return float(self.as_float().powi(i_32));
                }
                float(1.).div(self).powi(-i)
            }
            Num::Float(f) => float(self.as_float().powf(f)),
        }
    }

    pub fn mul(self, other: &Num) -> Num {
        match (self, other) {
            (Num::Int(i1), Num::Int(i2)) => int(i1 * i2),
            (n1, n2) => float(n1.as_float() * n2.as_float()),
        }
    }

    pub fn div(self, other: Num) -> Num {
        float(self.as_float() / other.as_float())
    }

    pub fn int_div(self, other: Num) -> Num {
        match (self, other) {
            (Num::Int(i1), Num::Int(i2)) if !i2.is_zero() => int(i1.div_floor(&i2)),
            (n1, n2) => {
                let f = (n1.as_float() / n2.as_float()).floor();
                let Some(i) = BigInt::from_f64(f) else {
                    return float(f);
                };
                int(i)
            }
        }
    }

    pub fn modulo(self, other: Num) -> Num {
        match (self, other) {
            (Num::Int(i1), Num::Int(i2)) if !i2.is_zero() => int(i1.mod_floor(&i2)),
            (n1, n2) => float(n1.as_float() % n2.as_float()),
        }
    }

    pub fn add(self, other: Num) -> Num {
        match (self, other) {
            (Num::Int(i1), Num::Int(i2)) => int(i1 + i2),
            (n1, n2) => float(n1.as_float() + n2.as_float()),
        }
    }

    pub fn sub(self, other: Num) -> Num {
        match (self, other) {
            (Num::Int(i1), Num::Int(i2)) => int(i1 - i2),
            (n1, n2) => float(n1.as_float() - n2.as_float()),
        }
    }

    pub fn sin(self) -> Num {
        float(self.as_float().sin())
    }

    pub fn cos(self) -> Num {
        float(self.as_float().cos())
    }

    pub fn tan(self) -> Num {
        float(self.as_float().tan())
    }

    pub fn asin(self) -> Num {
        float(self.as_float().asin())
    }

    pub fn acos(self) -> Num {
        float(self.as_float().acos())
    }

    pub fn atan(self) -> Num {
        float(self.as_float().atan())
    }

    pub fn sqrt(self) -> Num {
        match self {
            Num::Int(i) if !i.is_negative() => {
                let s = i.sqrt();
                if &s * &s == i {
                    return int(s);
                }
                float(int_to_float(&i).sqrt())
            }
            n => float(n.as_float().sqrt()),
        }
    }

    pub fn log(self) -> Num {
        float(self.as_float().ln())
    }

    pub fn log10(self) -> Num {
        float(self.as_float().log10())
    }

    pub fn log2(self) -> Num {
        float(self.as_float().log2())
    }

    pub fn floor(self) -> Num {
        match self {
            Num::Int(i) => int(i),
            Num::Float(f) => {
                let f = f.floor();
                if let Some(i) = BigInt::from_f64(f) {
                    int(i)
                } else {
                    float(f)
                }
            }
        }
    }

    pub fn ceil(self) -> Num {
        match self {
            Num::Int(i) => int(i),
            Num::Float(f) => {
                let f = f.ceil();
                if let Some(i) = BigInt::from_f64(f) {
                    int(i)
                } else {
                    float(f)
                }
            }
        }
    }

    pub fn round(self) -> Num {
        match self {
            Num::Int(i) => int(i),
            Num::Float(f) => {
                let f = f.round_ties_even();
                if let Some(i) = BigInt::from_f64(f) {
                    int(i)
                } else {
                    float(f)
                }
            }
        }
    }

    pub fn abs(self) -> Num {
        match self {
            Num::Int(i) => int(i.abs()),
            Num::Float(f) => float(f.abs()),
        }
    }
}

impl Display for Num {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
        match self {
            Num::Int(n) if n.bits() <= 100 => write!(f, "{n}"),
            n => {
                let n = n.as_float();
                if (-(1i128 << 100) as f64) < n && n < (1i128 << 100) as f64 {
                    return write!(f, "{}", format!("{n:.7}").trim_end_matches('0'));
                }
                write!(f, "{n:.7e}")
            }
        }
    }
}

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

    fn int(i: i128) -> Num {
        super::int(BigInt::from(i))
    }

    #[test]
    fn int_to_float_small() {
        for n in -(1 << 16)..1 << 16 {
            assert_eq!(
                int_to_float(&BigInt::from(n)),
                n as f64,
                "int_to_float({}) is not equal to {}",
                n,
                n as f64
            );
        }
    }

    #[test]
    fn int_to_float_medium() {
        let n = BigInt::from(2).pow(1023) + BigInt::from(2).pow(1022);
        assert_eq!(int_to_float(&n), n.to_f64().unwrap());
    }

    #[test]
    fn int_to_float_round_up() {
        assert_eq!(
            int_to_float(&BigInt::from(0x7f_ffff_ffff_ffffu64)),
            0x80_0000_0000_0000u64 as f64
        );
    }

    #[test]
    fn int_to_float_round_up_easy() {
        assert_eq!(
            int_to_float(&BigInt::from(0x40_0000_0000_0003u64)),
            0x40_0000_0000_0004u64 as f64
        );
    }

    #[test]
    fn int_to_float_round_up_overflow() {
        assert_eq!(
            int_to_float(&BigInt::from(0x1ff_ffff_ffff_fff9u64)),
            0x200_0000_0000_0000u64 as f64
        );
    }

    #[test]
    fn int_to_float_round_down() {
        assert_eq!(
            int_to_float(&BigInt::from(0x40_0000_0000_0001u64)),
            0x40_0000_0000_0000u64 as f64
        );
    }

    #[test]
    fn int_to_float_round_to_even_down() {
        assert_eq!(
            int_to_float(&BigInt::from(0x40_0000_0000_0002u64)),
            0x40_0000_0000_0000u64 as f64
        );
    }

    #[test]
    fn int_to_float_round_to_even_up() {
        assert_eq!(
            int_to_float(&BigInt::from(0x40_0000_0000_0006u64)),
            0x40_0000_0000_0008u64 as f64
        );
    }

    #[test]
    fn int_to_float_large() {
        assert_eq!(
            int_to_float(
                &((BigInt::from(0x1f_ffff_ffff_ffffu64) << 971) + (BigInt::from(1) << 970) - 1)
            ),
            f64::MAX
        );
    }

    #[test]
    fn int_to_float_overflow() {
        assert_eq!(
            int_to_float(
                &((BigInt::from(0x1f_ffff_ffff_ffffu64) << 971) + (BigInt::from(1) << 970))
            ),
            f64::INFINITY
        );
    }

    #[test]
    fn pow_positive_int() {
        assert_eq!(int(2).pow(int(16)), int(65536));
    }

    #[test]
    fn pow_negative_int() {
        assert_eq!(int(2).pow(int(-3)), float(0.125));
    }

    #[test]
    fn pow_float() {
        assert_eq!(int(4).pow(float(0.5)), float(2.));
    }

    #[test]
    fn pow_overflow() {
        assert_eq!(int(2).pow(int(1 << 126)), float(f64::INFINITY));
    }

    #[test]
    fn pow_underflow() {
        assert_eq!(int(2).pow(int(-(1 << 126))), float(0.));
    }

    #[test]
    fn modulo_pos() {
        assert_eq!(int(5).modulo(int(3)), int(2));
    }

    #[test]
    fn modulo_pos_neg() {
        assert_eq!(int(5).modulo(int(-3)), int(-1));
    }

    #[test]
    fn modulo_neg_pos() {
        assert_eq!(int(-5).modulo(int(3)), int(1));
    }

    #[test]
    fn modulo_neg() {
        assert_eq!(int(-5).modulo(int(-3)), int(-2));
    }

    #[test]
    fn modulo_float() {
        assert_eq!(float(5.).modulo(int(3)), float(2.));
    }

    #[test]
    fn sqrt() {
        for n in 0..65536 {
            assert_eq!(
                int(n * n).sqrt(),
                int(n),
                "int({}).sqrt() is not equal to {}",
                n * n,
                n
            );
        }
    }

    #[test]
    fn sqrt_big() {
        assert_eq!(int(1 << 126).sqrt(), int(1 << 63));
    }
}