improvise/crates/zset/tests/circuit.rs

//! The circuit's operators against recomputing from scratch, and the DBSP identities.
//! Evidence: `context/test-evidence/zset.md`.

use proptest::prelude::*;
use std::cell::Cell;
use std::ops::Bound;
use std::rc::Rc;
use zset::{
    AggregateState, Circuit, MemTrace, Progress, Rebuild, Stream, Trace, ZSet, differentiate,
    integrate,
};

type Z<K> = ZSet<K>;

fn pairs() -> impl Strategy<Value = Vec<(u8, i64)>> {
    prop::collection::vec((0u8..6, -3i64..=3), 0..6)
}

fn stream() -> impl Strategy<Value = Vec<Z<u8>>> {
    prop::collection::vec(pairs().prop_map(|p| p.into_iter().collect::<Z<u8>>()), 1..8)
}

fn kv() -> impl Strategy<Value = Vec<Z<(u8, u8)>>> {
    let one = prop::collection::vec(((0u8..4, 0u8..3), -3i64..=3), 0..6)
        .prop_map(|p| p.into_iter().collect::<Z<(u8, u8)>>());
    prop::collection::vec(one, 1..8)
}

/// Feeds `steps` to one input, one circuit step each, and checks after every step that
/// the sum of the outputs so far is `reference` of the sum of the inputs so far.
fn matches_recompute<K, L>(
    steps: &[Z<K>],
    build: impl FnOnce(&mut Circuit, &Stream<Z<K>>) -> Stream<Z<L>>,
    reference: impl Fn(&Z<K>) -> Z<L>,
) where
    K: Ord + Clone + std::fmt::Debug + 'static,
    L: Ord + Clone + std::fmt::Debug + 'static,
{
    let mut c = Circuit::new();
    let (input, s) = c.input::<K>();
    let out = build(&mut c, &s);
    let sink = c.output(&out);
    let (mut fed, mut got) = (Z::new(), Z::new());
    for (i, d) in steps.iter().enumerate() {
        input.push(d.clone());
        c.step();
        fed += d;
        got += sink.take();
        assert_eq!(got, reference(&fed), "[incremental] after step {i}");
    }
}

fn sum_by_key(z: &Z<(u8, u8)>) -> Z<(u8, i64)> {
    group(z, |vals| {
        let s: i64 = vals.iter().map(|(v, w)| i64::from(*v) * w).sum();
        (s != 0).then_some(s)
    })
}

/// The recompute-from-scratch aggregate: `f` over each key's integrated values.
fn group(z: &Z<(u8, u8)>, f: impl Fn(&[(u8, i64)]) -> Option<i64>) -> Z<(u8, i64)> {
    let mut out = Z::new();
    for k in z
        .iter()
        .map(|((k, _), _)| *k)
        .collect::<std::collections::BTreeSet<_>>()
    {
        let vals: Vec<(u8, i64)> = z
            .iter()
            .filter(|((a, _), _)| *a == k)
            .map(|((_, v), w)| (*v, w))
            .collect();
        if let Some(r) = f(&vals) {
            out.insert((k, r), 1);
        }
    }
    out
}

fn max_present(vals: &[(u8, i64)]) -> Option<i64> {
    vals.iter()
        .filter(|(_, w)| *w > 0)
        .map(|(v, _)| i64::from(*v))
        .max()
}

fn join_ref(a: &Z<(u8, u8)>, b: &Z<(u8, u8)>) -> Z<(u8, u8, u8)> {
    let mut out = Z::new();
    for ((k, x), wa) in a.iter() {
        for ((k2, y), wb) in b.iter() {
            if k == k2 {
                out.insert((*k, *x, *y), wa * wb);
            }
        }
    }
    out
}

