feat(zset): Circuit::begin, commit and rollback undo a failed step (improvise-bw4p.5)

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VdfEGXA7aJ6ftMsCc9PJze

crates/zset/src/circuit.rs

@@ -1,7 +1,9 @@
 //! The synchronous, single-threaded circuit. This module is the wiring (`Rc`, `RefCell`,
 //! boxed nodes); the operators' logic is in `ops.rs` and `lib.rs`.
 
-use crate::ops::{AggregateState, DistinctState, JoinState};
+use crate::ops::{
+    AggregateState, DistinctState, JoinState, Log, Transactional, Undoable, add_logged, unadd,
+};
 use crate::{MemTrace, Trace, ZSet};
 use std::cell::{Ref, RefCell};
 use std::rc::Rc;
@@ -12,6 +14,9 @@ use std::rc::Rc;
 pub(crate) struct Slot<T> {
     piece: T,
     whole: Option<T>,
+    /// Inside a transaction: the `(piece, whole)` the first publish replaced.
+    saved: Option<(T, Option<T>)>,
+    open: bool,
 }
 
 impl<T: Default> Default for Slot<T> {
@@ -19,6 +24,8 @@ impl<T: Default> Default for Slot<T> {
         Slot {
             piece: T::default(),
             whole: None,
+            saved: None,
+            open: false,
         }
     }
 }
@@ -48,14 +55,45 @@ impl<T: Piecewise> Stream<T> {
 
     fn publish(&self, v: T, cont: bool) {
         let mut slot = self.0.borrow_mut();
+        let slot = &mut *slot;
         if cont {
+            if slot.open && slot.saved.is_none() {
+                slot.saved = Some((slot.piece.clone(), slot.whole.clone()));
+            }
             let mut whole = slot.whole.take().unwrap_or_else(|| slot.piece.clone());
             whole.absorb(&v);
             slot.whole = Some(whole);
+            slot.piece = v;
         } else {
-            slot.whole = None;
+            let piece = std::mem::replace(&mut slot.piece, v);
+            let whole = slot.whole.take();
+            if slot.open && slot.saved.is_none() {
+                slot.saved = Some((piece, whole));
+            }
         }
-        slot.piece = v;
+    }
+}
+
+impl<T> Stream<T> {
+    fn begin(&self) {
+        let mut slot = self.0.borrow_mut();
+        slot.open = true;
+        slot.saved = None;
+    }
+
+    fn commit(&self) {
+        let mut slot = self.0.borrow_mut();
+        slot.open = false;
+        slot.saved = None;
+    }
+
+    fn rollback(&self) {
+        let mut slot = self.0.borrow_mut();
+        if let Some((piece, whole)) = slot.saved.take() {
+            slot.piece = piece;
+            slot.whole = whole;
+        }
+        slot.open = false;
     }
 }
 
@@ -89,7 +127,8 @@ impl<K: Ord> Input<K> {
     }
 }
 
-/// A sink that sums every step's delta until it is taken.
+/// A sink that sums every step's delta until it is taken. Take inside no transaction: a
+/// rollback subtracts the step's output, which a `take` in between has already removed.
 pub struct Output<K>(Rc<RefCell<ZSet<K>>>);
 
 impl<K: Ord> Output<K> {
@@ -98,80 +137,237 @@ impl<K: Ord> Output<K> {
     }
 }
 
