cubelinux-2/cubecode/src/trace.rs

//! Execution-trace artifacts: store a captured machine-code trace as a
//! first-class AI-artifact record (`Kind::Layer`) linked back to its source,
//! and replay it to reproduce the run — so an agent can persist & reproduce
//! interesting executions as ordinary CZYX records.
//!
//! A trace is the exact `Vec<Op>` a program actually executed ([`Vm::run_captured`]).
//! We persist it as a record body (via [`crate::opcode::encode`]) with
//! `doc_type = "layer"` (the AI-artifact kind) and a `linked_records` edge to
//! the source cell. Loading it back (via [`CodeCell::from_record`]) decodes the
//! ops and [`replay`] reproduces the observable behaviour with no access to the
//! source program.

use crate::{opcode, replay, Kind, Op, RunResult, Vm};
use cubecoords::{CubeHeader, Czyx};
use cubestore::{CubeBackend, CubeStore};

/// Store a captured execution trace as an AI-artifact record at `coord`,
/// linked to its `source` cell. `name` becomes the record title; the body is
/// the encoded trace; `doc_type` is `Kind::Layer`.
pub fn store_trace<B: CubeBackend>(
    store: &mut CubeStore<B>,
    coord: Czyx,
    source: Czyx,
    trace: &[Op],
    name: &str,
) {
    let mut h = CubeHeader::new();
    h.title = Some(name.to_string());
    h.doc_type = Some(Kind::Layer.as_str().to_string());
    h.linked_records = vec![source];
    h.refresh_flags();
    store.put_record(coord, &h, &opcode::encode(trace));
}

/// Load a stored trace record at `coord` and replay it, reproducing the run's
/// observable behaviour (data-stack result + `SYS_TRACE` output) from the
/// stored record — no access to the original program. Returns `None` if the
/// record is missing or its body is not a valid trace.
pub fn replay_trace<B: CubeBackend>(
    store: &CubeStore<B>,
    coord: Czyx,
) -> Option<(RunResult, Vec<u8>)> {
    let (_h, body) = store.get_record(&coord)?;
    let ops = opcode::decode(&body).ok()?;
    Some(replay(&ops))
}

/// Capture a run and persist it as a "golden": the executed trace stored as a
/// `Kind::Layer` record (linked to `source`) and the observed output stored as a
/// companion `golden` record (linked to the trace). Returns the observed output.
pub fn capture_golden<B: CubeBackend + Clone>(
    store: &mut CubeStore<B>,
    source: Czyx,
    trace_coord: Czyx,
    golden_coord: Czyx,
    name: &str,
) -> Vec<u8> {
    let mut vm = Vm::new((*store).clone());
    let (_r, trace) = vm.run_captured(source);
    let out = vm.output().to_vec();
    store_trace(store, trace_coord, source, &trace, name);
    let mut h = CubeHeader::new();
    h.title = Some(format!("golden:{name}"));
    h.doc_type = Some("golden".into());
    h.linked_records = vec![trace_coord];
    h.refresh_flags();
    store.put_record(golden_coord, &h, &out);
    out
}

/// Golden-trace regression check: re-run `source` fresh AND replay the stored
/// trace, and assert BOTH reproduce the recorded golden output.
///   * replay vs golden  -> the stored trace is a faithful reproduction;
///   * fresh run vs golden -> the program has not drifted (or it would change
///     the observed output).
/// Returns `Ok(())` if the behavior is unchanged, else `Err(reason)`.
pub fn verify_golden<B: CubeBackend + Clone>(
    store: &CubeStore<B>,
    source: Czyx,
    trace_coord: Czyx,
    golden_coord: Czyx,
) -> Result<(), String> {
    let golden = store
        .get_record(&golden_coord)
        .map(|(_, b)| b)
        .ok_or("golden record missing")?;
    let (_, replay_out) = replay_trace(store, trace_coord).ok_or("trace record missing/invalid")?;
    if replay_out != golden {
        return Err(format!(
            "replay output differs from golden ({} vs {} bytes)",
            replay_out.len(),
            golden.len()
        ));
    }
    let mut vm = Vm::new((*store).clone());
    let _ = vm.run(source);
    let fresh = vm.output().to_vec();
    if fresh != golden {
        return Err(format!(
            "fresh run differs from golden (program changed: {} vs {} bytes)",
            fresh.len(),
            golden.len()
        ));
    }
    Ok(())
}

/// A hop in the experiment-lineage graph: a record and the records it links to.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LineageHop {
    /// The record's coordinate.
    pub coord: Czyx,
    /// Its `doc_type` (kind: fn/kernel/layer/checkpoint/variant/golden/...).
    pub doc_type: String,
    /// Its human title, if any.
    pub title: String,
    /// The CZYX records this record links to (its edges).
    pub links: Vec<Czyx>,
}

