improvise/crates/improvise-incremental/src/group.rs

//! One group's maintained state and how it is read (ADR 009 clauses 4 and 8).
//! Reasoning: `context/code/improvise-incremental.md`, steps 4 and 6.

use crate::DeltaValue;
use improvise_core::formula::ast::AggFunc;
use improvise_core::model::cell::CellValue;
use numtower::ExactSum;
use std::collections::BTreeMap;

/// A group's cells: how many, how many are numbers, a multiset of their values and their exact
/// sum. An update touches one map entry, so it costs O(log d) in the distinct values.
#[derive(Debug, Clone, Default)]
pub struct Group {
    cells: i64,
    numbers: i64,
    values: BTreeMap<DeltaValue, i64>,
    sum: ExactSum,
    negative_zeros: i64,
}

impl Group {
    /// Adds `weight` copies of `value`; a negative weight retracts.
    pub fn apply(&mut self, value: &DeltaValue, weight: i64) {
        self.cells += weight;
        if matches!(value.value(), CellValue::Number(_)) {
            self.numbers += weight;
        }
        let held = self.values.get(value).copied().unwrap_or(0) + weight;
        match held {
            0 => self.values.remove(value),
            n => self.values.insert(value.clone(), n),
        };
    }

    /// Adds `weight` copies of the number `x` to the sum and the count of negative zeros, as stored:
    /// the multiset calls the two zeros equal, these must not.
    pub fn add_number(&mut self, x: f64, weight: i64) {
        self.sum.add(x, weight);
        if x == 0.0 && x.is_sign_negative() {
            self.negative_zeros += weight;
        }
    }

    pub fn cells(&self) -> i64 {
        self.cells
    }

    /// How many zeros the group holds, of either sign.
    fn zeros(&self) -> i64 {
        let zero = DeltaValue::new(&CellValue::Number(0.0));
        self.values.get(&zero).copied().unwrap_or(0)
    }

    pub fn is_empty(&self) -> bool {
        self.cells == 0
    }

    /// A group of exactly these values, for a fold over cells read one by one.
    pub fn of<'a>(values: impl IntoIterator<Item = &'a CellValue>) -> Self {
        let mut group = Group::default();
        for value in values {
            group.apply(&DeltaValue::new(value), 1);
            if let CellValue::Number(x) = value {
                group.add_number(*x, 1);
            }
        }
        group
    }

    /// The smallest number that is not NaN, if the group has one.
    fn smallest(&self) -> Option<f64> {
        self.values.keys().find_map(non_nan)
    }

    /// The largest number that is not NaN, if the group has one.
    fn largest(&self) -> Option<f64> {
        self.values.keys().rev().find_map(non_nan)
    }
}

fn non_nan(value: &DeltaValue) -> Option<f64> {
    match value.value() {
        CellValue::Number(n) if !n.is_nan() => Some(*n),
        _ => None,
    }
}

/// The aggregates a group answers, none of which looks at order.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Reduce {
    Count,
    Min,
    Max,
    Sum,
    Avg,
}

impl Reduce {
    pub fn of(func: &AggFunc) -> Self {
        match func {
            AggFunc::Count => Reduce::Count,
            AggFunc::Min => Reduce::Min,
            AggFunc::Max => Reduce::Max,
            AggFunc::Sum => Reduce::Sum,
            AggFunc::Avg => Reduce::Avg,
        }
    }
}

/// What eager's `aggregate_cell_values` returns for the union of `groups`: numbers reduce,
/// anything else is unanimous or `Error("mixed")`, and no cells is absent.
pub fn fold_groups(groups: &[&Group], func: Reduce) -> Option<CellValue> {
    let cells: i64 = groups.iter().map(|g| g.cells).sum();
    let numbers: i64 = groups.iter().map(|g| g.numbers).sum();
    match (cells, numbers == cells) {
        (0, _) => None,
        (_, true) => Some(CellValue::Number(reduce(groups, func, cells))),
        (_, false) => Some(unanimous(groups)),
    }
}

/// `Min` and `Max` skip NaN, and are infinite when nothing else is left, as eager's folds are.
fn reduce(groups: &[&Group], func: Reduce, cells: i64) -> f64 {
    match func {
        Reduce::Count => cells as f64,
        Reduce::Sum => total(groups).value(),
        Reduce::Avg => total(groups).value() / cells as f64,
        Reduce::Min => {
            let smallest = groups.iter().filter_map(|g| g.smallest());
            let min = smallest.fold(f64::INFINITY, f64::min);
            let negative: i64 = groups.iter().map(|g| g.negative_zeros).sum();
            signed_zero(min, negative > 0)
        }
        Reduce::Max => {
            let largest = groups.iter().filter_map(|g| g.largest());
            let max = largest.fold(f64::NEG_INFINITY, f64::max);
            let (zeros, negative): (i64, i64) = (
                groups.iter().map(|g| g.zeros()).sum(),
                groups.iter().map(|g| g.negative_zeros).sum(),
            );
            signed_zero(max, zeros > 0 && zeros == negative)
        }
    }
}

