improvise/crates/zset/src/lib.rs

//! Z-sets: maps from keys to non-zero `i64` weights (ADR 009 clause 1). An abelian group
//! wherever no weight leaves `i64`; past that it panics. See `context/code/zset.md`.

mod circuit;
mod edge;
mod ops;
mod trace;

pub use circuit::{Circuit, Input, Output, Progress, Rebuild, Stream};
pub use edge::{
    ChannelPipe, CodecError, Correction, Delta, Direct, EdgeCodec, EdgeSink, EdgeSource, Flow,
    Frame, Pipe, Recv, Refused, Request, Stamp, channel_pipes,
};
pub use ops::{AggregateState, DistinctState, JoinState};
pub use trace::{MemTrace, Trace};

use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
use std::iter::FromIterator;
use std::ops::{Add, AddAssign, Mul, Neg, Sub, SubAssign};

/// A map from `K` to a non-zero `i64` weight. A zero weight is never stored, so `==`
/// is the group's equality; iteration is in ascending key order.
#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct ZSet<K> {
    weights: BTreeMap<K, i64>,
}

#[track_caller]
fn overflow(op: &str, a: i64, b: i64) -> ! {
    panic!("zset: weight overflow in {op}: {a} and {b} do not combine within i64")
}

fn plus(a: i64, b: i64) -> i64 {
    a.checked_add(b).unwrap_or_else(|| overflow("add", a, b))
}

fn minus(a: i64, b: i64) -> i64 {
    a.checked_sub(b).unwrap_or_else(|| overflow("sub", a, b))
}

fn negated(a: i64) -> i64 {
    a.checked_neg().unwrap_or_else(|| overflow("neg", a, 0))
}

fn times(a: i64, n: i64) -> i64 {
    a.checked_mul(n).unwrap_or_else(|| overflow("scale", a, n))
}

impl<K> Default for ZSet<K> {
    fn default() -> Self {
        ZSet {
            weights: BTreeMap::new(),
        }
    }
}

impl<K: Ord> ZSet<K> {
    /// The zero of the group: no key has a weight.
    pub fn new() -> Self {
        Self::default()
    }

    /// One key at `weight`; a zero weight gives the zero Z-set.
    pub fn singleton(key: K, weight: i64) -> Self {
        let mut z = Self::new();
        z.insert(key, weight);
        z
    }

    /// The weight of `key`, zero if it is absent.
    pub fn weight(&self, key: &K) -> i64 {
        self.weights.get(key).copied().unwrap_or(0)
    }

    pub fn is_zero(&self) -> bool {
        self.weights.is_empty()
    }

    pub fn is_empty(&self) -> bool {
        self.is_zero()
    }

    /// The number of keys with a non-zero weight.
    pub fn len(&self) -> usize {
        self.weights.len()
    }

    /// Keys with their weights, in ascending key order. No weight is zero.
    pub fn iter(&self) -> impl Iterator<Item = (&K, i64)> {
        self.weights.iter().map(|(k, w)| (k, *w))
    }

    /// Adds `weight` to `key`'s, dropping the key at zero. Panics on `i64` overflow.
    pub fn insert(&mut self, key: K, weight: i64) {
        self.adjust(key, weight, false);
    }

    pub(crate) fn adjust(&mut self, key: K, weight: i64, subtract: bool) {
        match self.weights.entry(key) {
            Entry::Vacant(e) => {
                let from_zero = if subtract { negated(weight) } else { weight };
                if from_zero != 0 {
                    e.insert(from_zero);
                }
            }
            Entry::Occupied(mut e) => {
                let held = *e.get();
                let sum = if subtract {
                    minus(held, weight)
                } else {
                    plus(held, weight)
                };
                if sum == 0 {
                    e.remove();
                } else {
                    *e.get_mut() = sum;
                }
            }
        }
    }

    /// Maps every key through `f`; the weights of keys with one image sum.
    pub fn map<L: Ord>(&self, mut f: impl FnMut(&K) -> L) -> ZSet<L> {
        let mut out = ZSet::new();
        for (k, w) in &self.weights {
            out.insert(f(k), *w);
        }
        out
    }

    /// Maps every key to any number of images, each carrying the key's weight; the
    /// weights of equal images sum.
    pub fn flat_map<L: Ord, I: IntoIterator<Item = L>>(
        &self,
        mut f: impl FnMut(&K) -> I,
    ) -> ZSet<L> {
        let mut out = ZSet::new();
        for (k, w) in &self.weights {
            for image in f(k) {
                out.insert(image, *w);
            }
        }
        out
    }

