improvise/crates/zset/tests/rollback_model.rs

//! A model-based property over every public operator: a rollback leaves the circuit equal to one
//! that never ran the step. From the review of #355. Evidence: `context/test-evidence/zset.md`.

use proptest::prelude::*;
use proptest::test_runner::{Config, TestRunner};
use std::cell::Cell;
use std::collections::BTreeMap;
use std::ops::Bound;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use zset::{Circuit, Input, Output, Stream, Trace, ZSet};

type P = (u8, u8);
type Z = ZSet<P>;

thread_local! { static TRIP: Cell<Option<u32>> = const { Cell::new(None) }; }

fn maybe_trip() {
    TRIP.with(|t| {
        if let Some(n) = t.get() {
            if n == 0 {
                t.set(None);
                panic!("tripped");
            }
            t.set(Some(n - 1));
        }
    })
}

#[derive(Default)]
struct VecTrace<K, V>(BTreeMap<(K, V), i64>);

impl<K: Ord + Clone, V: Ord + Clone> Trace<K, V> for VecTrace<K, V> {
    fn apply(&mut self, k: &K, v: &V, w: i64) {
        maybe_trip();
        let key = (k.clone(), v.clone());
        let cur = self.0.get(&key).copied().unwrap_or(0);
        let n = cur.checked_add(w).expect("overflow in VecTrace");
        if n == 0 {
            self.0.remove(&key);
        } else {
            self.0.insert(key, n);
        }
    }
    fn lookup(&self, k: &K) -> Vec<(V, i64)> {
        self.0
            .iter()
            .filter(|((kk, _), _)| kk == k)
            .map(|((_, v), w)| (v.clone(), *w))
            .collect()
    }
    fn range(&self, lo: Bound<&K>, hi: Bound<&K>) -> Vec<(K, V, i64)> {
        self.0
            .iter()
            .filter(|((k, _), _)| std::ops::RangeBounds::contains(&(lo, hi), k))
            .map(|((k, v), w)| (k.clone(), v.clone(), *w))
            .collect()
    }
}

#[derive(Clone, Debug)]
enum Op {
    Integrate(usize),
    Delay(usize),
    Diff(usize),
    Distinct(usize),
    Plus(usize, usize),
    Map(usize),
    Filter(usize),
    FlatMap(usize),
    IndexBy(usize),
    Join(usize, usize),
    JoinCustom(usize, usize),
    Agg(usize),
    AggCustom(usize),
}

#[derive(Clone, Debug)]
enum Act {
    Push(bool, Z),
    Step(Option<u32>),
    Cont(Option<u32>),
    Begin,
    Commit,
    Rollback,
    Build(Op),
}

struct World {
    c: Circuit,
    inputs: [Input<P>; 2],
    streams: Vec<Stream<Z>>,
    sinks: Vec<Output<P>>,
}

fn agg_f(_: &u8, vals: &[(u8, i64)]) -> Option<u8> {
    maybe_trip();
    Some(
        vals.iter()
            .fold(0u8, |s, (v, w)| s.wrapping_add(v.wrapping_mul(*w as u8))),
    )
}

