improvise/crates/zset/src/circuit.rs

//! The synchronous circuit: `Send`, never `Sync`, no locks; it owns every edge's value in an
//! arena and its handles are ids. See `context/code/zset.md`, "Threads".

use crate::ops::{
    AggregateState, DistinctState, JoinState, Log, Transactional, Undoable, add_logged, unadd,
};
use crate::{MemTrace, Trace, ZSet};
use std::any::Any;
use std::marker::PhantomData;
use std::sync::atomic::{AtomicU64, Ordering};

/// What an edge carries: the latest piece, and the whole tick's value, kept only while a
/// tick is split (`Circuit::step_continue`); `None` means the piece is the whole.
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> {
    fn default() -> Self {
        Slot {
            piece: T::default(),
            whole: None,
            saved: None,
            open: false,
        }
    }
}

/// Values whose pieces of one tick add up to the tick's value.
pub trait Piecewise: Default + Clone + Send + 'static {
    fn absorb(&mut self, piece: &Self);
}

impl Piecewise for () {
    fn absorb(&mut self, _: &Self) {}
}

impl<K: Ord + Clone + Send + 'static> Piecewise for ZSet<K> {
    fn absorb(&mut self, piece: &Self) {
        *self += piece;
    }
}

impl<T: Piecewise> Slot<T> {
    fn publish(&mut self, v: T, cont: bool) {
        if cont {
            if self.open && self.saved.is_none() {
                self.saved = Some((self.piece.clone(), self.whole.clone()));
            }
            let mut whole = self.whole.take().unwrap_or_else(|| self.piece.clone());
            whole.absorb(&v);
            self.whole = Some(whole);
            self.piece = v;
        } else {
            let piece = std::mem::replace(&mut self.piece, v);
            let whole = self.whole.take();
            if self.open && self.saved.is_none() {
                self.saved = Some((piece, whole));
            }
        }
    }
}

impl<T> Slot<T> {
    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((piece, whole)) = self.saved.take() {
            self.piece = piece;
            self.whole = whole;
        }
        self.open = false;
    }
}

/// The index of one value in a circuit's arena, typed so a read cannot name the wrong type.
struct Id<T>(usize, PhantomData<fn(T) -> T>);

impl<T> Clone for Id<T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<T> Copy for Id<T> {}

/// Every edge's value, every input's pending delta and every output's sum, by id.
#[derive(Default)]
struct Arena(Vec<Box<dyn Any + Send>>);

impl Arena {
    fn alloc<T: Send + 'static>(&mut self, v: T) -> Id<T> {
        self.0.push(Box::new(v));
        Id(self.0.len() - 1, PhantomData)
    }

    fn get<T: 'static>(&self, id: Id<T>) -> &T {
        self.0[id.0].downcast_ref().expect("zset: arena type")
    }

    fn get_mut<T: 'static>(&mut self, id: Id<T>) -> &mut T {
        self.0[id.0].downcast_mut().expect("zset: arena type")
    }

    /// One value to read and another, allocated after it, to write.
    fn read_write<R: 'static, W: 'static>(&mut self, r: Id<R>, w: Id<W>) -> (&R, &mut W) {
        let (lo, hi) = self.0.split_at_mut(w.0);
        assert!(
            r.0 < w.0,
            "zset: the read entry comes before the written one"
        );
        let ty = "zset: arena type";
        (
            lo[r.0].downcast_ref().expect(ty),
            hi[0].downcast_mut().expect(ty),
        )
    }
}

/// Circuit identities; a handle remembers its circuit so a foreign one is refused.
static NEXT_CIRCUIT: AtomicU64 = AtomicU64::new(1);

type Edge<T> = Id<Slot<T>>;

/// The value an edge carries during the current step. Read it with `Circuit::current`.
pub struct Stream<T> {
    id: Edge<T>,
    circuit: u64,
}

impl<T> Clone for Stream<T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<T> Copy for Stream<T> {}

/// Where a caller feeds deltas, with `Circuit::push`; they enter the circuit at the next
/// `step`. It remembers which circuit made it, so a `Rebuild` can refuse another.
pub struct Input<K> {
    pending: Id<ZSet<K>>,
    circuit: u64,
}

impl<K> Clone for Input<K> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<K> Copy for Input<K> {}

