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

//! The cost-model-graph node contract (C10): the `CostNode` factor trait and the
//! `CostRunner<F>` adapter that wraps a factor into an engine `Node`.
//!
//! A cost node's only per-node difference is the **price-unit cost numerator** the
//! runner divides by the latched 1R distance. Everything else — gating on the PM
//! geometry (the co-temporality contract: the cost stream stays 1:1 with the
//! executor's record), the closed/open charge, the running `cum`, the 3-field
//! emit — is the runner's, written once. The 3-field cost record
//! (`COST_FIELD_NAMES`) is one source of truth, read by both the producer side
//! (`cost_node_builder`) and the `CostSum` aggregator; this mirrors the
//! `position_management::{FIELD_NAMES, WIDTH}` precedent and replaces what was a
//! by-convention triple lockstep.

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

/// The 3-field cost-in-R record every cost node emits, in slot order — one source
/// of truth for the producer schema and the `CostSum` aggregator.
pub const COST_WIDTH: usize = 3;
pub const COST_FIELD_NAMES: [&str; COST_WIDTH] =
    ["cost_in_r", "cum_cost_in_r", "open_cost_in_r"];

/// The PM-geometry input prefix every cost node gates on (`closed`, `open`,
/// `entry_price`, `stop_price`). A factor's own extra inputs are appended after
/// these, beginning at slot `GEOMETRY_WIDTH`.
pub const GEOMETRY_WIDTH: usize = 4;

fn geometry_input_ports() -> Vec<PortSpec> {
    vec![
        PortSpec { kind: ScalarKind::Bool, firing: Firing::Any, name: "closed".into() },
        PortSpec { kind: ScalarKind::Bool, firing: Firing::Any, name: "open".into() },
        PortSpec { kind: ScalarKind::F64, firing: Firing::Any, name: "entry_price".into() },
        PortSpec { kind: ScalarKind::F64, firing: Firing::Any, name: "stop_price".into() },
    ]
}

fn cost_output_fields() -> Vec<FieldSpec> {
    COST_FIELD_NAMES
        .iter()
        .map(|n| FieldSpec { name: (*n).into(), kind: ScalarKind::F64 })
        .collect()
}

/// When a cost factor charges its cost. `AtClose` — once per closed trade (commission,
/// flat cost, slippage): the per-trade default, charged on the close cycle. `PerHeldCycle`
/// — every cycle the position is held (carry, funding): the cost accrues over the hold.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ChargeMode {
    AtClose,
    PerHeldCycle,
}

/// A cost factor: the per-round-trip cost in *price units* (the numerator the
/// runner divides by the latched 1R distance). The only thing a cost node differs
/// in; the co-temporality skeleton is [`CostRunner`]'s, shared.
///
/// # Authoring a cost node
///
/// ```
/// use aura_core::Ctx;
/// use aura_std::{CostNode, CostRunner};
///
/// pub struct HalfSpreadCost {
///     half_spread: f64,
/// }
///
/// impl HalfSpreadCost {
///     pub fn new(half_spread: f64) -> CostRunner<HalfSpreadCost> {
///         assert!(half_spread >= 0.0, "HalfSpreadCost half_spread must be >= 0");
///         CostRunner::new(HalfSpreadCost { half_spread })
///     }
/// }
///
/// impl CostNode for HalfSpreadCost {
///     fn label(&self) -> String {
///         format!("HalfSpreadCost({})", self.half_spread)
///     }
///     fn cost_numerator(&mut self, _ctx: &Ctx<'_>) -> f64 {
///         self.half_spread // price units; the runner divides by the latched 1R distance
///     }
/// }
///
/// let _node = HalfSpreadCost::new(0.5); // a ready-to-wire cost node
/// ```
pub trait CostNode: 'static {
    /// One-line render label carrying the identifying param (C23).
    fn label(&self) -> String;
    /// Extra input ports beyond the 4 geometry inputs, appended at slot
    /// `GEOMETRY_WIDTH`. Default: none.
    fn extra_inputs(&self) -> Vec<PortSpec> {
        Vec::new()
    }
    /// When this factor charges. Default `AtClose` (the per-trade cost shape); a
    /// per-held-cycle (accrual) factor overrides this to `PerHeldCycle`.
    fn charge_mode(&self) -> ChargeMode {
        ChargeMode::AtClose
    }
    /// The round-trip cost in price units this cycle, BEFORE R-normalization and
    /// BEFORE the closed/open gate. Reads its extra inputs from `ctx` at
    /// `GEOMETRY_WIDTH + i`; an empty window during warm-up means 0 (the runner
    /// still emits the row — co-temporality).
    fn cost_numerator(&mut self, ctx: &Ctx<'_>) -> f64;
}

