Aura/crates/aura-std/src/gated_recorder.rs

//! `GatedRecorder` — a `Recorder` (C8 sink) that emits a recorded row only on
//! cycles where a designated `bool` gate column is true, plus the genuine final
//! row once on `finalize`. Fed a dense per-decision record gated on its
//! "closed-this-cycle" column, it yields exactly the rows a trade-ledger
//! reduction reads — the closed rows and the last row — in O(trades) memory
//! instead of O(cycles). C7-pure: it holds only owned state + the channel.

use aura_core::{Cell, Ctx, Firing, Node, NodeSchema, PortSpec, PrimitiveBuilder, Scalar, ScalarKind, Timestamp};
use std::sync::mpsc::Sender;

/// A recording sink that emits only rows whose `gate_col` (a `bool` column) is
/// true, plus the final row once on `finalize` (deduped against the last gated
/// emit). Warm-up behaves like `Recorder`: `None` until every column is warm,
/// recording nothing.
pub struct GatedRecorder {
    kinds: Vec<ScalarKind>,
    gate_col: usize,
    last: Option<(Timestamp, Vec<Scalar>)>,
    last_emitted: bool,
    tx: Sender<(Timestamp, Vec<Scalar>)>,
}

impl GatedRecorder {
    pub fn new(kinds: &[ScalarKind], gate_col: usize, tx: Sender<(Timestamp, Vec<Scalar>)>) -> Self {
        Self { kinds: kinds.to_vec(), gate_col, last: None, last_emitted: false, tx }
    }

    /// Param-generic recipe. `kinds` / `gate_col` / `tx` are non-param construction
    /// args (captured); the node declares no params. One input port per kind.
    pub fn builder(
        kinds: Vec<ScalarKind>,
        gate_col: usize,
        firing: Firing,
        tx: Sender<(Timestamp, Vec<Scalar>)>,
    ) -> PrimitiveBuilder {
        let inputs = kinds
            .iter()
            .enumerate()
            .map(|(i, &kind)| PortSpec { kind, firing, name: format!("col[{i}]") })
            .collect();
        let build_kinds = kinds.clone();
        PrimitiveBuilder::new(
            "GatedRecorder",
            NodeSchema {
                inputs,
                output: vec![],
                params: vec![],
                doc: "recording sink that persists its input while the gate is true, plus a final row at finalize",
            }, // sink: empty output (C8)
            move |_| Box::new(GatedRecorder::new(&build_kinds, gate_col, tx.clone())),
        )
    }
}

impl Node for GatedRecorder {
    fn lookbacks(&self) -> Vec<usize> {
        vec![1; self.kinds.len()]
    }

    fn eval(&mut self, ctx: Ctx<'_>) -> Option<&[Cell]> {
        let mut row = Vec::with_capacity(self.kinds.len());
        for (i, &kind) in self.kinds.iter().enumerate() {
            let scalar = match kind {
                ScalarKind::F64 => Scalar::f64(ctx.f64_in(i).get(0)?),
                ScalarKind::I64 => Scalar::i64(ctx.i64_in(i).get(0)?),
                ScalarKind::Bool => Scalar::bool(ctx.bool_in(i).get(0)?),
                ScalarKind::Timestamp => Scalar::ts(ctx.ts_in(i).get(0)?),
            };
            row.push(scalar);
        }
        let gated = row[self.gate_col].as_bool();
        if gated {
            let _ = self.tx.send((ctx.now(), row.clone()));
        }
        self.last = Some((ctx.now(), row));
        self.last_emitted = gated;
        None
    }

    fn finalize(&mut self) {
        // ensure a downstream `summarize_r` sees the true final row (e.g. an
        // open-at-end position) without double-counting a final row already
        // emitted as a closed trade.
        if !self.last_emitted
            && let Some(row) = self.last.take()
        {
            let _ = self.tx.send(row);
        }
    }

    fn label(&self) -> String {
        "GatedRecorder".to_string()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use aura_core::AnyColumn;
    use std::sync::mpsc;

    fn drive(seq: &[(bool, f64)]) -> Vec<(Timestamp, Vec<Scalar>)> {
        let (tx, rx) = mpsc::channel();
        let mut g = GatedRecorder::new(&[ScalarKind::Bool, ScalarKind::F64], 0, tx);
        let mut inputs = vec![
            AnyColumn::with_capacity(ScalarKind::Bool, 1),
            AnyColumn::with_capacity(ScalarKind::F64, 1),
        ];
        for (t, &(gate, payload)) in seq.iter().enumerate() {
            inputs[0].push(Scalar::bool(gate)).unwrap();
            inputs[1].push(Scalar::f64(payload)).unwrap();
            assert_eq!(g.eval(Ctx::new(&inputs, Timestamp(t as i64 + 1))), None);
        }
        g.finalize();
        rx.try_iter().collect()
    }

    #[test]
    fn emits_gated_rows_and_flushes_a_nongated_final_row() {
        // hold, closed, hold(final). emits the closed row, then flushes the final hold.
        let out = drive(&[(false, 1.0), (true, 2.0), (false, 3.0)]);
        assert_eq!(
            out,
            vec![
                (Timestamp(2), vec![Scalar::bool(true), Scalar::f64(2.0)]),
                (Timestamp(3), vec![Scalar::bool(false), Scalar::f64(3.0)]),
            ]
        );
    }

    #[test]
    fn does_not_double_emit_a_gated_final_row() {
        // final row IS gated -> emitted once by the gate; finalize must not re-emit.
        let out = drive(&[(false, 1.0), (true, 2.0)]);
        assert_eq!(out, vec![(Timestamp(2), vec![Scalar::bool(true), Scalar::f64(2.0)])]);
    }

    #[test]
    fn cold_warmup_records_nothing() {
        let (tx, rx) = mpsc::channel();
        let mut g = GatedRecorder::new(&[ScalarKind::Bool, ScalarKind::F64], 0, tx);
        let inputs = vec![
            AnyColumn::with_capacity(ScalarKind::Bool, 1),
            AnyColumn::with_capacity(ScalarKind::F64, 1),
        ];
        // both columns cold -> None, nothing recorded, finalize flushes nothing.
        assert_eq!(g.eval(Ctx::new(&inputs, Timestamp(1))), None);
        g.finalize();
        assert!(rx.try_recv().is_err());
    }
}