fn join_circuit(da: &[Z<(u8, u8)>], db: &[Z<(u8, u8)>]) -> Vec<Z<(u8, u8, u8)>> {
    let mut c = Circuit::new();
    let (ia, a) = c.input::<(u8, u8)>();
    let (ib, b) = c.input::<(u8, u8)>();
    let j = c.join_indexed(&a, &b, |k, x, y| (*k, *x, *y));
    let sink = c.output(&j);
    let mut deltas = Vec::new();
    for i in 0..da.len().max(db.len()) {
        ia.push(da.get(i).cloned().unwrap_or_default());
        ib.push(db.get(i).cloned().unwrap_or_default());
        c.step();
        deltas.push(sink.take());
    }
    deltas
}

proptest! {
    #[test]
    fn map_matches_recompute(steps in stream()) {
        matches_recompute(&steps, |c, s| c.map(s, |k| k % 3), |z| z.map(|k| k % 3));
    }

    #[test]
    fn filter_matches_recompute(steps in stream()) {
        matches_recompute(&steps, |c, s| c.filter(s, |k| k % 2 == 0), |z| z.filter(|k| k % 2 == 0));
    }

    #[test]
    fn flat_map_matches_recompute(steps in stream()) {
        let f = |k: &u8| vec![*k % 3, *k % 3 + 10, 20];
        matches_recompute(&steps, |c, s| c.flat_map(s, f), |z| z.flat_map(f));
    }

    #[test]
    fn index_by_matches_recompute(steps in stream()) {
        matches_recompute(&steps, |c, s| c.index_by(s, |k| k % 2), |z| z.map(|k| (k % 2, *k)));
    }

    #[test]
    fn distinct_matches_recompute(steps in stream()) {
        matches_recompute(&steps, |c, s| c.distinct(s), |z| z.distinct());
    }

    #[test]
    fn a_composed_pipeline_matches_recompute(steps in stream()) {
        let pipeline = |z: &Z<u8>| z.filter(|k| *k != 3).map(|k| k % 3).distinct();
        matches_recompute(
            &steps,
            |c, s| {
                let f = c.filter(s, |k| *k != 3);
                let m = c.map(&f, |k| k % 3);
                c.distinct(&m)
            },
            pipeline,
        );
    }

    #[test]
    fn aggregate_sum_matches_recompute(steps in kv()) {
        matches_recompute(
            &steps,
            |c, s| c.aggregate(s, |_, vals| {
                let t: i64 = vals.iter().map(|(v, w)| i64::from(*v) * w).sum();
                (t != 0).then_some(t)
            }),
            sum_by_key,
        );
    }

    #[test]
    fn aggregate_max_matches_recompute(steps in kv()) {
        matches_recompute(
            &steps,
            |c, s| c.aggregate(s, |_, vals| max_present(vals)),
            |z| group(z, max_present),
        );
    }

    #[test]
    fn join_matches_recompute(da in kv(), db in kv()) {
        let deltas = join_circuit(&da, &db);
        let (mut a, mut b, mut got) = (Z::new(), Z::new(), Z::new());
        for (i, d) in deltas.iter().enumerate() {
            a += da.get(i).cloned().unwrap_or_default();
            b += db.get(i).cloned().unwrap_or_default();
            got += d;
            prop_assert_eq!(&got, &join_ref(&a, &b), "[join] after step {}", i);
        }
    }

    #[test]
    fn join_is_bilinear(a1 in kv(), a2 in kv(), b1 in kv(), b2 in kv()) {
        let total = |s: &[Z<(u8, u8)>]| s.iter().fold(Z::new(), |acc, d| acc + d.clone());
        let (a1, a2, b1, b2) = (total(&a1), total(&a2), total(&b1), total(&b2));
        prop_assert_eq!(join_ref(&(&a1 + &a2), &b1), &join_ref(&a1, &b1) + &join_ref(&a2, &b1), "[bilinear-left]");
        prop_assert_eq!(join_ref(&a1, &(&b1 + &b2)), &join_ref(&a1, &b1) + &join_ref(&a1, &b2), "[bilinear-right]");
        // and the circuit's output is bilinear too: one step of a+a' equals two of a, a'.
        let once = join_circuit(&[&a1 + &a2], std::slice::from_ref(&b1));
        let twice = join_circuit(&[a1.clone(), a2.clone()], &[b1.clone(), Z::new()]);
        let sum_once = once.into_iter().fold(Z::new(), |x, d| x + d);
        let sum_twice = twice.into_iter().fold(Z::new(), |x, d| x + d);
        prop_assert_eq!(sum_once, sum_twice, "[bilinear-circuit]");
    }

    #[test]
    fn integrate_after_differentiate_is_the_identity(steps in stream()) {
        prop_assert_eq!(integrate(&differentiate(&steps)), steps.clone(), "[slice]");
        let mut c = Circuit::new();
        let (input, s) = c.input::<u8>();
        let d = c.differentiate(&s);
        let back = c.integrate(&d);
        let (sink_s, sink_back) = (c.output(&s), c.output(&back));
        let (mut seen_s, mut seen_back) = (Vec::new(), Vec::new());
        for d in &steps {
            input.push(d.clone());
            c.step();
            seen_s.push(s.current());
            seen_back.push(back.current());
        }
        let _ = (sink_s, sink_back);
        prop_assert_eq!(seen_back, seen_s, "[circuit-i-d]");
    }

    #[test]
    fn differentiate_after_integrate_is_the_identity(steps in stream()) {
        prop_assert_eq!(differentiate(&integrate(&steps)), steps.clone(), "[slice]");
        let mut c = Circuit::new();
        let (input, s) = c.input::<u8>();
        let i = c.integrate(&s);
        let back = c.differentiate(&i);
        for (n, d) in steps.iter().enumerate() {
            input.push(d.clone());
            c.step();
            prop_assert_eq!(back.current(), d.clone(), "[circuit-d-i] step {}", n);
        }
    }

    #[test]
    fn delay_is_the_previous_step(steps in stream()) {
        let mut c = Circuit::new();
        let (input, s) = c.input::<u8>();
        let d = c.delay(&s);
        let mut prev = Z::new();
        for (n, step) in steps.iter().enumerate() {
            input.push(step.clone());
            c.step();
            prop_assert_eq!(d.current(), prev.clone(), "[delay] step {}", n);
            prev = step.clone();
        }
    }

    /// Whatever the deadline does, the rebuilt output is what one delta would give.
    #[test]
    fn a_resumable_rebuild_equals_an_uninterrupted_step(
        snapshot in pairs(),
        chunk in 0usize..4,
        budget in 1usize..4,
    ) {
        let snapshot: Z<u8> = snapshot.into_iter().collect();
        let build = |c: &mut Circuit, s: &Stream<Z<u8>>| {
            let m = c.map(s, |k| k % 3);
            c.distinct(&m)
        };
        let mut whole = Circuit::new();
        let (input, s) = whole.input::<u8>();
        let out = build(&mut whole, &s);
        let sink = whole.output(&out);
        input.push(snapshot.clone());
        whole.step();
        let want = sink.take();

        let mut c = Circuit::new();
        let (input, s) = c.input::<u8>();
        let out = build(&mut c, &s);
        let sink = c.output(&out);
        let mut job = Rebuild::new(input, snapshot, chunk);
        let mut got = Z::new();
        let mut calls = 0;
        loop {
            calls += 1;
            prop_assert!(calls < 100, "[rebuild] never finished");
            let left = Cell::new(budget);
            let progress = job.step(&mut c, || {
                left.set(left.get() - 1);
                left.get() == 0
            });
            got += sink.take();
            if progress == Progress::Done { break; }
        }
        prop_assert_eq!(got, want, "[rebuild]");
    }

    #[test]
    fn a_trace_reads_what_was_applied(ops in prop::collection::vec(((0u8..5, 0u8..3), -3i64..=3), 0..12), lo in 0u8..6, hi in 0u8..6) {
        let mut t = MemTrace::new();
        let mut reference: Z<(u8, u8)> = Z::new();
        for ((k, v), w) in &ops {
            t.apply(k, v, *w);
            reference.insert((*k, *v), *w);
        }
        for k in 0..5u8 {
            let want: Vec<(u8, i64)> = reference.iter().filter(|((a, _), _)| *a == k).map(|((_, v), w)| (*v, w)).collect();
            prop_assert_eq!(t.lookup(&k), want, "[lookup] key {}", k);
        }
        let want: Vec<(u8, u8, i64)> = reference.iter().filter(|((a, _), _)| *a >= lo && *a < hi).map(|((a, v), w)| (*a, *v, w)).collect();
        prop_assert_eq!(t.range(Bound::Included(&lo), Bound::Excluded(&hi)), want, "[range]");
        let all: Vec<(u8, u8, i64)> = reference.iter().map(|((a, v), w)| (*a, *v, w)).collect();
        prop_assert_eq!(t.range(Bound::Unbounded, Bound::Unbounded), all, "[range-all]");
        prop_assert_eq!(t.keys(), reference.iter().map(|((a, _), _)| *a).collect::<std::collections::BTreeSet<_>>().len());
    }
}

