Skip to main content

lex_trace/
recorder.rs

1//! Trace recorder — implements `lex_bytecode::vm::Tracer` and builds a
2//! `TraceTree` as the VM executes.
3
4use indexmap::IndexMap;
5use lex_bytecode::vm::Tracer;
6use lex_bytecode::Value;
7use serde::{Deserialize, Serialize};
8use std::sync::{Arc, Mutex};
9
10#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
11pub struct RunId(pub String);
12
13impl RunId {
14    pub fn new(seed: &str) -> Self {
15        use sha2::{Digest, Sha256};
16        let mut h = Sha256::new();
17        h.update(seed.as_bytes());
18        h.update(format!("{:?}", std::time::SystemTime::now()).as_bytes());
19        let r = h.finalize();
20        let mut hex = String::with_capacity(64);
21        for b in r { hex.push_str(&format!("{:02x}", b)); }
22        RunId(hex)
23    }
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
27#[serde(rename_all = "snake_case")]
28pub enum TraceNodeKind { Call, Effect }
29
30#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
31pub struct TraceNode {
32    pub node_id: String,
33    pub kind: TraceNodeKind,
34    /// For `Call`: the function name. For `Effect`: `kind.op` (e.g. `io.print`).
35    pub target: String,
36    pub input: serde_json::Value,
37    /// `Some` on success; `None` if the node ended in error.
38    #[serde(default, skip_serializing_if = "Option::is_none")]
39    pub output: Option<serde_json::Value>,
40    #[serde(default, skip_serializing_if = "Option::is_none")]
41    pub error: Option<String>,
42    pub started_at: u64,
43    pub ended_at: u64,
44    #[serde(default)]
45    pub children: Vec<TraceNode>,
46}
47
48#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
49pub struct TraceTree {
50    pub run_id: String,
51    pub root_target: String,
52    pub root_input: serde_json::Value,
53    pub root_output: Option<serde_json::Value>,
54    pub root_error: Option<String>,
55    pub started_at: u64,
56    pub ended_at: u64,
57    pub nodes: Vec<TraceNode>,
58}
59
60impl TraceTree {
61    /// Find a node by `NodeId`, depth-first.
62    pub fn find(&self, node_id: &str) -> Option<&TraceNode> {
63        for n in &self.nodes {
64            if let Some(found) = find_in(n, node_id) { return Some(found); }
65        }
66        None
67    }
68}
69
70fn find_in<'a>(n: &'a TraceNode, target: &str) -> Option<&'a TraceNode> {
71    if n.node_id == target { return Some(n); }
72    for c in &n.children {
73        if let Some(f) = find_in(c, target) { return Some(f); }
74    }
75    None
76}
77
78/// Tracer that builds a `TraceTree`. The tree is shared via `Arc<Mutex>`
79/// so callers can read it after the VM finishes.
80pub struct Recorder {
81    state: Arc<Mutex<RecorderState>>,
82}
83
84pub(crate) struct RecorderState {
85    /// Open frames: each entry has its inputs filled in but `output`/
86    /// `error`/`ended_at` not yet known. Children of an open frame are
87    /// staged into a sibling buffer; on `exit`, they get attached to the
88    /// node that's closing.
89    open: Vec<OpenFrame>,
90    /// Top-level finished nodes (the call we're tracing might span the
91    /// whole VM run, so this is normally a single node tree).
92    completed: Vec<TraceNode>,
93    /// Effect overrides for replay; keyed by NodeId.
94    pub(crate) overrides: IndexMap<String, serde_json::Value>,
95}
96
97struct OpenFrame {
98    node: TraceNode,
99    /// Children that have completed under this frame.
100    children: Vec<TraceNode>,
101}
102
103impl Recorder {
104    pub fn new() -> Self {
105        Self {
106            state: Arc::new(Mutex::new(RecorderState {
107                open: Vec::new(),
108                completed: Vec::new(),
109                overrides: IndexMap::new(),
110            })),
111        }
112    }
113
114    /// Returned handle stays valid after the tracer is moved into the VM.
115    pub fn handle(&self) -> Handle {
116        Handle { state: Arc::clone(&self.state) }
117    }
118
119    /// Pre-load effect overrides for replay.
120    pub fn with_overrides(self, overrides: IndexMap<String, serde_json::Value>) -> Self {
121        self.state.lock().unwrap().overrides = overrides;
122        self
123    }
124}
125
126impl Default for Recorder { fn default() -> Self { Self::new() } }
127
128#[derive(Clone)]
129pub struct Handle {
130    state: Arc<Mutex<RecorderState>>,
131}
132
133impl Handle {
134    /// Drain the recorder into a finished `TraceTree`. Call after the VM
135    /// run returns. `root_target` and `root_input` describe the top-level
136    /// call (e.g. the `lex run` entry).
137    pub fn finalize(
138        &self,
139        root_target: impl Into<String>,
140        root_input: serde_json::Value,
141        root_output: Option<serde_json::Value>,
142        root_error: Option<String>,
143        started_at: u64,
144        ended_at: u64,
145    ) -> TraceTree {
146        let st = self.state.lock().unwrap();
147        TraceTree {
148            run_id: RunId::new(&format!("{}-{}", started_at, ended_at)).0,
149            root_target: root_target.into(),
150            root_input,
151            root_output,
152            root_error,
153            started_at,
154            ended_at,
155            nodes: st.completed.clone(),
156        }
157    }
158}
159
160fn now_unix() -> u64 {
161    use std::time::{SystemTime, UNIX_EPOCH};
162    SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_secs()).unwrap_or(0)
163}
164
165fn values_to_json(args: &[Value]) -> serde_json::Value {
166    serde_json::Value::Array(args.iter().map(value_to_json).collect())
167}
168
169fn value_to_json(v: &Value) -> serde_json::Value {
170    use serde_json::Value as J;
171    match v {
172        Value::Int(n) => J::from(*n),
173        Value::Float(f) => J::from(*f),
174        Value::Bool(b) => J::Bool(*b),
175        Value::Str(s) => J::String(s.to_string()),
176        Value::Bytes(b) => J::String(b.iter().map(|b| format!("{:02x}", b)).collect()),
177        Value::Unit => J::Null,
178        Value::List(items) => J::Array(items.iter().map(value_to_json).collect()),
179        Value::Tuple(items) => J::Array(items.iter().map(value_to_json).collect()),
180        Value::Record { fields, .. } => {
181            let mut m = serde_json::Map::new();
182            for (k, v) in fields.iter() { m.insert(k.to_string(), value_to_json(v)); }
183            J::Object(m)
184        }
185        // #464 step 2: escape analysis prevents this variant from
186        // reaching a trace boundary (the tracer only records args
187        // on Call/EffectCall — both escape sinks the analysis
188        // rejects). If we ever do see one, that's an analysis bug.
189        Value::StackRecord { .. } => J::String("<stack-record-unreachable>".into()),
190        // #464 tuple codegen: same reasoning as StackRecord above —
191        // the tracer only records args at escape sinks the analysis
192        // rejects, so a frame-local tuple can't reach here.
193        Value::StackTuple { .. } => J::String("<stack-tuple-unreachable>".into()),
194        // #463 slice 2a: arena-eligibility analysis (the request-scope
195        // variant of #464's escape pass) excludes the same Call /
196        // EffectCall sinks, so an arena handle can't reach the tracer
197        // either. Same defensive marker as the stack variants above.
198        Value::ArenaRecord { .. } => J::String("<arena-record-unreachable>".into()),
199        Value::ArenaTuple { .. } => J::String("<arena-tuple-unreachable>".into()),
200        Value::Variant { name, args } => {
201            let mut m = serde_json::Map::new();
202            m.insert("$variant".into(), J::String(name.clone()));
203            m.insert("args".into(), J::Array(args.iter().map(value_to_json).collect()));
204            J::Object(m)
205        }
206        Value::Closure { body_hash, .. } => {
207            // Render the first 4 bytes (8 hex chars) of the body hash
208            // (#222). Equivalent closures across source locations now
209            // produce the same trace token, so trace replay is stable
210            // when a developer moves a closure literal.
211            let prefix: String = body_hash.iter().take(4)
212                .map(|b| format!("{b:02x}")).collect();
213            J::String(format!("<closure {prefix}>"))
214        }
215        Value::F64Array { rows, cols, data } => {
216            let mut m = serde_json::Map::new();
217            m.insert("$f64_array".into(), J::Bool(true));
218            m.insert("rows".into(), J::from(*rows));
219            m.insert("cols".into(), J::from(*cols));
220            m.insert("data".into(), J::Array(data.iter().map(|f| J::from(*f)).collect()));
221            J::Object(m)
222        }
223        Value::Map(m) => {
224            let mut o = serde_json::Map::new();
225            o.insert("$map".into(), J::Bool(true));
226            o.insert("entries".into(), J::Array(m.iter().map(|(k, v)| {
227                J::Array(vec![value_to_json(&k.as_value()), value_to_json(v)])
228            }).collect()));
229            J::Object(o)
230        }
231        Value::Set(s) => {
232            let mut o = serde_json::Map::new();
233            o.insert("$set".into(), J::Bool(true));
234            o.insert("items".into(), J::Array(
235                s.iter().map(|k| value_to_json(&k.as_value())).collect()));
236            J::Object(o)
237        }
238        Value::Deque(items) => {
239            let mut o = serde_json::Map::new();
240            o.insert("$deque".into(), J::Bool(true));
241            o.insert("items".into(), J::Array(
242                items.iter().map(value_to_json).collect()));
243            J::Object(o)
244        }
245        Value::Actor(_) => J::String("<actor>".into()),
246        Value::Ticker(_) => J::String("<ticker>".into()),
247        Value::ArrowTable(t) => {
248            // Trace records the *shape*, not the data — full Arrow tables
249            // can be GB-scale. Replay through the agent API doesn't need
250            // the rows; if it does, capture them via `arrow.row_at`.
251            let mut o = serde_json::Map::new();
252            o.insert("$arrow_table".into(), J::Bool(true));
253            o.insert("nrows".into(), J::from(t.num_rows() as i64));
254            o.insert("ncols".into(), J::from(t.num_columns() as i64));
255            J::Object(o)
256        }
257    }
258}
259
260pub(crate) fn json_to_value(v: &serde_json::Value) -> Value {
261    use serde_json::Value as J;
262    match v {
263        J::Null => Value::Unit,
264        J::Bool(b) => Value::Bool(*b),
265        J::Number(n) => {
266            if let Some(i) = n.as_i64() { Value::Int(i) }
267            else if let Some(f) = n.as_f64() { Value::Float(f) }
268            else { Value::Unit }
269        }
270        J::String(s) => Value::Str(s.as_str().into()),
271        J::Array(items) => Value::List(items.iter().map(json_to_value).collect()),
272        J::Object(map) => {
273            // Detect the $variant shape we emit on the way out.
274            if let (Some(serde_json::Value::String(name)), Some(serde_json::Value::Array(args))) =
275                (map.get("$variant"), map.get("args"))
276            {
277                return Value::Variant {
278                    name: name.clone(),
279                    args: args.iter().map(json_to_value).collect(),
280                };
281            }
282            let mut out = indexmap::IndexMap::new();
283            for (k, v) in map { out.insert(k.clone(), json_to_value(v)); }
284            Value::record_dynamic(out)
285        }
286    }
287}
288
289impl Tracer for Recorder {
290    fn enter_call(&mut self, node_id: &str, name: &str, args: &[Value]) {
291        push_call_frame(&self.state, node_id, name, args);
292    }
293    fn enter_effect(&mut self, node_id: &str, kind: &str, op: &str, args: &[Value]) {
294        push_effect_frame(&self.state, node_id, kind, op, args);
295    }
296    fn exit_ok(&mut self, value: &Value) { exit_ok_frame(&self.state, value); }
297    fn exit_err(&mut self, message: &str) { exit_err_frame(&self.state, message); }
298    fn exit_call_tail(&mut self) { exit_tail_frame(&self.state); }
299    fn override_effect(&mut self, node_id: &str) -> Option<Value> {
300        lookup_override(&self.state, node_id)
301    }
302}
303
304/// Tracer impl for the recorder's shareable handle (#199). Multiple
305/// `Vm` instances driven against the same `Recorder` — for example,
306/// the spec-checker's per-`SpecExpr::Call` Vms — can each take their
307/// own `Box<dyn Tracer>` cloned from this handle, and the events
308/// will fold into the same trace tree.
309impl Tracer for Handle {
310    fn enter_call(&mut self, node_id: &str, name: &str, args: &[Value]) {
311        push_call_frame(&self.state, node_id, name, args);
312    }
313    fn enter_effect(&mut self, node_id: &str, kind: &str, op: &str, args: &[Value]) {
314        push_effect_frame(&self.state, node_id, kind, op, args);
315    }
316    fn exit_ok(&mut self, value: &Value) { exit_ok_frame(&self.state, value); }
317    fn exit_err(&mut self, message: &str) { exit_err_frame(&self.state, message); }
318    fn exit_call_tail(&mut self) { exit_tail_frame(&self.state); }
319    fn override_effect(&mut self, node_id: &str) -> Option<Value> {
320        lookup_override(&self.state, node_id)
321    }
322}
323
324// ---- Tracer body, factored so Recorder and Handle share it. ------
325
326fn push_call_frame(state: &Mutex<RecorderState>, node_id: &str, name: &str, args: &[Value]) {
327    let mut st = state.lock().unwrap();
328    st.open.push(OpenFrame {
329        node: TraceNode {
330            node_id: node_id.to_string(),
331            kind: TraceNodeKind::Call,
332            target: name.to_string(),
333            input: values_to_json(args),
334            output: None,
335            error: None,
336            started_at: now_unix(),
337            ended_at: 0,
338            children: Vec::new(),
339        },
340        children: Vec::new(),
341    });
342}
343
344fn push_effect_frame(state: &Mutex<RecorderState>, node_id: &str, kind: &str, op: &str, args: &[Value]) {
345    let mut st = state.lock().unwrap();
346    st.open.push(OpenFrame {
347        node: TraceNode {
348            node_id: node_id.to_string(),
349            kind: TraceNodeKind::Effect,
350            target: format!("{kind}.{op}"),
351            input: values_to_json(args),
352            output: None,
353            error: None,
354            started_at: now_unix(),
355            ended_at: 0,
356            children: Vec::new(),
357        },
358        children: Vec::new(),
359    });
360}
361
362fn exit_ok_frame(state: &Mutex<RecorderState>, value: &Value) {
363    let mut st = state.lock().unwrap();
364    if let Some(mut frame) = st.open.pop() {
365        frame.node.ended_at = now_unix();
366        frame.node.output = Some(value_to_json(value));
367        frame.node.children = frame.children;
368        attach_completed(&mut st, frame.node);
369    }
370}
371
372fn exit_err_frame(state: &Mutex<RecorderState>, message: &str) {
373    let mut st = state.lock().unwrap();
374    if let Some(mut frame) = st.open.pop() {
375        frame.node.ended_at = now_unix();
376        frame.node.error = Some(message.to_string());
377        frame.node.children = frame.children;
378        attach_completed(&mut st, frame.node);
379    }
380}
381
382fn exit_tail_frame(state: &Mutex<RecorderState>) {
383    let mut st = state.lock().unwrap();
384    if let Some(mut frame) = st.open.pop() {
385        frame.node.ended_at = now_unix();
386        frame.node.output = Some(serde_json::Value::Null);
387        frame.node.children = frame.children;
388        attach_completed(&mut st, frame.node);
389    }
390}
391
392fn lookup_override(state: &Mutex<RecorderState>, node_id: &str) -> Option<Value> {
393    let st = state.lock().unwrap();
394    st.overrides.get(node_id).map(json_to_value)
395}
396
397fn attach_completed(st: &mut RecorderState, node: TraceNode) {
398    if let Some(parent) = st.open.last_mut() {
399        parent.children.push(node);
400    } else {
401        st.completed.push(node);
402    }
403}