/// The shared co-temporality skeleton wrapping any [`CostNode`] factor into a
/// `Node`. Holds the only running state a cost node needs — `cum` and the output
/// buffer — so a factor impl stays pure.
pub struct CostRunner<F: CostNode> {
    factor: F,
    cum: f64,
    acc: f64,
    out: [Cell; COST_WIDTH],
}

impl<F: CostNode> CostRunner<F> {
    pub fn new(factor: F) -> Self {
        Self { factor, cum: 0.0, acc: 0.0, out: [Cell::from_f64(0.0); COST_WIDTH] }
    }
}

impl<F: CostNode> Node for CostRunner<F> {
    fn lookbacks(&self) -> Vec<usize> {
        vec![1; GEOMETRY_WIDTH + self.factor.extra_inputs().len()]
    }

    fn eval(&mut self, ctx: Ctx<'_>) -> Option<&[Cell]> {
        // Gate ONLY on the PM geometry (co-temporality): the cost stream stays 1:1
        // with the executor's record. A factor's not-yet-warm input contributes 0
        // (handled in `cost_numerator`), it does not withhold the row.
        let closed_w = ctx.bool_in(0);
        let open_w = ctx.bool_in(1);
        let entry_w = ctx.f64_in(2);
        let stop_w = ctx.f64_in(3);
        if closed_w.is_empty() || open_w.is_empty() || entry_w.is_empty() || stop_w.is_empty() {
            return None;
        }
        let closed = closed_w[0];
        let open = open_w[0];
        let latched = (entry_w[0] - stop_w[0]).abs();
        let numerator = self.factor.cost_numerator(&ctx);
        // Zero latched distance = no valid 1R denominator -> no cost. The
        // `numerator / latched` token form is preserved verbatim from the
        // pre-migration nodes for byte-identity (IEEE-754).
        let per = if latched > 0.0 { numerator / latched } else { 0.0 };
        let (cost_in_r, open_cost_in_r) = match self.factor.charge_mode() {
            // At-close (per-trade): charged once on the close cycle. VERBATIM the
            // pre-cycle-5 tokens — IEEE-754 byte-identity with the existing goldens.
            ChargeMode::AtClose => {
                let cost_in_r = if closed { per } else { 0.0 };
                let open_cost_in_r = if open { per } else { 0.0 };
                (cost_in_r, open_cost_in_r)
            }
            // Per-held-cycle (accrual): the carry accrues every cycle the position is
            // held; the accrued total realizes into `cum` at close, and marks the open
            // position (grows each held cycle) so the net_r_equity curve bleeds over
            // the hold rather than stepping at close.
            ChargeMode::PerHeldCycle => {
                if open || closed {
                    self.acc += per;
                }
                let cost_in_r = if closed { self.acc } else { 0.0 };
                let open_cost_in_r = if open { self.acc } else { 0.0 };
                if closed {
                    self.acc = 0.0;
                }
                (cost_in_r, open_cost_in_r)
            }
        };
        self.cum += cost_in_r;
        self.out = [
            Cell::from_f64(cost_in_r),
            Cell::from_f64(self.cum),
            Cell::from_f64(open_cost_in_r),
        ];
        Some(&self.out)
    }