/// A sink that sums every step's delta until `Circuit::take`. Take inside no transaction: a
/// rollback subtracts the step's output, which a `take` in between has already removed.
pub struct Output<K> {
    sum: Id<ZSet<K>>,
    circuit: u64,
}

impl<K> Clone for Output<K> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<K> Copy for Output<K> {}

trait Node: Send {
    /// `cont` is true for a later piece of a tick that has already begun.
    fn eval(&mut self, a: &mut Arena, cont: bool);
    fn begin(&mut self, a: &mut Arena);
    fn commit(&mut self, a: &mut Arena);
    fn rollback(&mut self, a: &mut Arena);
}

/// One input, one output, and the state `f` may change.
struct Unary<I, O, S, F> {
    src: Edge<I>,
    out: Edge<O>,
    state: S,
    f: F,
}

impl<I, O, S, F> Node for Unary<I, O, S, F>
where
    I: Piecewise,
    O: Piecewise,
    S: Transactional + Send,
    F: FnMut(&mut S, &I) -> O + Send,
{
    fn eval(&mut self, a: &mut Arena, cont: bool) {
        let v = (self.f)(&mut self.state, &a.get(self.src).piece);
        a.get_mut(self.out).publish(v, cont);
    }

    fn begin(&mut self, a: &mut Arena) {
        self.state.begin();
        a.get_mut(self.out).begin();
    }

    fn commit(&mut self, a: &mut Arena) {
        self.state.commit();
        a.get_mut(self.out).commit();
    }

    fn rollback(&mut self, a: &mut Arena) {
        self.state.rollback();
        a.get_mut(self.out).rollback();
    }
}

struct Binary<A, B, O, S, F> {
    a: Edge<A>,
    b: Edge<B>,
    out: Edge<O>,
    state: S,
    f: F,
}

impl<A, B, O, S, F> Node for Binary<A, B, O, S, F>
where
    A: Piecewise,
    B: Piecewise,
    O: Piecewise,
    S: Transactional + Send,
    F: FnMut(&mut S, &A, &B) -> O + Send,
{
    fn eval(&mut self, ar: &mut Arena, cont: bool) {
        let v = (self.f)(
            &mut self.state,
            &ar.get(self.a).piece,
            &ar.get(self.b).piece,
        );
        ar.get_mut(self.out).publish(v, cont);
    }

    fn begin(&mut self, a: &mut Arena) {
        self.state.begin();
        a.get_mut(self.out).begin();
    }

    fn commit(&mut self, a: &mut Arena) {
        self.state.commit();
        a.get_mut(self.out).commit();
    }

    fn rollback(&mut self, a: &mut Arena) {
        self.state.rollback();
        a.get_mut(self.out).rollback();
    }
}

/// 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: Edge<ZSet<K>>,
    out: Edge<ZSet<K>>,
    state: Held<K>,
}

impl<K: Ord + Clone + Send + 'static> Node for Delay<K> {
    fn eval(&mut self, a: &mut Arena, cont: bool) {
        let d = &a.get(self.src).piece;
        self.state.save();
        let v = if cont {
            self.state.held += d;
            ZSet::new()
        } else {
            std::mem::replace(&mut self.state.held, d.clone())
        };
        a.get_mut(self.out).publish(v, cont);
    }

    fn begin(&mut self, a: &mut Arena) {
        self.state.begin();
        a.get_mut(self.out).begin();
    }

    fn commit(&mut self, a: &mut Arena) {
        self.state.commit();
        a.get_mut(self.out).commit();
    }

    fn rollback(&mut self, a: &mut Arena) {
        self.state.rollback();
        a.get_mut(self.out).rollback();
    }
}

/// 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: Edge<ZSet<K>>,
    out: Edge<ZSet<K>>,
    state: Acc<K>,
}

impl<K: Ord + Clone + Send + 'static> Node for Integrate<K> {
    fn eval(&mut self, a: &mut Arena, cont: bool) {
        let d = &a.get(self.src).piece;
        add_logged(&mut self.state.acc, d, &mut self.state.log);
        let v = if cont {
            d.clone()
        } else {
            self.state.acc.clone()
        };
        a.get_mut(self.out).publish(v, cont);
    }

    fn begin(&mut self, a: &mut Arena) {
        self.state.begin();
        a.get_mut(self.out).begin();
    }

    fn commit(&mut self, a: &mut Arena) {
        self.state.commit();
        a.get_mut(self.out).commit();
    }

    fn rollback(&mut self, a: &mut Arena) {
        self.state.rollback();
        a.get_mut(self.out).rollback();
    }
}

