improvise/crates/numtower/src/sum.rs

//! Order-free, retractable sums (clause 8): ADR 009's f64 `ExactSum`, and `NumSum`,
//! which adds the exact partitions. The design: `context/code/numtower.md`, "Sums".

use std::collections::BTreeMap;

use num_bigint::{BigInt, BigUint};
use num_integer::Integer;
use num_traits::{One, Zero};

use crate::fixed::{I256, I320};
use crate::num::{MAX_SCALE, Num, NumRef, pow10};
use crate::ratio::{BigQ, ratio_to_f64};

/// The most non-decimal denominator groups a sum reads exactly (question 16: K = 8).
pub const K: usize = 8;

const LIMB_BITS: u32 = 32;
/// A finite f64 is `mant * 2^(shift - 1074)`, `shift <= 2045`; weights add 60 bits.
const LIMBS: usize = 70;
const CHUNK: u64 = 1 << 30;

/// An exact sum of weighted f64s, rounded once when read (ADR 009).
#[derive(Clone, Debug)]
pub struct ExactSum {
    limbs: [i64; LIMBS],
    pending: u64,
    nan: i128,
    pos_inf: i128,
    neg_inf: i128,
    /// Net weight of the `-0.0` addends, and of every other finite one: a zero sum is
    /// `-0.0` only when the first is positive and the second is nil (IEEE).
    neg_zeros: i128,
    others: i128,
}

/// Equal sums, not equal representations: carries propagated, `pending` ignored.
impl PartialEq for ExactSum {
    fn eq(&self, o: &ExactSum) -> bool {
        let (mut mine, mut theirs) = (self.limbs, o.limbs);
        propagate(&mut mine);
        propagate(&mut theirs);
        let counts = |s: &ExactSum| (s.nan, s.pos_inf, s.neg_inf);
        let zero = mine.iter().all(|l| *l == 0);
        let same_zero = !zero || self.negative_zero() == o.negative_zero();
        mine == theirs && counts(self) == counts(o) && same_zero
    }
}

impl Eq for ExactSum {}

impl Default for ExactSum {
    fn default() -> Self {
        ExactSum {
            limbs: [0; LIMBS],
            pending: 0,
            nan: 0,
            pos_inf: 0,
            neg_inf: 0,
            neg_zeros: 0,
            others: 0,
        }
    }
}

impl ExactSum {
    pub fn new() -> ExactSum {
        ExactSum::default()
    }

    /// Adds `x` with a Z-set `weight`; the weight is split into base-2^30 digits.
    pub fn add(&mut self, x: f64, weight: i64) {
        match (x.is_nan(), x.is_infinite(), x > 0.0) {
            (true, _, _) => self.nan += weight as i128,
            (false, true, true) => self.pos_inf += weight as i128,
            (false, true, false) => self.neg_inf += weight as i128,
            (false, false, _) => {
                match x == 0.0 && x.is_sign_negative() {
                    true => self.neg_zeros += weight as i128,
                    false => self.others += weight as i128,
                }
                let magnitude = weight.unsigned_abs();
                for j in 0..3 {
                    let digit = (magnitude >> (30 * j)) & (CHUNK - 1);
                    if digit != 0 {
                        let signed = if weight < 0 {
                            -(digit as i64)
                        } else {
                            digit as i64
                        };
                        self.add_digit(x, signed, 30 * j);
                    }
                }
            }
        }
    }

    fn add_digit(&mut self, x: f64, weight: i64, extra: u32) {
        if self.pending + weight.unsigned_abs() > CHUNK {
            propagate(&mut self.limbs);
            self.pending = 0;
        }
        let bits = x.to_bits();
        let biased = ((bits >> 52) & 0x7ff) as u32;
        let fraction = bits & ((1 << 52) - 1);
        let (mant, shift) = match biased {
            0 => (fraction, 0),
            e => (fraction | (1 << 52), e - 1),
        };
        let sign = if x.is_sign_negative() {
            -weight
        } else {
            weight
        };
        let at = shift + extra;
        let wide = u128::from(mant) << (at % LIMB_BITS);
        let k = (at / LIMB_BITS) as usize;
        for i in 0..3 {
            let part = ((wide >> (LIMB_BITS * i as u32)) & 0xffff_ffff) as i64;
            self.limbs[k + i] += sign * part;
        }
        self.pending += weight.unsigned_abs();
    }