impl World {
    fn new() -> Self {
        let mut c = Circuit::new();
        let (a, sa) = c.input::<P>();
        let (b, sb) = c.input::<P>();
        let sinks = vec![c.output(&sa), c.output(&sb)];
        World {
            c,
            inputs: [a, b],
            streams: vec![sa, sb],
            sinks,
        }
    }
    fn build(&mut self, op: &Op) {
        let c = &mut self.c;
        let st = &self.streams;
        struct S<'a>(&'a [Stream<Z>]);
        impl S<'_> {
            fn s(&self, i: usize) -> Stream<Z> {
                self.0[i % self.0.len()]
            }
        }
        let this = S(st);
        let out: Stream<Z> = match *op {
            Op::Integrate(i) => {
                let s = this.s(i);
                c.integrate(&s)
            }
            Op::Delay(i) => {
                let s = this.s(i);
                c.delay(&s)
            }
            Op::Diff(i) => {
                let s = this.s(i);
                c.differentiate(&s)
            }
            Op::Distinct(i) => {
                let s = this.s(i);
                c.distinct(&s)
            }
            Op::Plus(i, j) => {
                let (x, y) = (this.s(i), this.s(j));
                c.plus(&x, &y)
            }
            Op::Map(i) => {
                let s = this.s(i);
                c.map(&s, |p: &P| {
                    maybe_trip();
                    (p.1, p.0 % 4)
                })
            }
            Op::Filter(i) => {
                let s = this.s(i);
                c.filter(&s, |p: &P| {
                    maybe_trip();
                    p.0.is_multiple_of(2)
                })
            }
            Op::FlatMap(i) => {
                let s = this.s(i);
                c.flat_map(&s, |p: &P| {
                    maybe_trip();
                    vec![*p, (p.0, p.1.wrapping_add(1) % 4)]
                })
            }
            Op::IndexBy(i) => {
                let s = this.s(i);
                let ix = c.index_by(&s, |p: &P| p.1);
                c.map(&ix, |(g, p): &(u8, P)| (*g, p.0))
            }
            Op::Join(i, j) => {
                let (x, y) = (this.s(i), this.s(j));
                c.join_indexed(&x, &y, |k, a, b| (*k, (a ^ b) % 4))
            }
            Op::JoinCustom(i, j) => {
                let (x, y) = (this.s(i), this.s(j));
                c.join_indexed_in(
                    &x,
                    &y,
                    VecTrace::default(),
                    VecTrace::default(),
                    |k, a, b| (*k, a.wrapping_add(*b) % 4),
                )
            }
            Op::Agg(i) => {
                let s = this.s(i);
                c.aggregate(&s, agg_f)
            }
            Op::AggCustom(i) => {
                let s = this.s(i);
                c.aggregate_in(&s, VecTrace::default(), agg_f)
            }
        };
        let o = self.c.output(&out);
        self.streams.push(out);
        self.sinks.push(o);
    }
    fn view(&self) -> Vec<Z> {
        self.streams.iter().map(|s| self.c.current(s)).collect()
    }
    fn drain(&mut self) -> Vec<Z> {
        self.sinks.iter().map(|s| self.c.take(s)).collect()
    }
}

static CASES: AtomicUsize = AtomicUsize::new(0);
static RB_AFTER_PANIC: AtomicUsize = AtomicUsize::new(0);
static RB_CLEAN: AtomicUsize = AtomicUsize::new(0);
static RB_AFTER_CONT: AtomicUsize = AtomicUsize::new(0);
static COMMITS: AtomicUsize = AtomicUsize::new(0);
static BUILT_IN_TXN: AtomicUsize = AtomicUsize::new(0);
static DOUBLE_BEGIN: AtomicUsize = AtomicUsize::new(0);
static OUTSIDE_PANIC_STOP: AtomicUsize = AtomicUsize::new(0);
static MULTI_STEP_TXN: AtomicUsize = AtomicUsize::new(0);
static TRIPS: AtomicUsize = AtomicUsize::new(0);

fn weight() -> impl Strategy<Value = i64> {
    prop_oneof![
        8 => -3i64..=3,
        1 => Just(i64::MAX),
        1 => Just(i64::MIN),
        1 => Just(i64::MAX / 2 + 1),
        1 => Just(-(i64::MAX / 2) - 1),
    ]
}

fn delta() -> impl Strategy<Value = Z> {
    prop::collection::vec(((0u8..4, 0u8..4), weight()), 0..5).prop_map(|v| {
        let mut z = Z::new();
        for (k, w) in v {
            // build without overflowing: replace if it would
            if z.weight(&k).checked_add(w).is_some() {
                z.insert(k, w);
            }
        }
        z
    })
}