/// 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: Id<ZSet<K>>,
    out: Edge<ZSet<K>>,
    saved: Option<ZSet<K>>,
}

impl<K: Ord + Clone + Send + 'static> Node for Source<K> {
    fn eval(&mut self, a: &mut Arena, cont: bool) {
        let v = std::mem::take(a.get_mut(self.pending));
        a.get_mut(self.out).publish(v, cont);
    }

    fn begin(&mut self, a: &mut Arena) {
        self.saved = Some(a.get(self.pending).clone());
        a.get_mut(self.out).begin();
    }

    fn commit(&mut self, a: &mut Arena) {
        self.saved = None;
        a.get_mut(self.out).commit();
    }

    fn rollback(&mut self, a: &mut Arena) {
        if let Some(saved) = self.saved.take() {
            *a.get_mut(self.pending) = saved;
        }
        a.get_mut(self.out).rollback();
    }
}

/// Adds each tick to the output's sum, with the log that takes a transaction's additions
/// back. Its own edge carries nothing.
struct Sink<K> {
    src: Edge<ZSet<K>>,
    out: Edge<()>,
    sum: Id<ZSet<K>>,
    log: Log<(K, i64)>,
}

impl<K: Ord + Clone + Send + 'static> Node for Sink<K> {
    fn eval(&mut self, a: &mut Arena, cont: bool) {
        let (src, sum) = a.read_write(self.src, self.sum);
        add_logged(sum, &src.piece, &mut self.log);
        a.get_mut(self.out).publish((), cont);
    }

    fn begin(&mut self, a: &mut Arena) {
        self.log.begin();
        a.get_mut(self.out).begin();
    }

    fn commit(&mut self, a: &mut Arena) {
        self.log.commit();
        a.get_mut(self.out).commit();
    }

    fn rollback(&mut self, a: &mut Arena) {
        let sum = a.get_mut(self.sum);
        self.log.rollback(|e| unadd(sum, e));
        a.get_mut(self.out).rollback();
    }
}

/// Nodes run in creation order, a topological order: a node is built from existing streams.
/// `Send` (`tests/threads.rs`) and, having no lock, not `Sync`.
pub struct Circuit {
    nodes: Vec<Box<dyn Node>>,
    arena: Arena,
    ticks: u64,
    /// This circuit's identity: every handle it makes holds it.
    token: u64,
    /// `Some(ticks)` as of `begin`, while a transaction is open.
    open: Option<u64>,
    /// A step or piece ran since `begin`.
    stepped: bool,
    /// Rollbacks that undid a step or piece. It never goes down, unlike `ticks`, so a
    /// `Rebuild` can tell that a rollback took some of its chunks.
    undone: u64,
}

impl Default for Circuit {
    fn default() -> Self {
        Circuit {
            nodes: Vec::new(),
            arena: Arena::default(),
            ticks: 0,
            token: NEXT_CIRCUIT.fetch_add(1, Ordering::Relaxed),
            open: None,
            stepped: false,
            undone: 0,
        }
    }
}

type Z<K> = ZSet<K>;

impl Circuit {
    pub fn new() -> Self {
        Self::default()
    }

    /// Runs one clock tick: every node, once, in order.
    pub fn step(&mut self) {
        self.ticks += 1;
        self.run(false);
    }

    /// How many ticks `step` has begun. A `Rebuild` uses it to see a tick split.
    pub fn ticks(&self) -> u64 {
        self.ticks
    }

    /// Runs another piece of the tick the last `step` began; the clock does not advance.
    /// Pieces sum to one `step` of their combined input (`context/code/zset.md`).
    pub fn step_continue(&mut self) {
        self.run(true);
    }