    /// Adds another accumulator's contents.
    pub fn absorb(&mut self, o: &ExactSum) {
        let mut theirs = o.limbs;
        propagate(&mut theirs);
        propagate(&mut self.limbs);
        for (mine, t) in self.limbs.iter_mut().zip(theirs) {
            *mine += t;
        }
        propagate(&mut self.limbs);
        self.pending = 0;
        self.nan += o.nan;
        self.pos_inf += o.pos_inf;
        self.neg_inf += o.neg_inf;
        self.neg_zeros += o.neg_zeros;
        self.others += o.others;
    }

    fn negative_zero(&self) -> bool {
        self.neg_zeros > 0 && self.others == 0
    }

    /// The exact sum rounded to the nearest f64, ties to even.
    pub fn value(&self) -> f64 {
        match (self.nan != 0, self.pos_inf != 0, self.neg_inf != 0) {
            (true, _, _) | (false, true, true) => f64::NAN,
            (false, true, false) => f64::INFINITY,
            (false, false, true) => f64::NEG_INFINITY,
            (false, false, false) => match (self.negative_zero(), self.total()) {
                (true, total) if total.is_zero() => -0.0,
                (_, total) => ratio_to_f64(&total, &(BigInt::one() << 1074)),
            },
        }
    }

    /// The exact sum over `count`, rounded once, so the mean neither overflows nor rounds twice.
    /// Non-finite sums and a zero's sign read as `value` does. Panics on a zero `count`.
    pub fn mean(&self, count: u64) -> f64 {
        assert!(count > 0, "the mean of nothing");
        match (self.nan != 0, self.pos_inf != 0, self.neg_inf != 0) {
            (false, false, false) => match (self.negative_zero(), self.total()) {
                (true, total) if total.is_zero() => -0.0,
                (_, total) => {
                    let unit = BigInt::one() << 1074;
                    ratio_to_f64(&total, &(unit * BigInt::from(count)))
                }
            },
            _ => self.value(),
        }
    }

    /// The finite part as an integer in units of 2^-1074.
    fn total(&self) -> BigInt {
        let mut limbs = self.limbs;
        propagate(&mut limbs);
        let digits: Vec<u32> = limbs[..LIMBS - 1].iter().map(|&l| l as u32).collect();
        let low = BigInt::from(BigUint::from_slice(&digits));
        (BigInt::from(limbs[LIMBS - 1]) << (LIMB_BITS as usize * (LIMBS - 1))) + low
    }
}

/// Carry each limb into the next, leaving every limb but the top in [0, 2^32).
fn propagate(limbs: &mut [i64; LIMBS]) {
    for k in 0..LIMBS - 1 {
        let carry = limbs[k] >> LIMB_BITS;
        limbs[k] -= carry << LIMB_BITS;
        limbs[k + 1] += carry;
    }
}

/// A `SUM` over any `Num`s: exact when every term is exact and at most K groups remain.
#[derive(Clone, Debug, Default)]
pub struct NumSum {
    decimals: I320,
    groups: BTreeMap<i128, I256>,
    group_floats: ExactSum,
    inexact: ExactSum,
    /// Net weight per f64 bit pattern: a read is inexact while any term is present.
    inexact_terms: BTreeMap<u64, i128>,
}

impl NumSum {
    pub fn new() -> NumSum {
        NumSum::default()
    }

    pub fn insert(&mut self, x: &Num) {
        self.add(x, 1);
    }

    pub fn retract(&mut self, x: &Num) {
        self.add(x, -1);
    }