/// The groups' sums, absorbed into one: exact, so the same in any order and for any grouping.
fn total(groups: &[&Group]) -> ExactSum {
    let mut sum = ExactSum::new();
    groups.iter().for_each(|g| sum.absorb(&g.sum));
    sum
}

/// `x` with a zero given the sign `negative` says: `Min` is `-0.0` if any zero is, `Max` only if
/// all are. Why: `context/code/improvise-incremental.md`, step 4, item 10.
fn signed_zero(x: f64, negative: bool) -> f64 {
    match x == 0.0 && negative {
        true => -0.0,
        false => x,
    }
}

fn unanimous(groups: &[&Group]) -> CellValue {
    let mut distinct = groups.iter().flat_map(|g| g.values.keys());
    let first = distinct.next();
    match first.filter(|first| distinct.all(|v| v == *first)) {
        Some(only) => only.value().clone(),
        None => CellValue::Error("mixed".into()),
    }
}

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

    fn group(values: &[CellValue]) -> Group {
        Group::of(values)
    }

    fn num(n: f64) -> CellValue {
        CellValue::Number(n)
    }

    #[test]
    fn a_retraction_removes_the_value_it_added() {
        let mut g = group(&[num(1.0), num(2.0)]);
        g.apply(&DeltaValue::new(&num(1.0)), -1);
        assert_eq!(g.cells(), 1);
        assert_eq!(
            fold_groups(&[&g], Reduce::Min),
            Some(num(2.0)),
            "the retracted minimum is gone"
        );
    }

    #[test]
    fn an_emptied_group_is_absent_and_keeps_no_entries() {
        let mut g = group(&[num(1.0)]);
        g.apply(&DeltaValue::new(&num(1.0)), -1);
        assert!(g.is_empty());
        assert!(g.values.is_empty(), "no zero-weight entry is left behind");
        assert_eq!(fold_groups(&[&g], Reduce::Count), None);
    }

    #[test]
    fn min_and_max_skip_nan_and_are_infinite_over_nothing_else() {
        let some = group(&[num(f64::NAN), num(3.0), num(-1.0)]);
        assert_eq!(fold_groups(&[&some], Reduce::Min), Some(num(-1.0)));
        assert_eq!(fold_groups(&[&some], Reduce::Max), Some(num(3.0)));
        let only_nan = group(&[num(f64::NAN), num(-f64::NAN)]);
        assert_eq!(
            fold_groups(&[&only_nan], Reduce::Min),
            Some(num(f64::INFINITY))
        );
        assert_eq!(
            fold_groups(&[&only_nan], Reduce::Max),
            Some(num(f64::NEG_INFINITY))
        );
        assert_eq!(fold_groups(&[&only_nan], Reduce::Count), Some(num(2.0)));
    }

    #[test]
    fn a_fold_over_several_groups_is_the_fold_of_their_union() {
        let (a, b) = (group(&[num(5.0), num(9.0)]), group(&[num(2.0)]));
        assert_eq!(fold_groups(&[&a, &b], Reduce::Min), Some(num(2.0)));
        assert_eq!(fold_groups(&[&a, &b], Reduce::Max), Some(num(9.0)));
        assert_eq!(fold_groups(&[&a, &b], Reduce::Count), Some(num(3.0)));
    }

    #[test]
    fn text_is_unanimous_or_mixed_across_groups() {
        let text = |s: &str| CellValue::Text(s.into());
        let (a, b) = (group(&[text("x"), text("x")]), group(&[text("x")]));
        assert_eq!(fold_groups(&[&a, &b], Reduce::Min), Some(text("x")));
        let c = group(&[text("y")]);
        let mixed = Some(CellValue::Error("mixed".into()));
        assert_eq!(fold_groups(&[&a, &c], Reduce::Min), mixed);
        let with_number = group(&[text("x"), num(1.0)]);
        assert_eq!(fold_groups(&[&with_number], Reduce::Count), mixed);
        let nan_among_text = group(&[num(f64::NAN), text("x")]);
        assert_eq!(fold_groups(&[&nan_among_text], Reduce::Count), mixed);
    }
}