    /// Opens a transaction: state changes are logged until `commit` or `rollback`, at the
    /// cost of the keys touched. Panics if one is 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);
        self.stepped = false;
        for n in &mut self.nodes {
            n.begin(&mut self.arena);
        }
    }

    /// 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(&mut self.arena);
        }
    }

    /// Puts the state back to what it was at `begin`, even after a step panicked half way;
    /// pushes since are dropped. Panics if none 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(&mut self.arena);
        }
        self.ticks = ticks;
        self.undone += u64::from(self.stepped);
        self.stepped = false;
    }

    fn run(&mut self, cont: bool) {
        self.stepped |= self.open.is_some();
        for n in &mut self.nodes {
            n.eval(&mut self.arena, cont);
        }
    }

    #[track_caller]
    fn mine(&self, circuit: u64) {
        assert_eq!(
            circuit, self.token,
            "zset: a handle from another circuit was used here"
        );
    }

    /// 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(&mut self.arena);
        }
        self.nodes.push(node);
    }

    fn edge<T: Piecewise>(&mut self) -> Stream<T> {
        let id = self.arena.alloc(Slot::<T>::default());
        Stream {
            id,
            circuit: self.token,
        }
    }

    /// The value of the last tick, all its pieces together (the zero before the first).
    pub fn current<T: Clone + Send + 'static>(&self, s: &Stream<T>) -> T {
        self.mine(s.circuit);
        let slot = self.arena.get(s.id);
        slot.whole.as_ref().unwrap_or(&slot.piece).clone()
    }

    /// Feeds a delta to `input`; it enters at the next `step`.
    pub fn push<K: Ord + Send + 'static>(&mut self, input: &Input<K>, delta: ZSet<K>) {
        self.mine(input.circuit);
        *self.arena.get_mut(input.pending) += delta;
    }

    /// Takes the sum of the deltas `output` has seen since the last take.
    pub fn take<K: Ord + Send + 'static>(&mut self, output: &Output<K>) -> ZSet<K> {
        self.mine(output.circuit);
        std::mem::take(self.arena.get_mut(output.sum))
    }

    fn unary_with<I: Piecewise, O: Piecewise, S: Transactional + Send + 'static>(
        &mut self,
        src: &Stream<I>,
        state: S,
        f: impl FnMut(&mut S, &I) -> O + Send + 'static,
    ) -> Stream<O> {
        self.mine(src.circuit);
        let out = self.edge::<O>();
        self.add(Box::new(Unary {
            src: src.id,
            out: out.id,
            state,
            f,
        }));
        out
    }

    fn unary<I: Piecewise, O: Piecewise>(
        &mut self,
        src: &Stream<I>,
        mut f: impl FnMut(&I) -> O + Send + 'static,
    ) -> Stream<O> {
        self.unary_with(src, (), move |_, d| f(d))
    }

    fn binary_with<A: Piecewise, B: Piecewise, O: Piecewise, S: Transactional + Send + 'static>(
        &mut self,
        a: &Stream<A>,
        b: &Stream<B>,
        state: S,
        f: impl FnMut(&mut S, &A, &B) -> O + Send + 'static,
    ) -> Stream<O> {
        self.mine(a.circuit);
        self.mine(b.circuit);
        let out = self.edge::<O>();
        self.add(Box::new(Binary {
            a: a.id,
            b: b.id,
            out: out.id,
            state,
            f,
        }));
        out
    }

    fn binary<A: Piecewise, B: Piecewise, O: Piecewise>(
        &mut self,
        a: &Stream<A>,
        b: &Stream<B>,
        mut f: impl FnMut(&A, &B) -> O + Send + 'static,
    ) -> Stream<O> {
        self.binary_with(a, b, (), move |_, x, y| f(x, y))
    }

    pub fn input<K: Ord + Clone + Send + 'static>(&mut self) -> (Input<K>, Stream<Z<K>>) {
        let pending = self.arena.alloc(Z::<K>::default());
        let out = self.edge::<Z<K>>();
        self.add(Box::new(Source {
            pending,
            out: out.id,
            saved: None,
        }));
        (
            Input {
                pending,
                circuit: self.token,
            },
            out,
        )
    }

    pub fn output<K: Ord + Clone + Send + 'static>(&mut self, s: &Stream<Z<K>>) -> Output<K> {
        self.mine(s.circuit);
        let sum = self.arena.alloc(Z::<K>::default());
        let out = self.edge::<()>();
        self.add(Box::new(Sink {
            src: s.id,
            out: out.id,
            sum,
            log: Log::default(),
        }));
        Output {
            sum,
            circuit: self.token,
        }
    }

    /// out(t) = in(t - 1); out(0) is zero.
    pub fn delay<K: Ord + Clone + Send + 'static>(&mut self, s: &Stream<Z<K>>) -> Stream<Z<K>> {
        self.mine(s.circuit);
        let out = self.edge::<Z<K>>();
        self.add(Box::new(Delay {
            src: s.id,
            out: out.id,
            state: Held {
                held: ZSet::new(),
                open: false,
                saved: None,
            },
        }));
        out
    }

    /// out(t) = in(0) + .. + in(t).
    pub fn integrate<K: Ord + Clone + Send + 'static>(&mut self, s: &Stream<Z<K>>) -> Stream<Z<K>> {
        self.mine(s.circuit);
        let out = self.edge::<Z<K>>();
        self.add(Box::new(Integrate {
            src: s.id,
            out: out.id,
            state: Acc {
                acc: ZSet::new(),
                log: Log::default(),
            },
        }));
        out
    }

    /// out(t) = in(t) - in(t - 1).
    pub fn differentiate<K: Ord + Clone + Send + 'static>(
        &mut self,
        s: &Stream<Z<K>>,
    ) -> Stream<Z<K>> {
        let prev = self.delay(s);
        self.binary(s, &prev, |a: &Z<K>, b: &Z<K>| a - b)
    }

    pub fn plus<K: Ord + Clone + Send + 'static>(
        &mut self,
        a: &Stream<Z<K>>,
        b: &Stream<Z<K>>,
    ) -> Stream<Z<K>> {
        self.binary(a, b, |a: &Z<K>, b: &Z<K>| a + b)
    }

    pub fn map<K: Ord + Clone + Send + 'static, L: Ord + Clone + Send + 'static>(
        &mut self,
        s: &Stream<Z<K>>,
        mut f: impl FnMut(&K) -> L + Send + 'static,
    ) -> Stream<Z<L>> {
        self.unary(s, move |d: &Z<K>| d.map(&mut f))
    }

    pub fn filter<K: Ord + Clone + Send + 'static>(
        &mut self,
        s: &Stream<Z<K>>,
        mut keep: impl FnMut(&K) -> bool + Send + 'static,
    ) -> Stream<Z<K>> {
        self.unary(s, move |d: &Z<K>| d.filter(&mut keep))
    }

    /// Each key becomes any number of images, each with the key's weight.
    pub fn flat_map<
        K: Ord + Clone + Send + 'static,
        L: Ord + Clone + Send + 'static,
        I: IntoIterator<Item = L>,
    >(
        &mut self,
        s: &Stream<Z<K>>,
        mut f: impl FnMut(&K) -> I + Send + 'static,
    ) -> Stream<Z<L>> {
        self.unary(s, move |d: &Z<K>| d.flat_map(&mut f))
    }

    /// Keys each element by `key`, giving the `(key, element)` pairs a join reads.
    pub fn index_by<K: Ord + Clone + Send + 'static, G: Ord + Clone + Send + 'static>(
        &mut self,
        s: &Stream<Z<K>>,
        mut key: impl FnMut(&K) -> G + Send + 'static,
    ) -> Stream<Z<(G, K)>> {
        self.unary(s, move |d: &Z<K>| d.map(|k| (key(k), k.clone())))
    }

    /// Incremental `distinct`.
    pub fn distinct<K: Ord + Clone + Send + 'static>(&mut self, s: &Stream<Z<K>>) -> Stream<Z<K>> {
        self.unary_with(s, DistinctState::new(), |st, d: &Z<K>| st.step(d))
    }

    /// The bilinear delta join over in-memory traces.
    pub fn join_indexed<K, A, B, O>(
        &mut self,
        a: &Stream<Z<(K, A)>>,
        b: &Stream<Z<(K, B)>>,
        f: impl FnMut(&K, &A, &B) -> O + Send + 'static,
    ) -> Stream<Z<O>>
    where
        K: Ord + Clone + Send + 'static,
        A: Ord + Clone + Send + 'static,
        B: Ord + Clone + Send + 'static,
        O: Ord + Clone + Send + 'static,
    {
        self.join_indexed_in(a, b, MemTrace::new(), MemTrace::new(), f)
    }

    /// `join_indexed` over traces the caller chose.
    pub fn join_indexed_in<K, A, B, O, TA, TB>(
        &mut self,
        a: &Stream<Z<(K, A)>>,
        b: &Stream<Z<(K, B)>>,
        left: TA,
        right: TB,
        mut f: impl FnMut(&K, &A, &B) -> O + Send + 'static,
    ) -> Stream<Z<O>>
    where
        K: Ord + Clone + Send + 'static,
        A: Ord + Clone + Send + 'static,
        B: Ord + Clone + Send + 'static,
        O: Ord + Clone + Send + 'static,
        TA: Trace<K, A> + Send + 'static,
        TB: Trace<K, B> + Send + 'static,
    {
        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)
        })
    }

    /// Incremental group-by over in-memory state.
    pub fn aggregate<K, V, R>(
        &mut self,
        s: &Stream<Z<(K, V)>>,
        f: impl FnMut(&K, &[(V, i64)]) -> Option<R> + Send + 'static,
    ) -> Stream<Z<(K, R)>>
    where
        K: Ord + Clone + Send + 'static,
        V: Ord + Clone + Send + 'static,
        R: Ord + Clone + Send + 'static,
    {
        self.aggregate_in(s, MemTrace::new(), f)
    }

    /// `aggregate` over a trace the caller chose.
    pub fn aggregate_in<K, V, R, T>(
        &mut self,
        s: &Stream<Z<(K, V)>>,
        groups: T,
        mut f: impl FnMut(&K, &[(V, i64)]) -> Option<R> + Send + 'static,
    ) -> Stream<Z<(K, R)>>
    where
        K: Ord + Clone + Send + 'static,
        V: Ord + Clone + Send + 'static,
        R: Ord + Clone + Send + 'static,
        T: Trace<K, V> + Send + 'static,
    {
        let st = AggregateState::new(Undoable::new(groups));
        self.unary_with(s, st, move |st, d: &Z<(K, V)>| st.step(d, &mut f))
    }
}