    /// Keeps the keys `keep` accepts, with their weights. It looks at the key only,
    /// which is what makes it linear.
    pub fn filter(&self, mut keep: impl FnMut(&K) -> bool) -> Self
    where
        K: Clone,
    {
        self.weights
            .iter()
            .filter(|(k, _)| keep(k))
            .map(|(k, w)| (k.clone(), *w))
            .collect()
    }

    /// Weight 1 for every key whose weight is positive; every other key is dropped.
    pub fn distinct(&self) -> Self
    where
        K: Clone,
    {
        self.weights
            .iter()
            .filter(|(_, w)| **w > 0)
            .map(|(k, _)| (k.clone(), 1))
            .collect()
    }

    /// Every weight times `n`; zero gives the zero Z-set. Panics on `i64` overflow.
    pub fn scale(&self, n: i64) -> Self
    where
        K: Clone,
    {
        if n == 0 {
            return Self::new();
        }
        let weights = self
            .weights
            .iter()
            .map(|(k, w)| (k.clone(), times(*w, n)))
            .collect();
        ZSet { weights }
    }
}

impl<K: Ord> FromIterator<(K, i64)> for ZSet<K> {
    /// Weights of repeated keys sum.
    fn from_iter<I: IntoIterator<Item = (K, i64)>>(iter: I) -> Self {
        let mut z = ZSet::new();
        z.extend(iter);
        z
    }
}

impl<K: Ord> Extend<(K, i64)> for ZSet<K> {
    fn extend<I: IntoIterator<Item = (K, i64)>>(&mut self, iter: I) {
        for (k, w) in iter {
            self.insert(k, w);
        }
    }
}

impl<K: Ord> IntoIterator for ZSet<K> {
    type Item = (K, i64);
    type IntoIter = std::collections::btree_map::IntoIter<K, i64>;

    fn into_iter(self) -> Self::IntoIter {
        self.weights.into_iter()
    }
}

impl<K: Ord> AddAssign for ZSet<K> {
    fn add_assign(&mut self, other: Self) {
        self.extend(other);
    }
}

impl<K: Ord + Clone> AddAssign<&ZSet<K>> for ZSet<K> {
    fn add_assign(&mut self, other: &ZSet<K>) {
        self.extend(other.iter().map(|(k, w)| (k.clone(), w)));
    }
}

impl<K: Ord> Add for ZSet<K> {
    type Output = ZSet<K>;

    fn add(mut self, other: Self) -> Self {
        self += other;
        self
    }
}

impl<K: Ord + Clone> Add for &ZSet<K> {
    type Output = ZSet<K>;

    fn add(self, other: Self) -> ZSet<K> {
        let mut z = self.clone();
        z += other;
        z
    }
}

impl<K: Ord> Neg for ZSet<K> {
    type Output = ZSet<K>;

    fn neg(self) -> Self {
        let weights = self
            .weights
            .into_iter()
            .map(|(k, w)| (k, negated(w)))
            .collect();
        ZSet { weights }
    }
}

impl<K: Ord + Clone> Neg for &ZSet<K> {
    type Output = ZSet<K>;

    fn neg(self) -> ZSet<K> {
        self.clone().neg()
    }
}

impl<K: Ord> Sub for ZSet<K> {
    type Output = ZSet<K>;

    fn sub(mut self, other: Self) -> Self {
        self -= other;
        self
    }
}

impl<K: Ord + Clone> Sub for &ZSet<K> {
    type Output = ZSet<K>;

    fn sub(self, other: Self) -> ZSet<K> {
        self.clone() - other.clone()
    }
}

impl<K: Ord> SubAssign for ZSet<K> {
    fn sub_assign(&mut self, other: Self) {
        for (k, w) in other {
            self.adjust(k, w, true);
        }
    }
}

impl<K: Ord + Clone> Mul<i64> for &ZSet<K> {
    type Output = ZSet<K>;

    fn mul(self, n: i64) -> ZSet<K> {
        self.scale(n)
    }
}

/// The running sums of a stream of deltas: element `i` is `d_0 + ... + d_i`.
pub fn integrate<K: Ord + Clone>(deltas: &[ZSet<K>]) -> Vec<ZSet<K>> {
    let mut state = ZSet::new();
    deltas
        .iter()
        .map(|d| {
            state += d;
            state.clone()
        })
        .collect()
}

