Aura/crates/aura-strategy/src/cost_sum.rs

//! `CostSum` — the output node of a C10 cost-model graph: it sums `n_costs`
//! cost-in-R records per-field into one aggregate record, so any number of cost
//! nodes collapses to the single 3-field cost stream the net-R seam already
//! consumes (`summarize_r` + the `net_r_equity` tap stay unchanged). Each cost
//! node contributes the 3-field `{cost_in_r, cum_cost_in_r, open_cost_in_r}`
//! record; the aggregate is the per-field sum. `n_costs = 1` is the identity, so
//! the cost path is uniform whether one or several cost nodes are wired.

use aura_core::{
    Cell, Ctx, FieldSpec, Firing, Node, NodeSchema, PortSpec, PrimitiveBuilder, ScalarKind,
};

/// The cost-record field triple and its width come from the shared cost contract
/// (`crate::cost::{COST_FIELD_NAMES, COST_WIDTH}`) — one source of truth, read by
/// both the producer side (`cost_node_builder`) and this aggregator. The input-name
/// loop, the output schema, the lookback vector, and the eval accumulator all read
/// it, so the producer↔aggregator field match is structural, not by-convention.
use crate::cost::{cost_port, COST_FIELD_NAMES, COST_WIDTH};

/// Per-field sum of `n_costs` cost-in-R records. Inputs are
/// `cost[k].{cost_in_r,cum_cost_in_r,open_cost_in_r}` for `k in 0..n_costs`, in
/// slot order (3 per cost node); the 3-field output mirrors a single cost record.
/// Emits `None` until every input leg is present (mode-A as-of join, like LinComb).
pub struct CostSum {
    n_costs: usize,
    out: [Cell; COST_WIDTH],
}

impl CostSum {
    pub fn new(n_costs: usize) -> Self {
        assert!(n_costs >= 1, "CostSum needs at least one cost input");
        Self { n_costs, out: [Cell::from_f64(0.0); COST_WIDTH] }
    }

    /// The param-generic recipe. `n_costs` is topology (fixed per blueprint, C19),
    /// captured by the build closure (no per-build params). The input names are a
    /// lockstep contract with the connect side (`cost[k].<field>`).
    pub fn builder(n_costs: usize) -> PrimitiveBuilder {
        let mut inputs = Vec::with_capacity(n_costs * COST_WIDTH);
        for k in 0..n_costs {
            for field in COST_FIELD_NAMES {
                inputs.push(PortSpec {
                    kind: ScalarKind::F64,
                    firing: Firing::Any,
                    // Owned String (PortSpec.name is String); the NAME comes from
                    // the shared cost_port helper so the aggregator's input names
                    // and cost_graph's agg-side names share ONE shape definition.
                    name: cost_port(k, field).to_string(),
                });
            }
        }
        PrimitiveBuilder::new(
            "CostSum",
            NodeSchema {
                inputs,
                output: COST_FIELD_NAMES
                    .iter()
                    .map(|n| FieldSpec { name: (*n).into(), kind: ScalarKind::F64 })
                    .collect(),
                params: vec![],
                doc: "sums cost-model contributions into one cost-in-R stream",
            },
            move |_| Box::new(CostSum::new(n_costs)),
        )
    }
}

impl Node for CostSum {
    fn lookbacks(&self) -> Vec<usize> {
        vec![1; self.n_costs * COST_WIDTH]
    }

    fn eval(&mut self, ctx: Ctx<'_>) -> Option<&[Cell]> {
        let mut acc = [0.0_f64; COST_WIDTH]; // per-field accumulator, COST_FIELD_NAMES order
        for k in 0..self.n_costs {
            for (f, slot) in acc.iter_mut().enumerate() {
                let w = ctx.f64_in(k * COST_WIDTH + f);
                if w.is_empty() {
                    return None;
                }
                *slot += w[0];
            }
        }
        self.out = [Cell::from_f64(acc[0]), Cell::from_f64(acc[1]), Cell::from_f64(acc[2])];
        Some(&self.out)
    }

    fn label(&self) -> String {
        format!("CostSum({})", self.n_costs)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use aura_core::{AnyColumn, Scalar, Timestamp};

    fn f64_cols(n: usize) -> Vec<AnyColumn> {
        (0..n).map(|_| AnyColumn::with_capacity(ScalarKind::F64, 1)).collect()
    }

    #[test]
    fn two_records_sum_per_field() {
        let mut s = CostSum::new(2);
        let mut inputs = f64_cols(6);
        // cost[0] = [0.5, 0.5, 0.0]; cost[1] = [0.375, 1.0, 0.2]
        for (i, v) in [0.5, 0.5, 0.0, 0.375, 1.0, 0.2].into_iter().enumerate() {
            inputs[i].push(Scalar::f64(v)).unwrap();
        }
        // per-field sum: [0.875, 1.5, 0.2]
        assert_eq!(
            s.eval(Ctx::new(&inputs, Timestamp(0))),
            Some([Cell::from_f64(0.875), Cell::from_f64(1.5), Cell::from_f64(0.2)].as_slice())
        );
    }

    #[test]
    fn n_one_is_identity() {
        let mut s = CostSum::new(1);
        let mut inputs = f64_cols(3);
        for (i, v) in [0.5, 1.25, 0.3].into_iter().enumerate() {
            inputs[i].push(Scalar::f64(v)).unwrap();
        }
        assert_eq!(
            s.eval(Ctx::new(&inputs, Timestamp(0))),
            Some([Cell::from_f64(0.5), Cell::from_f64(1.25), Cell::from_f64(0.3)].as_slice())
        );
    }

    #[test]
    fn withholds_until_every_leg_present() {
        let mut s = CostSum::new(2);
        let mut inputs = f64_cols(6);
        // only the first cost node's three fields present -> withhold
        for col in inputs.iter_mut().take(3) {
            col.push(Scalar::f64(1.0)).unwrap();
        }
        assert_eq!(s.eval(Ctx::new(&inputs, Timestamp(0))), None);
    }

    #[test]
    fn input_slots_are_named_cost_index_field() {
        let s = CostSum::builder(2);
        let names: Vec<String> = s.schema().inputs.iter().map(|p| p.name.clone()).collect();
        assert_eq!(
            names,
            [
                "cost[0].cost_in_r", "cost[0].cum_cost_in_r", "cost[0].open_cost_in_r",
                "cost[1].cost_in_r", "cost[1].cum_cost_in_r", "cost[1].open_cost_in_r",
            ]
        );
    }

    #[test]
    fn label_carries_the_arity() {
        assert_eq!(CostSum::new(2).label(), "CostSum(2)");
    }

    #[test]
    #[should_panic(expected = "CostSum needs at least one cost input")]
    fn new_panics_on_zero() {
        let _ = CostSum::new(0);
    }
}