-trait Node {
+trait Node: Transactional {
     /// `cont` is true for a later piece of a tick that has already begun.
     fn eval(&mut self, cont: bool);
 }
 
-struct Unary<I, O, F> {
+/// One input, one output, and the state `f` may change.
+struct Unary<I, O, S, F> {
     src: Stream<I>,
     out: Stream<O>,
+    state: S,
     f: F,
 }
 
-impl<I: Piecewise, O: Piecewise, F: FnMut(&I) -> O> Node for Unary<I, O, F> {
+impl<I, O, S: Transactional, F> Transactional for Unary<I, O, S, F> {
+    fn begin(&mut self) {
+        self.state.begin();
+        self.out.begin();
+    }
+
+    fn commit(&mut self) {
+        self.state.commit();
+        self.out.commit();
+    }
+
+    fn rollback(&mut self) {
+        self.state.rollback();
+        self.out.rollback();
+    }
+}
+
+impl<I: Piecewise, O: Piecewise, S: Transactional, F: FnMut(&mut S, &I) -> O> Node
+    for Unary<I, O, S, F>
+{
     fn eval(&mut self, cont: bool) {
-        let v = (self.f)(&self.src.piece());
+        let v = (self.f)(&mut self.state, &self.src.piece());
         self.out.publish(v, cont);
     }
 }
 
-struct Binary<A, B, O, F> {
+struct Binary<A, B, O, S, F> {
     a: Stream<A>,
     b: Stream<B>,
     out: Stream<O>,
+    state: S,
     f: F,
 }
 
-impl<A: Piecewise, B: Piecewise, O: Piecewise, F: FnMut(&A, &B) -> O> Node for Binary<A, B, O, F> {
+impl<A, B, O, S: Transactional, F> Transactional for Binary<A, B, O, S, F> {
+    fn begin(&mut self) {
+        self.state.begin();
+        self.out.begin();
+    }
+
+    fn commit(&mut self) {
+        self.state.commit();
+        self.out.commit();
+    }
+
+    fn rollback(&mut self) {
+        self.state.rollback();
+        self.out.rollback();
+    }
+}
+
+impl<A: Piecewise, B: Piecewise, O: Piecewise, S: Transactional, F: FnMut(&mut S, &A, &B) -> O> Node
+    for Binary<A, B, O, S, F>
+{
     fn eval(&mut self, cont: bool) {
-        let v = (self.f)(&self.a.piece(), &self.b.piece());
+        let v = (self.f)(&mut self.state, &self.a.piece(), &self.b.piece());
         self.out.publish(v, cont);
     }
 }
 
+/// What a `Delay` holds, and what it held when a transaction began.
+struct Held<K> {
+    held: ZSet<K>,
+    open: bool,
+    saved: Option<ZSet<K>>,
+}
+
+impl<K: Ord + Clone> Held<K> {
+    /// Notes the value to go back to, once per transaction. `held` is one tick's delta.
+    fn save(&mut self) {
+        if self.open && self.saved.is_none() {
+            self.saved = Some(self.held.clone());
+        }
+    }
+}
+
+impl<K> Transactional for Held<K> {
+    fn begin(&mut self) {
+        self.open = true;
+        self.saved = None;
+    }
+
+    fn commit(&mut self) {
+        self.open = false;
+        self.saved = None;
+    }
+
+    fn rollback(&mut self) {
+        if let Some(held) = self.saved.take() {
+            self.held = held;
+        }
+        self.open = false;
+    }
+}
+
 /// out(t) = in(t - 1). A continued piece adds to what is held and emits nothing.
 struct Delay<K> {
     src: Stream<ZSet<K>>,
     out: Stream<ZSet<K>>,
-    held: ZSet<K>,
+    state: Held<K>,
+}
+
+impl<K> Transactional for Delay<K> {
+    fn begin(&mut self) {
+        self.state.begin();
+        self.out.begin();
+    }
+
+    fn commit(&mut self) {
+        self.state.commit();
+        self.out.commit();
+    }
+
+    fn rollback(&mut self) {
+        self.state.rollback();
+        self.out.rollback();
+    }
 }
 
 impl<K: Ord + Clone> Node for Delay<K> {
     fn eval(&mut self, cont: bool) {
         let d = self.src.piece();
+        self.state.save();
         let v = if cont {
-            self.held += &*d;
+            self.state.held += &*d;
             ZSet::new()
         } else {
-            std::mem::replace(&mut self.held, d.clone())
+            std::mem::replace(&mut self.state.held, d.clone())
         };
         drop(d);
         self.out.publish(v, cont);
     }
 }
 
+/// The integral, and the log that takes a transaction's additions back.
+struct Acc<K> {
+    acc: ZSet<K>,
+    log: Log<(K, i64)>,
+}
+
+impl<K: Ord> Transactional for Acc<K> {
+    fn begin(&mut self) {
+        self.log.begin();
+    }
+
+    fn commit(&mut self) {
+        self.log.commit();
+    }
+
+    fn rollback(&mut self) {
+        let acc = &mut self.acc;
+        self.log.rollback(|e| unadd(acc, e));
+    }
+}
+
 /// out(t) = in(0) + .. + in(t). A continued piece emits only its own input.
 struct Integrate<K> {
     src: Stream<ZSet<K>>,
     out: Stream<ZSet<K>>,
-    acc: ZSet<K>,
+    state: Acc<K>,
+}
+
+impl<K: Ord> Transactional for Integrate<K> {
+    fn begin(&mut self) {
+        self.state.begin();
+        self.out.begin();
+    }
+
+    fn commit(&mut self) {
+        self.state.commit();
+        self.out.commit();
+    }
+
+    fn rollback(&mut self) {
+        self.state.rollback();
+        self.out.rollback();
+    }
 }
 
 impl<K: Ord + Clone> Node for Integrate<K> {
     fn eval(&mut self, cont: bool) {
         let d = self.src.piece();
-        self.acc += &*d;
-        let v = if cont { d.clone() } else { self.acc.clone() };
+        add_logged(&mut self.state.acc, &d, &mut self.state.log);
+        let v = if cont {
+            d.clone()
+        } else {
+            self.state.acc.clone()
+        };
         drop(d);
         self.out.publish(v, cont);
     }
 }
 
-/// Moves what was pushed into the circuit's stream.
+/// Moves what was pushed into the circuit's stream. Inside a transaction it remembers what
+/// was pending, so a rollback gives back the delta the failed step took.
 struct Source<K> {
     pending: Rc<RefCell<ZSet<K>>>,
     out: Stream<ZSet<K>>,
+    saved: Option<ZSet<K>>,
+}
+
+impl<K: Ord + Clone> Transactional for Source<K> {
+    fn begin(&mut self) {
+        self.saved = Some(self.pending.borrow().clone());
+        self.out.begin();
+    }
+
+    fn commit(&mut self) {
+        self.saved = None;
+        self.out.commit();
+    }
+
+    fn rollback(&mut self) {
+        if let Some(saved) = self.saved.take() {
+            *self.pending.borrow_mut() = saved;
+        }
+        self.out.rollback();
+    }
 }
 
 impl<K: Ord + Clone> Node for Source<K> {
@@ -181,6 +377,27 @@ impl<K: Ord + Clone> Node for Source<K> {
     }
 }
 
+/// The sink's sum, and the log that takes a transaction's additions back.
+struct Sink<K> {
+    sum: Rc<RefCell<ZSet<K>>>,
+    log: Log<(K, i64)>,
+}
+
+impl<K: Ord> Transactional for Sink<K> {
+    fn begin(&mut self) {
+        self.log.begin();
+    }
+
+    fn commit(&mut self) {
+        self.log.commit();
+    }
+
+    fn rollback(&mut self) {
+        let sum = &mut *self.sum.borrow_mut();
+        self.log.rollback(|e| unadd(sum, e));
+    }
+}
+
 /// Nodes run in creation order, which is a topological order because a node can only
 /// be built from streams that already exist. There are no cycles to schedule.
 #[derive(Default)]
@@ -189,6 +406,8 @@ pub struct Circuit {
     ticks: u64,
     /// This circuit's identity: every `Input` it makes holds a clone.
     token: Rc<()>,
+    /// `Some(ticks)` as of `begin`, while a transaction is open.
+    open: Option<u64>,
 }
 
 type Z<K> = ZSet<K>;
@@ -215,21 +434,101 @@ impl Circuit {
         self.run(true);
     }
 
+    /// Opens a transaction: from here until `commit` or `rollback`, every state change is
+    /// logged so `rollback` can take it back. The log costs the keys the steps touch, not
+    /// the circuit's size. Panics if one is already open. See `context/code/zset.md`.
+    pub fn begin(&mut self) {
+        assert!(
+            self.open.is_none(),
+            "zset: a transaction is already open on this circuit"
+        );
+        self.open = Some(self.ticks);
+        for n in &mut self.nodes {
+            n.begin();
+        }
+    }
+
+    /// Keeps what the transaction did. Panics if none is open.
+    pub fn commit(&mut self) {
+        assert!(
+            self.open.take().is_some(),
+            "zset: no transaction is open on this circuit"
+        );
+        for n in &mut self.nodes {
+            n.commit();
+        }
+    }
+
+    /// Puts the circuit's state back to what it was at `begin`, including after a step
+    /// that panicked half way. Pushes made since `begin` are dropped; what was pending at
+    /// `begin` is pending again. Panics if no transaction is open.
+    pub fn rollback(&mut self) {
+        let ticks = self
+            .open
+            .take()
+            .expect("zset: no transaction is open on this circuit");
+        for n in &mut self.nodes {
+            n.rollback();
+        }
+        self.ticks = ticks;
+    }
+
     fn run(&mut self, cont: bool) {
         for n in &mut self.nodes {
             n.eval(cont);
         }
     }
 
+    /// Adds a node; one built inside a transaction starts recording at once.
+    fn add(&mut self, mut node: Box<dyn Node>) {
+        if self.open.is_some() {
+            node.begin();
+        }
+        self.nodes.push(node);
+    }
+
+    fn unary_with<I: Piecewise + 'static, O: Piecewise + 'static, S: Transactional + 'static>(
+        &mut self,
+        src: &Stream<I>,
+        state: S,
+        f: impl FnMut(&mut S, &I) -> O + 'static,
+    ) -> Stream<O> {
+        let out = Stream(Rc::default());
+        self.add(Box::new(Unary {
+            src: src.clone(),
+            out: out.clone(),
+            state,
+            f,
+        }));
+        out
+    }
+
     fn unary<I: Piecewise + 'static, O: Piecewise + 'static>(
         &mut self,
         src: &Stream<I>,
-        f: impl FnMut(&I) -> O + 'static,
+        mut f: impl FnMut(&I) -> O + 'static,
+    ) -> Stream<O> {
+        self.unary_with(src, (), move |_, d| f(d))
+    }
+
+    fn binary_with<
+        A: Piecewise + 'static,
+        B: Piecewise + 'static,
+        O: Piecewise + 'static,
+        S: Transactional + 'static,
+    >(
+        &mut self,
+        a: &Stream<A>,
+        b: &Stream<B>,
+        state: S,
+        f: impl FnMut(&mut S, &A, &B) -> O + 'static,
     ) -> Stream<O> {
         let out = Stream(Rc::default());
-        self.nodes.push(Box::new(Unary {
-            src: src.clone(),
+        self.add(Box::new(Binary {
+            a: a.clone(),
+            b: b.clone(),
             out: out.clone(),
+            state,
             f,
         }));
         out
@@ -239,33 +538,30 @@ impl Circuit {
         &mut self,
         a: &Stream<A>,
         b: &Stream<B>,
-        f: impl FnMut(&A, &B) -> O + 'static,
+        mut f: impl FnMut(&A, &B) -> O + 'static,
     ) -> Stream<O> {
-        let out = Stream(Rc::default());
-        self.nodes.push(Box::new(Binary {
-            a: a.clone(),
-            b: b.clone(),
-            out: out.clone(),
-            f,
-        }));
-        out
+        self.binary_with(a, b, (), move |_, x, y| f(x, y))
     }
 
     pub fn input<K: Ord + Clone + 'static>(&mut self) -> (Input<K>, Stream<Z<K>>) {
         let pending: Rc<RefCell<Z<K>>> = Rc::default();
         let out = Stream(Rc::default());
-        self.nodes.push(Box::new(Source {
+        self.add(Box::new(Source {
             pending: Rc::clone(&pending),
             out: out.clone(),
+            saved: None,
         }));
         (Input(pending, Rc::clone(&self.token)), out)
     }
 
     pub fn output<K: Ord + Clone + 'static>(&mut self, s: &Stream<Z<K>>) -> Output<K> {
         let sink = Output(Rc::default());
-        let acc = Rc::clone(&sink.0);
-        self.unary(s, move |d: &Z<K>| {
-            *acc.borrow_mut() += d;
+        let state = Sink {
+            sum: Rc::clone(&sink.0),
+            log: Log::default(),
+        };
+        self.unary_with(s, state, |sink, d: &Z<K>| {
+            add_logged(&mut sink.sum.borrow_mut(), d, &mut sink.log);
         });
         sink
     }
@@ -273,10 +569,14 @@ impl Circuit {
     /// out(t) = in(t - 1); out(0) is zero.
     pub fn delay<K: Ord + Clone + 'static>(&mut self, s: &Stream<Z<K>>) -> Stream<Z<K>> {
         let out = Stream(Rc::default());
-        self.nodes.push(Box::new(Delay {
+        self.add(Box::new(Delay {
             src: s.clone(),
             out: out.clone(),
-            held: ZSet::new(),
+            state: Held {
+                held: ZSet::new(),
+                open: false,
+                saved: None,
+            },
         }));
         out
     }
@@ -284,10 +584,13 @@ impl Circuit {
     /// out(t) = in(0) + .. + in(t).
     pub fn integrate<K: Ord + Clone + 'static>(&mut self, s: &Stream<Z<K>>) -> Stream<Z<K>> {
         let out = Stream(Rc::default());
-        self.nodes.push(Box::new(Integrate {
+        self.add(Box::new(Integrate {
             src: s.clone(),
             out: out.clone(),
-            acc: ZSet::new(),
+            state: Acc {
+                acc: ZSet::new(),
+                log: Log::default(),
+            },
         }));
         out
     }
@@ -346,8 +649,7 @@ impl Circuit {
 
     /// Incremental `distinct`.
     pub fn distinct<K: Ord + Clone + 'static>(&mut self, s: &Stream<Z<K>>) -> Stream<Z<K>> {
-        let mut st = DistinctState::new();
-        self.unary(s, move |d: &Z<K>| st.step(d))
+        self.unary_with(s, DistinctState::new(), |st, d: &Z<K>| st.step(d))
     }
 
     /// The bilinear delta join over in-memory traces.
@@ -383,8 +685,8 @@ impl Circuit {
         TA: Trace<K, A> + 'static,
         TB: Trace<K, B> + 'static,
     {
-        let mut st = JoinState::new(left, right);
-        self.binary(a, b, move |da: &Z<(K, A)>, db: &Z<(K, B)>| {
+        let st = JoinState::new(Undoable::new(left), Undoable::new(right));
+        self.binary_with(a, b, st, move |st, da: &Z<(K, A)>, db: &Z<(K, B)>| {
             st.step(da, db, &mut f)
         })
     }
@@ -416,8 +718,8 @@ impl Circuit {
         R: Ord + Clone + 'static,
         T: Trace<K, V> + 'static,
     {
-        let mut st = AggregateState::new(groups);
-        self.unary(s, move |d: &Z<(K, V)>| st.step(d, &mut f))
+        let st = AggregateState::new(Undoable::new(groups));
+        self.unary_with(s, st, move |st, d: &Z<(K, V)>| st.step(d, &mut f))
     }
 }
 

crates/zset/src/lib.rs

@@ -91,7 +91,7 @@ impl<K: Ord> ZSet<K> {
         self.adjust(key, weight, false);
     }
 
-    fn adjust(&mut self, key: K, weight: i64, subtract: bool) {
+    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 };

crates/zset/src/ops.rs

@@ -4,15 +4,90 @@
 use crate::{Trace, ZSet, plus, times};
 use std::collections::BTreeMap;
 
+/// State a transaction can undo (`Circuit::begin`): `begin` starts recording, `commit`
+/// forgets the record, `rollback` replays it backwards. The cost is the delta, not the state.
+pub(crate) trait Transactional {
+    fn begin(&mut self);
+    fn commit(&mut self);
+    fn rollback(&mut self);
+}
+
+impl Transactional for () {
+    fn begin(&mut self) {}
+    fn commit(&mut self) {}
+    fn rollback(&mut self) {}
+}
+
+/// An undo log that records only while a transaction is open.
+#[derive(Clone, Debug)]
+pub(crate) struct Log<E> {
+    on: bool,
+    entries: Vec<E>,
+}
+
+impl<E> Default for Log<E> {
+    fn default() -> Self {
+        Log {
+            on: false,
+            entries: Vec::new(),
+        }
+    }
+}
+
+impl<E> Log<E> {
+    pub(crate) fn begin(&mut self) {
+        self.on = true;
+        self.entries.clear();
+    }
+
+    pub(crate) fn commit(&mut self) {
+        self.on = false;
+        self.entries = Vec::new();
+    }
+
+    /// Records `entry()` if a transaction is open; otherwise `entry` is not even called.
+    pub(crate) fn record(&mut self, entry: impl FnOnce() -> E) {
+        if self.on {
+            self.entries.push(entry());
+        }
+    }
+
+    /// Hands the entries to `undo`, newest first, and closes the log.
+    pub(crate) fn rollback(&mut self, mut undo: impl FnMut(E)) {
+        self.on = false;
+        for e in std::mem::take(&mut self.entries).into_iter().rev() {
+            undo(e);
+        }
+    }
+}
+
+/// `z += d`, one key at a time, logging each key that was applied. A key whose addition
+/// overflows panics before it changes anything, so the log is exactly what was applied.
+pub(crate) fn add_logged<K: Ord + Clone>(z: &mut ZSet<K>, d: &ZSet<K>, log: &mut Log<(K, i64)>) {
+    for (k, w) in d.iter() {
+        z.insert(k.clone(), w);
+        log.record(|| (k.clone(), w));
+    }
+}
+
+/// Takes one logged addition back. The result is the weight before it, which was in range.
+pub(crate) fn unadd<K: Ord>(z: &mut ZSet<K>, (k, w): (K, i64)) {
+    z.adjust(k, w, true);
+}
+
 /// Incremental `distinct`: out = distinct(held + d) - distinct(held), per touched key.
 #[derive(Clone, Debug, Default)]
 pub struct DistinctState<K> {
     held: ZSet<K>,
+    log: Log<(K, i64)>,
 }
 
 impl<K: Ord + Clone> DistinctState<K> {
     pub fn new() -> Self {
-        DistinctState { held: ZSet::new() }
+        DistinctState {
+            held: ZSet::new(),
+            log: Log::default(),
+        }
     }
 
     pub fn step(&mut self, d: &ZSet<K>) -> ZSet<K> {
@@ -22,11 +97,81 @@ impl<K: Ord + Clone> DistinctState<K> {
             let after = plus(before, w);
             out.insert(k.clone(), i64::from(after > 0) - i64::from(before > 0));
             self.held.insert(k.clone(), w);
+            self.log.record(|| (k.clone(), w));
         }
         out
     }
 }
 
+impl<K: Ord> Transactional for DistinctState<K> {
+    fn begin(&mut self) {
+        self.log.begin();
+    }
+
+    fn commit(&mut self) {
+        self.log.commit();
+    }
+
+    fn rollback(&mut self) {
+        let held = &mut self.held;
+        self.log.rollback(|e| unadd(held, e));
+    }
+}
+
+/// A trace that logs what it applies, so a transaction can take it back.
+#[derive(Clone, Debug)]
+pub(crate) struct Undoable<T, K, V> {
+    inner: T,
+    log: Log<(K, V, i64)>,
+}
+
+impl<T, K, V> Undoable<T, K, V> {
+    pub(crate) fn new(inner: T) -> Self {
+        Undoable {
+            inner,
+            log: Log::default(),
+        }
+    }
+}
+
+impl<T: Trace<K, V>, K: Clone, V: Clone> Trace<K, V> for Undoable<T, K, V> {
+    fn apply(&mut self, k: &K, v: &V, w: i64) {
+        self.inner.apply(k, v, w);
+        self.log.record(|| (k.clone(), v.clone(), w));
+    }
+
+    fn lookup(&self, k: &K) -> Vec<(V, i64)> {
+        self.inner.lookup(k)
+    }
+
+    fn range(&self, lo: std::ops::Bound<&K>, hi: std::ops::Bound<&K>) -> Vec<(K, V, i64)> {
+        self.inner.range(lo, hi)
+    }
+}
+
+impl<T: Trace<K, V>, K, V> Transactional for Undoable<T, K, V> {
+    fn begin(&mut self) {
+        self.log.begin();
+    }
+
+    fn commit(&mut self) {
+        self.log.commit();
+    }
+
+    fn rollback(&mut self) {
+        let inner = &mut self.inner;
+        self.log.rollback(|(k, v, w)| match w.checked_neg() {
+            Some(n) => inner.apply(&k, &v, n),
+            // -i64::MIN is 2^63: add it as i64::MAX then 1. The weight before the apply
+            // was non-negative, so neither step leaves i64.
+            None => {
+                inner.apply(&k, &v, i64::MAX);
+                inner.apply(&k, &v, 1);
+            }
+        });
+    }
+}
+
 /// The bilinear delta join of two keyed streams:
 /// out = da x b_prev + a_prev x db + da x db, never a rescan of either side.
 #[derive(Clone, Debug, Default)]
@@ -75,6 +220,22 @@ impl<TA, TB> JoinState<TA, TB> {
     }
 }
 
+impl<TA: Transactional, TB: Transactional> Transactional for JoinState<TA, TB> {
+    fn begin(&mut self) {
+        self.left.begin();
+        self.right.begin();
+    }
+
+    fn commit(&mut self) {
+        self.left.commit();
+        self.right.commit();
+    }
+
+    fn rollback(&mut self) {
+        self.left.rollback();
+        self.right.rollback();
+    }
+}
 /// Incremental group-by: `f` folds one key's integrated values to at most one result.
 /// An empty group has no row, and `f` is not asked about it. Only touched keys are re-folded: out = -(k, f(old)) + (k, f(new)).
 #[derive(Clone, Debug, Default)]
@@ -122,3 +283,17 @@ impl<T> AggregateState<T> {
         out
     }
 }
+
+impl<T: Transactional> Transactional for AggregateState<T> {
+    fn begin(&mut self) {
+        self.groups.begin();
+    }
+
+    fn commit(&mut self) {
+        self.groups.commit();
+    }
+
+    fn rollback(&mut self) {
+        self.groups.rollback();
+    }
+}

crates/zset/src/trace.rs

@@ -7,7 +7,8 @@ use std::ops::Bound;
 /// An integrated `ZSet<(K, V)>`, indexed by `K`. Reads return owned values, so an
 /// implementation may page or compress what it holds. See `context/code/zset.md`.
 pub trait Trace<K, V> {
-    /// Adds `w` to the weight of `(k, v)`.
+    /// Adds `w` to the weight of `(k, v)`. A panic must leave the trace as it was: a
+    /// transaction (`Circuit::begin`) logs an apply only once it has returned.
     fn apply(&mut self, k: &K, v: &V, w: i64);
     /// The values at `k` with their non-zero weights, ascending by value.
     fn lookup(&self, k: &K) -> Vec<(V, i64)>;

crates/zset/tests/cost.rs

@@ -81,3 +81,63 @@ fn an_integrate_stream_holds_its_integral_twice_when_the_tick_is_whole() {
     drop((c, input, s, i));
     assert_eq!(live(), base, "[leak] every key is dropped with the circuit");
 }
+
+type E = (Counted, u8);
+
+#[derive(Clone, Copy)]
+enum Mode {
+    Plain,
+    Commit,
+    Rollback,
+}
+
+/// Every stateful operator over `Counted` keys, loaded with `n` keys; then one step of a
+/// two-key delta in `mode`. Returns the key clones that step (and `begin`/`end`) made.
+fn clones_of_one_step(n: u32, mode: Mode) -> usize {
+    let mut c = Circuit::new();
+    let (a, sa) = c.input::<E>();
+    let (b, sb) = c.input::<E>();
+    let d = c.distinct(&sa);
+    let dl = c.delay(&sa);
+    let j = c.join_indexed(&sa, &sb, |k, x, y| (k.0, *x, *y));
+    let g = c.aggregate(&sa, |_, vals| Some(vals.len() as i64));
+    let sinks = (c.output(&d), c.output(&dl), c.output(&j), c.output(&g));
+    let rows =
+        |lo: u32, hi: u32| -> ZSet<E> { (lo..hi).map(|i| ((Counted::new(i), 0), 1)).collect() };
+    a.push(rows(0, n));
+    b.push(rows(0, n));
+    c.step();
+    for z in [&sinks.0, &sinks.1] {
+        z.take();
+    }
+    sinks.2.take();
+    sinks.3.take();
+    // A small step after the load: a `delay` holds the previous tick's input, so a
+    // transaction opened right after a big load would copy that load once.
+    a.push(rows(n, n + 2));
+    b.push(rows(n, n + 2));
+    c.step();
+    let before = clones();
+    if !matches!(mode, Mode::Plain) {
+        c.begin();
+    }
+    a.push(rows(n + 2, n + 4));
+    b.push(rows(n + 2, n + 4));
+    c.step();
+    match mode {
+        Mode::Plain => {}
+        Mode::Commit => c.commit(),
+        Mode::Rollback => c.rollback(),
+    }
+    clones() - before
+}
+
+#[test]
+fn a_transaction_costs_the_keys_a_step_touches_and_not_the_size_of_the_circuit() {
+    for mode in [Mode::Commit, Mode::Rollback] {
+        let extra = |n| clones_of_one_step(n, mode) - clones_of_one_step(n, Mode::Plain);
+        let (small, large) = (extra(10), extra(300));
+        assert!(small > 0, "[cost] the log copies the keys it records");
+        assert_eq!(small, large, "[cost] the overhead grew with the state");
+    }
+}

crates/zset/tests/rollback.proptest-regressions

@@ -0,0 +1,8 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc b611223baeb182175151c758f4f6877b5897e5b067f9e5010d14e20b1c6c4418 # shrinks to pre = [], post = [(ZSet { weights: {} }, ZSet { weights: {} })], carry = ZSet { weights: {} }, bad = (ZSet { weights: {} }, ZSet { weights: {} }), fail = false
+cc 71d785a9a4f55f12c26afef2c65babad2f29c1f2a30becf55936928d570254b4 # shrinks to pre = [], post = [(ZSet { weights: {} }, ZSet { weights: {} })], carry = ZSet { weights: {} }, bad = (ZSet { weights: {} }, ZSet { weights: {} }), fail = false

crates/zset/tests/rollback.rs

@@ -0,0 +1,279 @@
+//! `Circuit::begin`/`rollback`/`commit` (ADR 009 clause 11): a failed step is undone, and
+//! the circuit then behaves as if the step never ran. Evidence: `context/test-evidence/zset.md`.
+
+use proptest::prelude::*;
+use std::panic::{AssertUnwindSafe, catch_unwind};
+use zset::{Circuit, Input, ZSet};
+
+type Z<K> = ZSet<K>;
+type P = (u8, u8);
+
+/// A circuit using every stateful operator, and two closures: every stream's
+/// `current()`, and every sink drained.
+struct Rig {
+    c: Circuit,
+    a: Input<P>,
+    b: Input<P>,
+    view: Box<dyn Fn() -> String>,
+    drain: Box<dyn Fn() -> String>,
+}
+
+fn rig() -> Rig {
+    let mut c = Circuit::new();
+    let (a, sa) = c.input::<P>();
+    let (b, sb) = c.input::<P>();
+    let ia = c.integrate(&sa);
+    let da = c.delay(&sa);
+    let df = c.differentiate(&sa);
+    let ds = c.distinct(&sa);
+    let j = c.join_indexed(&sa, &sb, |k, x, y| (*k, *x, *y));
+    let ij = c.integrate(&j);
+    let ag = c.aggregate(&sa, |_, vals| {
+        Some(vals.iter().fold(0i64, |s, (v, w)| {
+            s.wrapping_add(i64::from(*v).wrapping_mul(*w))
+        }))
+    });
+    let pl = c.plus(&ia, &sb);
+    let sinks = (
+        c.output(&ia),
+        c.output(&da),
+        c.output(&df),
+        c.output(&ds),
+        c.output(&j),
+        c.output(&ij),
+        c.output(&ag),
+        c.output(&pl),
+    );
+    let (s1, s2) = (sa.clone(), sb.clone());
+    let view = move || {
+        format!(
+            "{:?} {:?} {:?} {:?} {:?} {:?} {:?} {:?} {:?} {:?} {:?}",
+            s1.current(),
+            s2.current(),
+            ia.current(),
+            da.current(),
+            df.current(),
+            ds.current(),
+            j.current(),
+            ij.current(),
+            ag.current(),
+            pl.current(),
+            s1.current().len(),
+        )
+    };
+    let drain = move || {
+        format!(
+            "{:?} {:?} {:?} {:?} {:?} {:?} {:?} {:?}",
+            sinks.0.take(),
+            sinks.1.take(),
+            sinks.2.take(),
+            sinks.3.take(),
+            sinks.4.take(),
+            sinks.5.take(),
+            sinks.6.take(),
+            sinks.7.take(),
+        )
+    };
+    Rig {
+        c,
+        a,
+        b,
+        view: Box::new(view),
+        drain: Box::new(drain),
+    }
+}
+
+impl Rig {
+    fn tick(&mut self, a: &Z<P>, b: &Z<P>) {
+        self.a.push(a.clone());
+        self.b.push(b.clone());
+        self.c.step();
+    }
+}
+
+fn delta() -> impl Strategy<Value = Z<P>> {
+    prop::collection::vec(((0u8..4, 0u8..3), -3i64..=3), 0..6).prop_map(|p| p.into_iter().collect())
+}
+
+fn ticks(n: std::ops::Range<usize>) -> impl Strategy<Value = Vec<(Z<P>, Z<P>)>> {
+    prop::collection::vec((delta(), delta()), n)
+}
+
+/// The key that sits near `i64::MAX` once preloaded, so a later `+5` overflows it.
+const BIG: P = (9, 9);
+
+fn poison() -> Z<P> {
+    // `(0, 0)` sorts first, so operators have applied it before `BIG` overflows.
+    [((0, 0), 1), (BIG, 5)].into_iter().collect()
+}
+
+proptest! {
+    /// A step that fails half way, then a rollback, leaves the circuit as it was; the ticks
+    /// after run as if the failed step had never been attempted.
+    #[test]
+    fn a_rolled_back_step_leaves_the_circuit_as_it_was_and_the_ticks_after_unchanged(
+        pre in ticks(0..4),
+        post in ticks(1..4),
+        carry in delta(),
+        bad in (delta(), delta()),
+        fail in any::<bool>(),
+    ) {
+        let (mut x, mut y) = (rig(), rig());
+        let preload: Z<P> = [(BIG, i64::MAX - 1)].into_iter().collect();
+        for (a, b) in pre.iter().chain([&(preload, Z::new())]) {
+            x.tick(a, b);
+            y.tick(a, b);
+            prop_assert_eq!((x.view)(), (y.view)(), "[pre] state");
+            prop_assert_eq!((x.drain)(), (y.drain)(), "[pre] output");
+        }
+        x.a.push(carry.clone());
+        y.a.push(carry);
+        let (view_before, drain_before) = ((x.view)(), (x.drain)());
+
+        x.c.begin();
+        let mut a = bad.0.clone();
+        if fail {
+            a += poison();
+        }
+        x.a.push(a);
+        x.b.push(bad.1.clone());
+        let outcome = catch_unwind(AssertUnwindSafe(|| x.c.step()));
+        prop_assert_eq!(outcome.is_err(), fail, "[panic] the poison step panics, the others do not");
+        x.c.rollback();
+
+        prop_assert_eq!((x.view)(), view_before, "[rollback] every stream's current()");
+        prop_assert_eq!((x.drain)(), drain_before, "[rollback] every sink");
+        for (i, (a, b)) in post.iter().enumerate() {
+            x.tick(a, b);
+            y.tick(a, b);
+            prop_assert_eq!((x.view)(), (y.view)(), "[post] state, tick {}", i);
+            prop_assert_eq!((x.drain)(), (y.drain)(), "[post] output, tick {}", i);
+        }
+    }
+
+    /// `commit` keeps what the steps did; a transaction that commits is not observable.
+    #[test]
+    fn a_committed_transaction_equals_the_same_steps_outside_one(
+        steps in ticks(1..6),
+        split in 0usize..6,
+    ) {
+        let (mut x, mut y) = (rig(), rig());
+        for (i, (a, b)) in steps.iter().enumerate() {
+            let open = i >= split;
+            if open {
+                x.c.begin();
+            }
+            x.tick(a, b);
+            if open {
+                x.c.commit();
+            }
+            y.tick(a, b);
+            prop_assert_eq!((x.view)(), (y.view)(), "[commit] state, tick {}", i);
+            prop_assert_eq!((x.drain)(), (y.drain)(), "[commit] output, tick {}", i);
+        }
+    }
+}
+
+#[test]
+fn an_operator_built_inside_a_transaction_goes_back_to_empty() {
+    let mut c = Circuit::new();
+    let (input, s) = c.input::<u8>();
+    c.begin();
+    let i = c.integrate(&s);
+    input.push([(1u8, 1)].into_iter().collect());
+    c.step();
+    c.rollback();
+    input.push([(2u8, 1)].into_iter().collect());
+    c.step();
+    let want: Z<u8> = [(2u8, 1)].into_iter().collect();
+    assert_eq!(
+        i.current(),
+        want,
+        "[built] the integral forgot the rolled-back tick"
+    );
+}
+
+/// One key at `MAX` that a transaction lowers by one and raises again. Undone oldest first,
+/// the first undo would take it to `MAX + 1`; newest first, nothing leaves `i64`.
+#[test]
+fn rollback_undoes_newest_first_so_no_undo_leaves_i64() {
+    let mut c = Circuit::new();
+    let (input, s) = c.input::<P>();
+    let (i, d) = (c.integrate(&s), c.distinct(&s));
+    let g = c.aggregate(&s, |_, vals| Some(vals.len() as i64));
+    let one = |w: i64| -> Z<P> { [((1, 1), w)].into_iter().collect() };
+    input.push(one(i64::MAX));
+    c.step();
+    let before = (i.current(), d.current(), g.current());
+    c.begin();
+    for w in [-1, 1] {
+        input.push(one(w));
+        c.step();
+    }
+    c.rollback();
+    assert_eq!(
+        (i.current(), d.current(), g.current()),
+        before,
+        "[order] streams"
+    );
+    input.push(one(-1));
+    c.step();
+    assert_eq!(
+        i.current(),
+        one(i64::MAX - 1),
+        "[order] the integral is back at MAX first"
+    );
+}
+
+/// `i64::MIN` has no negation: a trace undoes it in two applications.
+#[test]
+fn a_minimum_weight_applied_to_every_kind_of_state_is_taken_back() {
+    let mut x = rig();
+    let mut y = rig();
+    let min: Z<P> = [((1, 1), i64::MIN)].into_iter().collect();
+    x.c.begin();
+    x.tick(&min, &Z::new());
+    x.c.rollback();
+    let right: Z<P> = [((1, 2), 1)].into_iter().collect();
+    x.tick(&Z::new(), &right);
+    y.tick(&Z::new(), &right);
+    assert_eq!((x.view)(), (y.view)(), "[min] state");
+    assert_eq!(
+        (x.drain)(),
+        (y.drain)(),
+        "[min] output: the join saw no left side"
+    );
+}
+
+#[test]
+fn rollback_undoes_the_clock() {
+    let mut c = Circuit::new();
+    let (input, _s) = c.input::<u8>();
+    c.step();
+    c.begin();
+    input.push([(1u8, 1)].into_iter().collect());
+    c.step();
+    assert_eq!(c.ticks(), 2);
+    c.rollback();
+    assert_eq!(c.ticks(), 1, "[ticks] the rolled-back step is not counted");
+}
+
+#[test]
+#[should_panic(expected = "already open")]
+fn a_second_begin_is_refused_loudly() {
+    let mut c = Circuit::new();
+    c.begin();
+    c.begin();
+}
+
+#[test]
+#[should_panic(expected = "no transaction is open")]
+fn a_rollback_with_nothing_open_is_refused_loudly() {
+    Circuit::new().rollback();
+}
+
+#[test]
+#[should_panic(expected = "no transaction is open")]
+fn a_commit_with_nothing_open_is_refused_loudly() {
+    Circuit::new().commit();
+}