fn op() -> impl Strategy<Value = Op> {
    let i = 0usize..64;
    prop_oneof![
        i.clone().prop_map(Op::Integrate),
        i.clone().prop_map(Op::Delay),
        i.clone().prop_map(Op::Diff),
        i.clone().prop_map(Op::Distinct),
        (i.clone(), i.clone()).prop_map(|(a, b)| Op::Plus(a, b)),
        i.clone().prop_map(Op::Map),
        i.clone().prop_map(Op::Filter),
        i.clone().prop_map(Op::FlatMap),
        i.clone().prop_map(Op::IndexBy),
        (i.clone(), i.clone()).prop_map(|(a, b)| Op::Join(a, b)),
        (i.clone(), i.clone()).prop_map(|(a, b)| Op::JoinCustom(a, b)),
        i.clone().prop_map(Op::Agg),
        i.prop_map(Op::AggCustom),
    ]
}

fn act() -> impl Strategy<Value = Act> {
    let trip = prop_oneof![2 => Just(None), 1 => (0u32..40).prop_map(Some)];
    prop_oneof![
        4 => (any::<bool>(), delta()).prop_map(|(b, z)| Act::Push(b, z)),
        4 => trip.clone().prop_map(Act::Step),
        2 => trip.prop_map(Act::Cont),
        2 => Just(Act::Begin),
        1 => Just(Act::Commit),
        2 => Just(Act::Rollback),
        1 => op().prop_map(Act::Build),
    ]
}

fn apply_plain(w: &mut World, a: &Act) {
    match a {
        Act::Push(b, z) => w.c.push(&w.inputs[*b as usize], z.clone()),
        Act::Step(_) => w.c.step(),
        Act::Cont(_) => w.c.step_continue(),
        Act::Build(op) => w.build(op),
        _ => unreachable!(),
    }
}

fn check(x: &mut World, y: &mut World, at: &str) -> Result<(), TestCaseError> {
    prop_assert_eq!(x.c.ticks(), y.c.ticks(), "ticks {}", at);
    prop_assert_eq!(x.view(), y.view(), "view {}", at);
    prop_assert_eq!(x.drain(), y.drain(), "sinks {}", at);
    Ok(())
}