#[test]
fn a_rebuild_that_is_out_of_time_reports_pending_after_one_chunk() {
    let mut c = Circuit::new();
    let (input, s) = c.input::<u8>();
    let sink = c.output(&s);
    let mut job = Rebuild::new(input, (0..6u8).map(|k| (k, 1)).collect(), 2);
    assert_eq!(job.step(&mut c, || true), Progress::Pending);
    assert_eq!(sink.take().len(), 2, "exactly one chunk ran");
    assert_eq!(job.step(&mut c, || false), Progress::Done);
    assert_eq!(sink.take().len(), 4);
    assert_eq!(job.step(&mut c, || true), Progress::Done);
}

#[test]
fn an_unbounded_range_that_is_empty_reads_nothing() {
    let mut t = MemTrace::new();
    t.apply(&1u8, &1u8, 1);
    assert!(t.range(Bound::Excluded(&1), Bound::Included(&1)).is_empty());
    assert!(t.range(Bound::Included(&2), Bound::Included(&1)).is_empty());
    assert_eq!(t.range(Bound::Included(&1), Bound::Included(&1)).len(), 1);
}

#[test]
fn the_join_reads_state_through_the_trace_it_was_given() {
    struct Counting(MemTrace<u8, u8>, Rc<Cell<usize>>);
    impl Trace<u8, u8> for Counting {
        fn apply(&mut self, k: &u8, v: &u8, w: i64) {
            self.0.apply(k, v, w)
        }
        fn lookup(&self, k: &u8) -> Vec<(u8, i64)> {
            self.1.set(self.1.get() + 1);
            self.0.lookup(k)
        }
        fn range(&self, lo: Bound<&u8>, hi: Bound<&u8>) -> Vec<(u8, u8, i64)> {
            self.0.range(lo, hi)
        }
    }
    let looked = Rc::new(Cell::new(0));
    let mut c = Circuit::new();
    let (ia, a) = c.input::<(u8, u8)>();
    let (ib, b) = c.input::<(u8, u8)>();
    let j = c.join_indexed_in(
        &a,
        &b,
        Counting(MemTrace::new(), looked.clone()),
        MemTrace::new(),
        |k, x, y| (*k, *x, *y),
    );
    ia.push(Z::singleton((1, 1), 1));
    ib.push(Z::singleton((1, 2), 1));
    c.step();
    assert_eq!(j.current(), Z::singleton((1, 1, 2), 1));
    assert!(looked.get() > 0);
}

#[test]
fn aggregate_state_refolds_only_the_touched_keys() {
    let mut st = AggregateState::new(MemTrace::<u8, u8>::new());
    let folds = Cell::new(0);
    let f = |_: &u8, vals: &[(u8, i64)]| {
        folds.set(folds.get() + 1);
        Some(vals.len() as i64)
    };
    let big: Z<(u8, u8)> = (0..50u8).map(|k| ((k, 0), 1)).collect();
    st.step(&big, f);
    folds.set(0);
    let out = st.step(&Z::singleton((7, 1), 1), f);
    assert_eq!(folds.get(), 2, "one key, folded before and after");
    assert_eq!(out, Z::from_iter([((7, 1), -1), ((7, 2), 1)]));
}