    /// Adds `x` with a Z-set `weight`. Exact and associative: the order cannot matter.
    pub fn add(&mut self, x: &Num, weight: i64) {
        match x.kind() {
            NumRef::Inexact(f) => {
                self.inexact.add(f, weight);
                let net = self.inexact_terms.entry(f.to_bits()).or_insert(0);
                *net += weight as i128;
                if *net == 0 {
                    self.inexact_terms.remove(&f.to_bits());
                }
            }
            NumRef::Exact(d) => {
                let unit = pow10(MAX_SCALE - d.scale()) as u128;
                self.add_scaled(d.mantissa(), unit, weight);
            }
            NumRef::Fraction(r) => match at_scale_28(*r.denom()) {
                Some(unit) => self.add_scaled(*r.numer(), unit, weight),
                None => self.add_to_group(*r.numer(), *r.denom(), weight),
            },
        }
    }

    fn add_scaled(&mut self, n: i128, unit: u128, weight: i64) {
        let magnitude = weight.unsigned_abs() as u128;
        let scaled = match unit.checked_mul(magnitude) {
            Some(both) => I320::from_i128(n).checked_mul_u128(both),
            None => I320::from_i128(n)
                .checked_mul_u128(unit)
                .and_then(|t| t.checked_mul_u128(magnitude)),
        };
        let term = scaled
            .and_then(|t| if weight < 0 { t.checked_neg() } else { Some(t) })
            .expect("a weighted term is below 2^284");
        self.decimals = self
            .decimals
            .checked_add(&term)
            .expect("320 bits holds 2^36 terms");
    }

    fn add_to_group(&mut self, n: i128, d: i128, weight: i64) {
        let old = self.groups.get(&d).copied().unwrap_or(I256::ZERO);
        let new = old
            .checked_add(&I256::mul_i128(n, weight as i128))
            .expect("256 bits holds 2^64 terms");
        let value = |num: I256| ratio_to_f64(&num.to_bigint(), &BigInt::from(d));
        if !old.is_zero() {
            self.group_floats.add(value(old), -1);
        }
        match new.is_zero() {
            true => {
                self.groups.remove(&d);
            }
            false => {
                self.group_floats.add(value(new), 1);
                self.groups.insert(d, new);
            }
        }
    }

    /// The current non-decimal denominator groups.
    pub fn groups(&self) -> usize {
        self.groups.len()
    }

    /// One read: exact over at most K groups, else the f64 accumulators (clause 8).
    pub fn read(&self) -> Num {
        match (!self.inexact_terms.is_empty(), self.groups.len() > K) {
            (false, false) => Num::settle(&self.exact_total()),
            (_, true) => {
                let mut acc = self.inexact.clone();
                acc.absorb(&self.group_floats);
                acc.add(self.decimal_total().to_f64(), 1);
                Num::inexact(acc.value())
            }
            (true, false) => {
                let mut acc = self.inexact.clone();
                acc.add(self.exact_total().to_f64(), 1);
                Num::inexact(acc.value())
            }
        }
    }

    fn decimal_total(&self) -> BigQ {
        BigQ::new(self.decimals.to_bigint(), BigInt::from(pow10(MAX_SCALE)))
    }

    fn exact_total(&self) -> BigQ {
        let decimal = (!self.decimals.is_zero())
            .then(|| (self.decimals.to_bigint(), BigInt::from(pow10(MAX_SCALE))));
        let terms: Vec<(BigInt, BigInt)> = decimal
            .into_iter()
            .chain(
                self.groups
                    .iter()
                    .map(|(d, n)| (n.to_bigint(), BigInt::from(*d))),
            )
            .collect();
        let (num, den) = combine_over_lcm(&terms);
        BigQ::new(num, den)
    }
}

/// `Σ n/d` over the lcm of the denominators, not yet reduced: the read reduces once.
pub(crate) fn combine_over_lcm(terms: &[(BigInt, BigInt)]) -> (BigInt, BigInt) {
    let lcm = terms.iter().fold(BigInt::one(), |l, (_, d)| l.lcm(d));
    let num = terms
        .iter()
        .fold(BigInt::zero(), |s, (n, d)| s + n * (&lcm / d));
    (num, lcm)
}