fn run(ops: &[Op], acts: &[Act]) -> Result<(), TestCaseError> {
    CASES.fetch_add(1, Relaxed);
    let (mut x, mut y) = (World::new(), World::new());
    for o in ops {
        x.build(o);
        y.build(o);
    }
    let mut open = false;
    let mut failed = false;
    let mut buf: Vec<Act> = vec![];
    let mut steps_in_txn = 0;
    let mut cont_in_txn = false;
    for (n, a) in acts.iter().enumerate() {
        let at = format!("after act {n} {a:?}");
        match a {
            Act::Begin => {
                if open {
                    let r = catch_unwind(AssertUnwindSafe(|| x.c.begin()));
                    prop_assert!(r.is_err(), "double begin must panic");
                    DOUBLE_BEGIN.fetch_add(1, Relaxed);
                } else {
                    x.c.begin();
                    open = true;
                    failed = false;
                    steps_in_txn = 0;
                    cont_in_txn = false;
                }
            }
            Act::Commit if open && !failed => {
                x.c.commit();
                open = false;
                for b in buf.drain(..) {
                    apply_plain(&mut y, &b);
                }
                COMMITS.fetch_add(1, Relaxed);
                if steps_in_txn > 1 {
                    MULTI_STEP_TXN.fetch_add(1, Relaxed);
                }
                check(&mut x, &mut y, &at)?;
            }
            Act::Commit | Act::Rollback if open => {
                x.c.rollback();
                open = false;
                if failed {
                    RB_AFTER_PANIC.fetch_add(1, Relaxed);
                } else {
                    RB_CLEAN.fetch_add(1, Relaxed);
                }
                if cont_in_txn {
                    RB_AFTER_CONT.fetch_add(1, Relaxed);
                }
                if steps_in_txn > 1 {
                    MULTI_STEP_TXN.fetch_add(1, Relaxed);
                }
                for b in buf.drain(..) {
                    if let Act::Build(op) = b {
                        y.build(&op);
                    }
                }
                check(&mut x, &mut y, &at)?;
            }
            Act::Commit | Act::Rollback => {}
            _ if open && failed => {} // only rollback after a failed step
            Act::Step(t) | Act::Cont(t) if open => {
                TRIP.with(|c| c.set(*t));
                let r = catch_unwind(AssertUnwindSafe(|| apply_plain(&mut x, a)));
                let tripped = TRIP.with(|c| c.get()).is_none() && t.is_some();
                TRIP.with(|c| c.set(None));
                steps_in_txn += 1;
                if matches!(a, Act::Cont(_)) {
                    cont_in_txn = true;
                }
                if r.is_err() {
                    failed = true;
                    if tripped {
                        TRIPS.fetch_add(1, Relaxed);
                    }
                } else {
                    buf.push(a.clone());
                }
            }
            _ if open => {
                if matches!(a, Act::Build(_)) {
                    BUILT_IN_TXN.fetch_add(1, Relaxed);
                }
                let r = catch_unwind(AssertUnwindSafe(|| apply_plain(&mut x, a)));
                if r.is_err() {
                    failed = true;
                } else {
                    buf.push(a.clone());
                }
            }
            _ => {
                let r = catch_unwind(AssertUnwindSafe(|| apply_plain(&mut x, a)));
                if r.is_err() {
                    OUTSIDE_PANIC_STOP.fetch_add(1, Relaxed);
                    return Ok(());
                }
                apply_plain(&mut y, a);
                check(&mut x, &mut y, &at)?;
            }
        }
    }
    if open {
        x.c.rollback();
        for b in buf.drain(..) {
            if let Act::Build(op) = b {
                y.build(&op);
            }
        }
        if failed {
            RB_AFTER_PANIC.fetch_add(1, Relaxed);
        } else {
            RB_CLEAN.fetch_add(1, Relaxed);
        }
        check(&mut x, &mut y, "final rollback")?;
    }
    // ticks after: both behave the same on a fresh tick
    for i in 0..3u8 {
        let z: Z = [((i, i), 1), ((i + 1, 0), -1)].into_iter().collect();
        let r = catch_unwind(AssertUnwindSafe(|| {
            x.c.push(&x.inputs[0], z.clone());
            x.c.push(&x.inputs[1], z.clone());
            x.c.step()
        }));
        if r.is_err() {
            return Ok(());
        }
        y.c.push(&y.inputs[0], z.clone());
        y.c.push(&y.inputs[1], z);
        y.c.step();
        check(&mut x, &mut y, &format!("post tick {i}"))?;
    }
    Ok(())
}

#[test]
fn reviewer_rollback_matches_a_circuit_that_never_ran_the_step() {
    let cases: u32 = std::env::var("ATTACK_CASES")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(256);
    std::panic::set_hook(Box::new(|_| {}));
    let mut runner = TestRunner::new(Config {
        cases,
        ..Config::default()
    });
    let strat = (
        prop::collection::vec(op(), 1..12),
        prop::collection::vec(act(), 1..40),
    );
    let r = runner.run(&strat, |(ops, acts)| run(&ops, &acts));
    eprintln!(
        "cases={} rb_after_panic={} rb_clean={} rb_after_cont={} commits={} built_in_txn={} double_begin={} outside_panic_stop={} multi_step_txn={} trips={}",
        CASES.load(Relaxed),
        RB_AFTER_PANIC.load(Relaxed),
        RB_CLEAN.load(Relaxed),
        RB_AFTER_CONT.load(Relaxed),
        COMMITS.load(Relaxed),
        BUILT_IN_TXN.load(Relaxed),
        DOUBLE_BEGIN.load(Relaxed),
        OUTSIDE_PANIC_STOP.load(Relaxed),
        MULTI_STEP_TXN.load(Relaxed),
        TRIPS.load(Relaxed)
    );
    let _ = std::panic::take_hook();
    if let Err(e) = r {
        eprintln!("FAILURE: {e:?}");
        panic!("attack failed");
    }
}