/// Whether a resumable job finished.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Progress {
    Pending,
    Done,
}

/// A snapshot load as a sliced job: one logical tick (`context/code/zset.md`). A pending
/// job owns the clock: a `Circuit::step` before it is Done makes `step` panic, for good.
pub struct Rebuild<K> {
    input: Input<K>,
    rest: std::collections::btree_map::IntoIter<K, i64>,
    chunk: usize,
    /// The circuit's tick count right after this job's first chunk, once it has run one.
    began_at: Option<u64>,
    /// The circuit's `undone` count when this job first ran, even if that chunk panicked.
    undone_at: Option<u64>,
}

impl<K: Ord + Send + 'static> Rebuild<K> {
    /// `chunk` keys per circuit step; zero is read as one.
    pub fn new(input: Input<K>, snapshot: ZSet<K>, chunk: usize) -> Self {
        Rebuild {
            input,
            rest: snapshot.into_iter(),
            chunk: chunk.max(1),
            began_at: None,
            undone_at: None,
        }
    }

    /// Whether `circuit` is the one that made this job's input.
    pub fn belongs_to(&self, circuit: &Circuit) -> bool {
        self.input.circuit == circuit.token
    }

    /// Whether a `Circuit::step` split this job's tick, or a rollback undid a step after it
    /// first ran (finished or not). An interrupted job never finishes: start another.
    pub fn interrupted(&self, circuit: &Circuit) -> bool {
        let split = self.rest.len() > 0 && self.began_at.is_some_and(|t| t != circuit.ticks());
        split || self.undone_at.is_some_and(|u| u != circuit.undone)
    }

    /// Runs chunks until done or `expired()`, at least one per call. A predicate, not an
    /// `Instant` (wasm32). An empty snapshot takes its one tick; a finished job none.
    pub fn step(&mut self, circuit: &mut Circuit, mut expired: impl FnMut() -> bool) -> Progress {
        assert!(
            self.belongs_to(circuit),
            "zset: the rebuild's input belongs to another circuit"
        );
        self.undone_at.get_or_insert(circuit.undone);
        assert!(
            !self.interrupted(circuit),
            "zset: the rebuild was interrupted by a Circuit::step or a rollback before it finished"
        );
        loop {
            let chunk: ZSet<K> = self.rest.by_ref().take(self.chunk).collect();
            if chunk.is_zero() && self.began_at.is_some() {
                return Progress::Done;
            }
            circuit.stepped |= circuit.open.is_some();
            circuit.push(&self.input, chunk);
            if self.began_at.is_some() {
                circuit.step_continue();
            } else {
                circuit.step();
                self.began_at = Some(circuit.ticks());
            }
            if self.rest.len() == 0 {
                return Progress::Done;
            }
            if expired() {
                return Progress::Pending;
            }
        }
    }
}