/// `10^28 / d` when `d = 2^p 5^q` with `max(p, q) <= 28`: the term joins the decimal partition.
fn at_scale_28(d: i128) -> Option<u128> {
    let p = d.trailing_zeros();
    let mut rest = d >> p;
    let mut q = 0;
    while rest % 5 == 0 {
        rest /= 5;
        q += 1;
    }
    (rest == 1 && p.max(q) <= MAX_SCALE).then(|| (pow10(MAX_SCALE) / d) as u128)
}

#[cfg(test)]
mod tests {
    use super::*;
    use num_rational::BigRational;
    use proptest::prelude::*;
    use proptest::test_runner::{Config, RngAlgorithm, TestRng, TestRunner};

    type Combine = fn(&[(BigInt, BigInt)]) -> (BigInt, BigInt);

    /// Up to K groups of wide numerators whose denominators share factors.
    fn groups() -> impl Strategy<Value = Vec<(BigInt, BigInt)>> {
        let shared = prop::sample::select(vec![1i128, 3, 7 * 11, 1 << 20, 3i128.pow(40)]);
        let group =
            (any::<i128>(), any::<i64>(), shared, 1i128..(1 << 60)).prop_map(|(n, w, s, d)| {
                (
                    BigInt::from(n) * BigInt::from(w),
                    BigInt::from(s) * BigInt::from(d),
                )
            });
        prop::collection::vec(group, 1..=K)
    }

    fn check_grouped_read(combine: Combine, deterministic: bool) -> Result<(), String> {
        let config = Config {
            cases: 256,
            failure_persistence: None,
            ..Config::default()
        };
        let mut runner = match deterministic {
            true => {
                TestRunner::new_with_rng(config, TestRng::deterministic_rng(RngAlgorithm::ChaCha))
            }
            false => TestRunner::new(config),
        };
        let result = runner.run(&groups(), |terms| {
            let (num, den) = combine(&terms);
            let lcm = terms.iter().fold(BigInt::one(), |l, (_, d)| l.lcm(d));
            prop_assert_eq!(&den, &lcm, "the read did not combine over the lcm");
            let oracle: BigRational = terms
                .iter()
                .map(|(n, d)| BigRational::new(n.clone(), d.clone()))
                .sum();
            let read = BigQ::new(num, den);
            let got = BigRational::new(read.num.clone(), read.den.clone());
            prop_assert_eq!(got, oracle, "the read is not the sum");
            prop_assert!(read.num.gcd(&read.den).is_one(), "the read is not reduced");
            Ok(())
        });
        result.map_err(|e| e.to_string())
    }

    #[test]
    fn the_grouped_read_combines_over_the_lcm_and_reduces_once() {
        check_grouped_read(combine_over_lcm, false).unwrap();
    }

    fn over_the_product(terms: &[(BigInt, BigInt)]) -> (BigInt, BigInt) {
        let den = terms.iter().fold(BigInt::one(), |p, (_, d)| p * d);
        let num = terms
            .iter()
            .fold(BigInt::zero(), |s, (n, d)| s + n * (&den / d));
        (num, den)
    }

    #[test]
    fn the_grouped_read_combines_over_the_lcm_and_reduces_once_catches_the_product() {
        let err = check_grouped_read(over_the_product, true).expect_err("the rival passed");
        assert!(
            err.contains("the read did not combine over the lcm"),
            "{err}"
        );
    }

    #[test]
    fn a_fraction_with_a_decimal_denominator_joins_the_decimal_partition() {
        let mut s = NumSum::new();
        s.insert(&Num::ratio(10i128.pow(35), 1).unwrap());
        s.insert(&Num::ratio(3, 1 << 20).unwrap());
        assert_eq!(s.groups(), 0);
        s.insert(&Num::ratio(1, 1i128 << 40).unwrap());
        s.insert(&Num::ratio(1, 3).unwrap());
        assert_eq!(
            s.groups(),
            2,
            "2^40 has scale 40, past 28: a group of its own"
        );
        s.retract(&Num::ratio(1, 3).unwrap());
        assert_eq!(
            s.groups(),
            1,
            "a group whose numerator returns to zero is dropped"
        );
    }
}