/// Walk the linked-record graph starting at `start`, following
/// `linked_records` edges (breadth-first) up to `depth` levels. Returns the
/// ordered list of reachable records — the experiment lineage chain
/// (source -> data -> weights -> trace -> golden, etc.). Each record is the
/// kind the CUBE associates; provenance is a graph walk, not a lookup.
pub fn lineage<B: CubeBackend>(
    store: &CubeStore<B>,
    start: Czyx,
    depth: usize,
) -> Vec<LineageHop> {
    let mut out = Vec::new();
    let mut frontier = vec![start];
    let mut seen: std::collections::HashSet<Czyx> = std::collections::HashSet::new();
    for _ in 0..=depth {
        if frontier.is_empty() {
            break;
        }
        let mut next = Vec::new();
        for c in frontier {
            if !seen.insert(c) {
                continue;
            }
            if let Some((h, _)) = store.get_record(&c) {
                let links = h.linked_records.clone();
                out.push(LineageHop {
                    coord: c,
                    doc_type: h.doc_type.clone().unwrap_or_default(),
                    title: h.title.clone().unwrap_or_default(),
                    links: links.clone(),
                });
                next.extend(links);
            }
        }
        frontier = next;
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{SYS_TRACE, Vm};
    use cubestore::HashBackend;

    #[test]
    fn trace_stored_as_layer_and_replayed_without_source() {
        let mut store = CubeStore::new(HashBackend::new());
        let entry = Czyx::new(1, 1, 1, 7);
        let ops = [
            Op::Const(0), Op::Const(1), Op::Syscall(SYS_TRACE),
            Op::Const(1), Op::Const(1), Op::Syscall(SYS_TRACE),
            Op::Const(2), Op::Const(2), Op::Syscall(SYS_TRACE),
            Op::Halt,
        ];
        let mut h = CubeHeader::new();
        h.title = Some("solitaire".into());
        h.doc_type = Some("fn".into());
        h.refresh_flags();
        store.put_record(entry, &h, &opcode::encode(&ops));

        // 1) Capture the run as a machine-code trace.
        let mut vm = Vm::new(store.clone());
        let (_r, trace) = vm.run_captured(entry);
        let produced = vm.output().to_vec();
        assert!(!trace.is_empty());

        // 2) Store it as a Kind::Layer AI-artifact record linked to the source.
        let trace_coord = Czyx::new(9, 0, 0, 1);
        store_trace(&mut store, trace_coord, entry, &trace, "solitaire-run-1");

        // 3) The record is a "layer" artifact with an edge back to the source.
        let (th, _tb) = store.get_record(&trace_coord).unwrap();
        assert_eq!(th.doc_type.as_deref(), Some("layer"));
        assert_eq!(th.linked_records, vec![entry]);

        // 4) Reproduce the run purely from the stored trace record.
        let (result, replayed) = replay_trace(&store, trace_coord).unwrap();
        assert_eq!(result, RunResult::Halted { top: None });
        assert_eq!(replayed, produced, "stored trace must reproduce the run");
        assert!(replayed.contains(&b'\n'));
    }

    #[test]
    fn golden_regression_detects_drift() {
        let mut store = CubeStore::new(HashBackend::new());
        let entry = Czyx::new(1, 1, 1, 9);
        let mk = |val: u8| {
            vec![Op::Const(val), Op::Syscall(SYS_TRACE), Op::Halt]
        };
        let mut h = CubeHeader::new();
        h.doc_type = Some("fn".into());
        h.refresh_flags();
        store.put_record(entry, &h, &opcode::encode(&mk(7)));

        let trace_c = Czyx::new(9, 0, 0, 1);
        let golden_c = Czyx::new(9, 0, 0, 2);
        let out = capture_golden(&mut store, entry, trace_c, golden_c, "regress");
        assert_eq!(out, b"\x07\n");
        assert!(verify_golden(&store, entry, trace_c, golden_c).is_ok());

        // Drift the program: change the emitted value -> fresh run differs.
        store.put_record(entry, &h, &opcode::encode(&mk(9)));
        let r = verify_golden(&store, entry, trace_c, golden_c);
        assert!(r.is_err(), "program drift must be detected: {r:?}");
        // The stored trace is immutable: it still replays the ORIGINAL golden.
        let (_, rep) = replay_trace(&store, trace_c).unwrap();
        assert_eq!(rep, b"\x07\n", "stored trace still reproduces the original run");
    }

    #[test]
    fn lineage_walks_source_data_weights_trace() {
        let mut store = CubeStore::new(HashBackend::new());
        // Build experiment lineage: source -> data -> weights -> trace -> golden.
        let mut mk = |coord: Czyx, doc: &str, title: &str, links: Vec<Czyx>| {
            let mut h = CubeHeader::new();
            h.title = Some(title.into());
            h.doc_type = Some(doc.into());
            h.linked_records = links;
            h.refresh_flags();
            store.put_record(coord, &h, &[]);
        };
        let source = Czyx::new(1, 1, 1, 1);
        let data = Czyx::new(2, 0, 0, 1);
        let weights = Czyx::new(3, 0, 0, 1);
        let trace = Czyx::new(4, 0, 0, 1);
        let golden = Czyx::new(5, 0, 0, 1);
        mk(source, "fn", "train", vec![data, weights]);
        mk(data, "data", "batch", vec![]);
        mk(weights, "layer", "model-v1", vec![]);
        mk(trace, "layer", "trace:train", vec![source]);
        mk(golden, "golden", "golden:train", vec![trace]);

        // Walk from the trace back to its lineage.
        let hops = lineage(&store, trace, 4);
        assert!(hops.iter().any(|h| h.coord == trace && h.doc_type == "layer"));
        assert!(hops.iter().any(|h| h.coord == source && h.doc_type == "fn"));
        // Depth-limited: from a leaf with no links, only itself.
        let leaf = lineage(&store, data, 3);
        assert_eq!(leaf.len(), 1);
        assert_eq!(leaf[0].coord, data);
        assert_eq!(leaf[0].links, vec![]);
    }
}