/// The differences of a stream of states: element 0 is `s_0` (the state before it is
/// zero), and element `i` is `s_i - s_(i-1)`.
pub fn differentiate<K: Ord + Clone>(states: &[ZSet<K>]) -> Vec<ZSet<K>> {
    let mut previous = ZSet::new();
    states
        .iter()
        .map(|s| {
            let d = s - &previous;
            previous = s.clone();
            d
        })
        .collect()
}

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

    #[test]
    fn a_zero_weight_is_never_stored() {
        let z = ZSet::from_iter([("a", 2), ("a", -2), ("b", 0), ("c", 1)]);
        assert_eq!(z.len(), 1);
        assert_eq!(z.weight(&"a"), 0);
        assert_eq!(z, ZSet::singleton("c", 1));
        assert!(ZSet::singleton("x", 0).is_zero());
    }

    #[test]
    fn distinct_clamps_positive_weights_and_drops_the_rest() {
        let z = ZSet::from_iter([("a", 5), ("b", -3), ("c", 1)]);
        assert_eq!(z.distinct(), ZSet::from_iter([("a", 1), ("c", 1)]));
    }

    #[test]
    fn map_sums_colliding_images() {
        let z = ZSet::from_iter([(1, 2), (2, -2), (3, 4)]);
        assert_eq!(z.map(|k| k % 2), ZSet::from_iter([(0, -2), (1, 6)]));
        assert!(ZSet::from_iter([(1, 2), (2, -2)]).map(|_| 0).is_zero());
    }

    #[test]
    fn the_differentiate_of_a_stream_starts_from_zero() {
        let s = vec![ZSet::singleton(1, 1), ZSet::singleton(2, 1)];
        let d = differentiate(&s);
        assert_eq!(d[0], s[0]);
        assert_eq!(d[1], ZSet::from_iter([(1, -1), (2, 1)]));
        assert_eq!(integrate(&d), s);
    }

    #[test]
    #[should_panic(expected = "weight overflow in add")]
    fn adding_past_i64_max_panics() {
        let _ = ZSet::singleton(1, i64::MAX) + ZSet::singleton(1, 1);
    }

    #[test]
    #[should_panic(expected = "weight overflow in add")]
    fn adding_past_i64_min_panics() {
        let _ = ZSet::singleton(1, i64::MIN) + ZSet::singleton(1, -1);
    }

    #[test]
    #[should_panic(expected = "weight overflow in neg")]
    fn negating_i64_min_panics() {
        let _ = -ZSet::singleton(1, i64::MIN);
    }

    #[test]
    #[should_panic(expected = "weight overflow in scale")]
    fn scaling_past_i64_panics() {
        let _ = ZSet::singleton(1, i64::MAX).scale(2);
    }

    #[test]
    fn subtracting_a_set_from_itself_at_i64_min_is_zero() {
        let a = ZSet::singleton(1, i64::MIN);
        assert!(a.clone().sub(a.clone()).is_zero());
        assert!((&a - &a).is_zero());
        assert!(differentiate(&[a.clone(), a]).pop().unwrap().is_zero());
    }

    #[test]
    #[should_panic(expected = "weight overflow in neg")]
    fn subtracting_i64_min_from_nothing_panics() {
        let _ = ZSet::<u8>::new() - ZSet::singleton(1, i64::MIN);
    }

    #[test]
    fn overflow_depends_on_the_order_of_evaluation_and_never_gives_a_wrong_value() {
        let (max, one, minus) = (
            ZSet::singleton(1, i64::MAX),
            ZSet::singleton(1, 1),
            ZSet::singleton(1, -1),
        );
        let right = max.clone() + (one.clone() + minus.clone());
        assert_eq!(right, max);
        let left = std::panic::catch_unwind(|| (max.clone() + one.clone()) + minus.clone());
        assert!(left.is_err());
    }

    #[test]
    #[should_panic(expected = "weight overflow in sub")]
    fn subtracting_past_i64_max_panics() {
        let _ = ZSet::singleton(1, i64::MAX) - ZSet::singleton(1, -1);
    }

    #[test]
    #[should_panic(expected = "weight overflow in sub")]
    fn subtracting_past_i64_min_panics() {
        let _ = ZSet::singleton(1, i64::MIN) - ZSet::singleton(1, 1);
    }

    #[test]
    #[should_panic(expected = "weight overflow in scale")]
    fn scaling_i64_min_by_minus_one_panics() {
        let _ = ZSet::singleton(1, i64::MIN).scale(-1);
    }

    #[test]
    fn scaling_by_zero_never_overflows() {
        assert!(ZSet::singleton(1, i64::MAX).scale(0).is_zero());
    }
}