    fn label(&self) -> String {
        self.factor.label()
    }
}

/// Assemble a cost-node `PrimitiveBuilder`: the 4 geometry inputs ++ the factor's
/// extra inputs, the standard 3-field cost output, the given params, and a build
/// closure. The single home for the cost-node schema shape.
pub fn cost_node_builder(
    name: &'static str,
    extra_inputs: Vec<PortSpec>,
    params: Vec<ParamSpec>,
    build: impl Fn(&[Cell]) -> Box<dyn Node> + 'static,
) -> PrimitiveBuilder {
    let mut inputs = geometry_input_ports();
    inputs.extend(extra_inputs);
    PrimitiveBuilder::new(name, NodeSchema { inputs, output: cost_output_fields(), params }, build)
}

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

    /// A test-only factor: a constant numerator, no extra inputs.
    struct StubCost(f64);
    impl CostNode for StubCost {
        fn label(&self) -> String {
            format!("StubCost({})", self.0)
        }
        fn cost_numerator(&mut self, _ctx: &Ctx<'_>) -> f64 {
            self.0
        }
    }

    /// A test-only factor with one extra f64 input, read at `GEOMETRY_WIDTH`.
    struct StubExtra(f64);
    impl CostNode for StubExtra {
        fn label(&self) -> String {
            "StubExtra".into()
        }
        fn extra_inputs(&self) -> Vec<PortSpec> {
            vec![PortSpec { kind: ScalarKind::F64, firing: Firing::Any, name: "x".into() }]
        }
        fn cost_numerator(&mut self, ctx: &Ctx<'_>) -> f64 {
            let w = ctx.f64_in(GEOMETRY_WIDTH);
            let x = if w.is_empty() { 0.0 } else { w[0] };
            self.0 * x
        }
    }

    /// A test-only factor charging per held cycle (accrual), constant numerator.
    struct StubPerHeld(f64);
    impl CostNode for StubPerHeld {
        fn label(&self) -> String {
            format!("StubPerHeld({})", self.0)
        }
        fn charge_mode(&self) -> ChargeMode {
            ChargeMode::PerHeldCycle
        }
        fn cost_numerator(&mut self, _ctx: &Ctx<'_>) -> f64 {
            self.0
        }
    }

    fn geom_cols() -> Vec<AnyColumn> {
        vec![
            AnyColumn::with_capacity(ScalarKind::Bool, 1), // closed
            AnyColumn::with_capacity(ScalarKind::Bool, 1), // open
            AnyColumn::with_capacity(ScalarKind::F64, 1),  // entry
            AnyColumn::with_capacity(ScalarKind::F64, 1),  // stop
        ]
    }

    #[test]
    fn withholds_until_geometry_present() {
        let mut r = CostRunner::new(StubCost(2.0));
        let inputs = geom_cols(); // empty
        assert_eq!(r.eval(Ctx::new(&inputs, Timestamp(0))), None);
    }

    #[test]
    fn charges_numerator_over_latched_on_close() {
        let mut r = CostRunner::new(StubCost(2.0));
        let mut inputs = geom_cols();
        inputs[0].push(Scalar::bool(true)).unwrap();
        inputs[1].push(Scalar::bool(false)).unwrap();
        inputs[2].push(Scalar::f64(100.0)).unwrap();
        inputs[3].push(Scalar::f64(96.0)).unwrap(); // latched 4 -> 2/4 = 0.5
        assert_eq!(
            r.eval(Ctx::new(&inputs, Timestamp(0))),
            Some([Cell::from_f64(0.5), Cell::from_f64(0.5), Cell::from_f64(0.0)].as_slice())
        );
    }

    #[test]
    fn open_emits_would_be_cost_not_in_cum() {
        let mut r = CostRunner::new(StubCost(2.0));
        let mut inputs = geom_cols();
        inputs[0].push(Scalar::bool(false)).unwrap();
        inputs[1].push(Scalar::bool(true)).unwrap(); // open
        inputs[2].push(Scalar::f64(100.0)).unwrap();
        inputs[3].push(Scalar::f64(96.0)).unwrap();
        assert_eq!(
            r.eval(Ctx::new(&inputs, Timestamp(0))),
            Some([Cell::from_f64(0.0), Cell::from_f64(0.0), Cell::from_f64(0.5)].as_slice())
        );
    }

    #[test]
    fn zero_latched_no_cost() {
        let mut r = CostRunner::new(StubCost(2.0));
        let mut inputs = geom_cols();
        inputs[0].push(Scalar::bool(true)).unwrap();
        inputs[1].push(Scalar::bool(false)).unwrap();
        inputs[2].push(Scalar::f64(100.0)).unwrap();
        inputs[3].push(Scalar::f64(100.0)).unwrap(); // latched 0
        assert_eq!(
            r.eval(Ctx::new(&inputs, Timestamp(0))),
            Some([Cell::from_f64(0.0), Cell::from_f64(0.0), Cell::from_f64(0.0)].as_slice())
        );
    }

    #[test]
    fn cum_accumulates() {
        let mut r = CostRunner::new(StubCost(2.0));
        let mut a = geom_cols();
        a[0].push(Scalar::bool(true)).unwrap();
        a[1].push(Scalar::bool(false)).unwrap();
        a[2].push(Scalar::f64(100.0)).unwrap();
        a[3].push(Scalar::f64(96.0)).unwrap(); // 0.5
        let _ = r.eval(Ctx::new(&a, Timestamp(0)));
        let mut b = geom_cols();
        b[0].push(Scalar::bool(true)).unwrap();
        b[1].push(Scalar::bool(false)).unwrap();
        b[2].push(Scalar::f64(100.0)).unwrap();
        b[3].push(Scalar::f64(98.0)).unwrap(); // latched 2 -> 1.0; cum 1.5
        assert_eq!(
            r.eval(Ctx::new(&b, Timestamp(1))),
            Some([Cell::from_f64(1.0), Cell::from_f64(1.5), Cell::from_f64(0.0)].as_slice())
        );
    }

    #[test]
    fn per_held_cycle_accrues_open_cost_and_dumps_at_close() {
        // per = 2/4 = 0.5 each cycle. A hold of open, open, close.
        let mut r = CostRunner::new(StubPerHeld(2.0));
        // Cycle 1: held open -> accrue 0.5 into open_cost; nothing into cum yet.
        let mut a = geom_cols();
        a[0].push(Scalar::bool(false)).unwrap(); // not closed
        a[1].push(Scalar::bool(true)).unwrap(); // open
        a[2].push(Scalar::f64(100.0)).unwrap();
        a[3].push(Scalar::f64(96.0)).unwrap(); // latched 4 -> per 0.5
        assert_eq!(
            r.eval(Ctx::new(&a, Timestamp(0))),
            Some([Cell::from_f64(0.0), Cell::from_f64(0.0), Cell::from_f64(0.5)].as_slice())
        );
        // Cycle 2: still held open -> open_cost GROWS to 1.0 (the accrual bleed), cum still 0.
        let mut b = geom_cols();
        b[0].push(Scalar::bool(false)).unwrap();
        b[1].push(Scalar::bool(true)).unwrap();
        b[2].push(Scalar::f64(100.0)).unwrap();
        b[3].push(Scalar::f64(96.0)).unwrap();
        assert_eq!(
            r.eval(Ctx::new(&b, Timestamp(1))),
            Some([Cell::from_f64(0.0), Cell::from_f64(0.0), Cell::from_f64(1.0)].as_slice())
        );
        // Cycle 3: close -> accrue the closing cycle, DUMP the total 1.5 into cost_in_r,
        // cum steps to 1.5, open_cost back to 0.
        let mut c = geom_cols();
        c[0].push(Scalar::bool(true)).unwrap(); // closed
        c[1].push(Scalar::bool(false)).unwrap(); // not open
        c[2].push(Scalar::f64(100.0)).unwrap();
        c[3].push(Scalar::f64(96.0)).unwrap();
        assert_eq!(
            r.eval(Ctx::new(&c, Timestamp(2))),
            Some([Cell::from_f64(1.5), Cell::from_f64(1.5), Cell::from_f64(0.0)].as_slice())
        );
    }

    #[test]
    fn per_held_cycle_resets_accrual_between_trades() {
        // Two trades (open, close each). acc must reset at the first close so the
        // second trade's dumped total is its OWN accrual, not cumulative.
        let mut r = CostRunner::new(StubPerHeld(2.0)); // per 0.5
        let mut o1 = geom_cols();
        o1[0].push(Scalar::bool(false)).unwrap();
        o1[1].push(Scalar::bool(true)).unwrap();
        o1[2].push(Scalar::f64(100.0)).unwrap();
        o1[3].push(Scalar::f64(96.0)).unwrap();
        let _ = r.eval(Ctx::new(&o1, Timestamp(0))); // open_cost 0.5
        let mut c1 = geom_cols();
        c1[0].push(Scalar::bool(true)).unwrap();
        c1[1].push(Scalar::bool(false)).unwrap();
        c1[2].push(Scalar::f64(100.0)).unwrap();
        c1[3].push(Scalar::f64(96.0)).unwrap();
        assert_eq!(
            r.eval(Ctx::new(&c1, Timestamp(1))),
            Some([Cell::from_f64(1.0), Cell::from_f64(1.0), Cell::from_f64(0.0)].as_slice())
        ); // trade 1 total 1.0, cum 1.0, acc reset
        let mut o2 = geom_cols();
        o2[0].push(Scalar::bool(false)).unwrap();
        o2[1].push(Scalar::bool(true)).unwrap();
        o2[2].push(Scalar::f64(100.0)).unwrap();
        o2[3].push(Scalar::f64(96.0)).unwrap();
        let _ = r.eval(Ctx::new(&o2, Timestamp(2))); // fresh open_cost 0.5
        let mut c2 = geom_cols();
        c2[0].push(Scalar::bool(true)).unwrap();
        c2[1].push(Scalar::bool(false)).unwrap();
        c2[2].push(Scalar::f64(100.0)).unwrap();
        c2[3].push(Scalar::f64(96.0)).unwrap();
        assert_eq!(
            r.eval(Ctx::new(&c2, Timestamp(3))),
            Some([Cell::from_f64(1.0), Cell::from_f64(2.0), Cell::from_f64(0.0)].as_slice())
        ); // trade 2's OWN total 1.0 (acc reset proved), cum running 2.0
    }

    #[test]
    fn extra_input_cold_contributes_zero_but_row_emits() {
        // Co-temporality: a not-yet-warm factor input -> 0 cost, but a row IS emitted.
        let mut r = CostRunner::new(StubExtra(0.5));
        let mut inputs = geom_cols();
        inputs.push(AnyColumn::with_capacity(ScalarKind::F64, 1)); // x, empty
        inputs[0].push(Scalar::bool(true)).unwrap();
        inputs[1].push(Scalar::bool(false)).unwrap();
        inputs[2].push(Scalar::f64(100.0)).unwrap();
        inputs[3].push(Scalar::f64(96.0)).unwrap();
        assert_eq!(
            r.eval(Ctx::new(&inputs, Timestamp(0))),
            Some([Cell::from_f64(0.0), Cell::from_f64(0.0), Cell::from_f64(0.0)].as_slice())
        );
    }

    #[test]
    fn extra_input_warm_scales_numerator() {
        let mut r = CostRunner::new(StubExtra(0.5));
        let mut inputs = geom_cols();
        inputs.push(AnyColumn::with_capacity(ScalarKind::F64, 1));
        inputs[0].push(Scalar::bool(true)).unwrap();
        inputs[1].push(Scalar::bool(false)).unwrap();
        inputs[2].push(Scalar::f64(100.0)).unwrap();
        inputs[3].push(Scalar::f64(96.0)).unwrap(); // latched 4
        inputs[4].push(Scalar::f64(3.0)).unwrap(); // x=3 -> 0.5*3=1.5 -> 1.5/4 = 0.375
        assert_eq!(
            r.eval(Ctx::new(&inputs, Timestamp(0))),
            Some([Cell::from_f64(0.375), Cell::from_f64(0.375), Cell::from_f64(0.0)].as_slice())
        );
    }

    #[test]
    fn lookbacks_count_geometry_plus_extra() {
        assert_eq!(CostRunner::new(StubCost(1.0)).lookbacks(), vec![1; GEOMETRY_WIDTH]);
        assert_eq!(CostRunner::new(StubExtra(1.0)).lookbacks(), vec![1; GEOMETRY_WIDTH + 1]);
    }

    #[test]
    fn runner_label_delegates_to_factor() {
        assert_eq!(CostRunner::new(StubCost(2.0)).label(), "StubCost(2)");
    }

    #[test]
    fn geometry_width_matches_port_count() {
        assert_eq!(GEOMETRY_WIDTH, geometry_input_ports().len());
    }

    #[test]
    fn cost_output_fields_are_the_triple() {
        let names: Vec<String> = cost_output_fields().into_iter().map(|f| f.name).collect();
        assert_eq!(names, COST_FIELD_NAMES.to_vec());
    }

    #[test]
    fn cost_node_builder_assembles_geometry_prefix_then_extras() {
        // The builder's input-assembly contract, at the schema level: the 4 geometry
        // ports first (in slot order), the factor's extra inputs appended beginning
        // at GEOMETRY_WIDTH, the shared 3-field cost output, and the given params
        // passed through verbatim.
        let extra = vec![PortSpec { kind: ScalarKind::F64, firing: Firing::Any, name: "vol".into() }];
        let params = vec![ParamSpec { name: "k".into(), kind: ScalarKind::F64 }];
        let builder = cost_node_builder(
            "Probe",
            extra.clone(),
            params.clone(),
            |_p| -> Box<dyn Node> { Box::new(CostRunner::new(StubCost(0.0))) },
        );
        let schema = builder.schema();
        // The first GEOMETRY_WIDTH inputs are the geometry prefix verbatim...
        assert_eq!(schema.inputs[..GEOMETRY_WIDTH], geometry_input_ports()[..]);
        // ...and the factor's extras begin exactly at slot GEOMETRY_WIDTH.
        assert_eq!(schema.inputs[GEOMETRY_WIDTH..], extra[..]);
        assert_eq!(schema.inputs.len(), GEOMETRY_WIDTH + extra.len());
        // The output is the shared 3-field cost triple; params pass through.
        assert_eq!(schema.output, cost_output_fields());
        assert_eq!(schema.params, params);
    }

    #[test]
    fn producer_and_aggregator_share_the_triple() {
        // The structural lockstep: producer output and aggregator output/inputs all
        // read COST_FIELD_NAMES (one source), replacing the by-convention triple.
        let prod: Vec<String> =
            ConstantCost::builder().schema().output.iter().map(|f| f.name.clone()).collect();
        assert_eq!(prod, COST_FIELD_NAMES.to_vec());
        let agg_out: Vec<String> =
            CostSum::builder(1).schema().output.iter().map(|f| f.name.clone()).collect();
        assert_eq!(agg_out, COST_FIELD_NAMES.to_vec());
        let agg_in: Vec<String> =
            CostSum::builder(1).schema().inputs.iter().map(|p| p.name.clone()).collect();
        let expected: Vec<String> =
            COST_FIELD_NAMES.iter().map(|f| format!("cost[0].{f}")).collect();
        assert_eq!(agg_in, expected);
    }
}