1use 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 pub target: String,
36 pub input: serde_json::Value,
37 #[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 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
78pub struct Recorder {
81 state: Arc<Mutex<RecorderState>>,
82}
83
84pub(crate) struct RecorderState {
85 open: Vec<OpenFrame>,
90 completed: Vec<TraceNode>,
93 pub(crate) overrides: IndexMap<String, serde_json::Value>,
95}
96
97struct OpenFrame {
98 node: TraceNode,
99 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 pub fn handle(&self) -> Handle {
116 Handle { state: Arc::clone(&self.state) }
117 }
118
119 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 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.clone()),
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 { m.insert(k.clone(), value_to_json(v)); }
183 J::Object(m)
184 }
185 Value::Variant { name, args } => {
186 let mut m = serde_json::Map::new();
187 m.insert("$variant".into(), J::String(name.clone()));
188 m.insert("args".into(), J::Array(args.iter().map(value_to_json).collect()));
189 J::Object(m)
190 }
191 Value::Closure { body_hash, .. } => {
192 let prefix: String = body_hash.iter().take(4)
197 .map(|b| format!("{b:02x}")).collect();
198 J::String(format!("<closure {prefix}>"))
199 }
200 Value::F64Array { rows, cols, data } => {
201 let mut m = serde_json::Map::new();
202 m.insert("$f64_array".into(), J::Bool(true));
203 m.insert("rows".into(), J::from(*rows));
204 m.insert("cols".into(), J::from(*cols));
205 m.insert("data".into(), J::Array(data.iter().map(|f| J::from(*f)).collect()));
206 J::Object(m)
207 }
208 Value::Map(m) => {
209 let mut o = serde_json::Map::new();
210 o.insert("$map".into(), J::Bool(true));
211 o.insert("entries".into(), J::Array(m.iter().map(|(k, v)| {
212 J::Array(vec![value_to_json(&k.as_value()), value_to_json(v)])
213 }).collect()));
214 J::Object(o)
215 }
216 Value::Set(s) => {
217 let mut o = serde_json::Map::new();
218 o.insert("$set".into(), J::Bool(true));
219 o.insert("items".into(), J::Array(
220 s.iter().map(|k| value_to_json(&k.as_value())).collect()));
221 J::Object(o)
222 }
223 Value::Deque(items) => {
224 let mut o = serde_json::Map::new();
225 o.insert("$deque".into(), J::Bool(true));
226 o.insert("items".into(), J::Array(
227 items.iter().map(value_to_json).collect()));
228 J::Object(o)
229 }
230 }
231}
232
233pub(crate) fn json_to_value(v: &serde_json::Value) -> Value {
234 use serde_json::Value as J;
235 match v {
236 J::Null => Value::Unit,
237 J::Bool(b) => Value::Bool(*b),
238 J::Number(n) => {
239 if let Some(i) = n.as_i64() { Value::Int(i) }
240 else if let Some(f) = n.as_f64() { Value::Float(f) }
241 else { Value::Unit }
242 }
243 J::String(s) => Value::Str(s.clone()),
244 J::Array(items) => Value::List(items.iter().map(json_to_value).collect()),
245 J::Object(map) => {
246 if let (Some(serde_json::Value::String(name)), Some(serde_json::Value::Array(args))) =
248 (map.get("$variant"), map.get("args"))
249 {
250 return Value::Variant {
251 name: name.clone(),
252 args: args.iter().map(json_to_value).collect(),
253 };
254 }
255 let mut out = indexmap::IndexMap::new();
256 for (k, v) in map { out.insert(k.clone(), json_to_value(v)); }
257 Value::Record(out)
258 }
259 }
260}
261
262impl Tracer for Recorder {
263 fn enter_call(&mut self, node_id: &str, name: &str, args: &[Value]) {
264 push_call_frame(&self.state, node_id, name, args);
265 }
266 fn enter_effect(&mut self, node_id: &str, kind: &str, op: &str, args: &[Value]) {
267 push_effect_frame(&self.state, node_id, kind, op, args);
268 }
269 fn exit_ok(&mut self, value: &Value) { exit_ok_frame(&self.state, value); }
270 fn exit_err(&mut self, message: &str) { exit_err_frame(&self.state, message); }
271 fn exit_call_tail(&mut self) { exit_tail_frame(&self.state); }
272 fn override_effect(&mut self, node_id: &str) -> Option<Value> {
273 lookup_override(&self.state, node_id)
274 }
275}
276
277impl Tracer for Handle {
283 fn enter_call(&mut self, node_id: &str, name: &str, args: &[Value]) {
284 push_call_frame(&self.state, node_id, name, args);
285 }
286 fn enter_effect(&mut self, node_id: &str, kind: &str, op: &str, args: &[Value]) {
287 push_effect_frame(&self.state, node_id, kind, op, args);
288 }
289 fn exit_ok(&mut self, value: &Value) { exit_ok_frame(&self.state, value); }
290 fn exit_err(&mut self, message: &str) { exit_err_frame(&self.state, message); }
291 fn exit_call_tail(&mut self) { exit_tail_frame(&self.state); }
292 fn override_effect(&mut self, node_id: &str) -> Option<Value> {
293 lookup_override(&self.state, node_id)
294 }
295}
296
297fn push_call_frame(state: &Mutex<RecorderState>, node_id: &str, name: &str, args: &[Value]) {
300 let mut st = state.lock().unwrap();
301 st.open.push(OpenFrame {
302 node: TraceNode {
303 node_id: node_id.to_string(),
304 kind: TraceNodeKind::Call,
305 target: name.to_string(),
306 input: values_to_json(args),
307 output: None,
308 error: None,
309 started_at: now_unix(),
310 ended_at: 0,
311 children: Vec::new(),
312 },
313 children: Vec::new(),
314 });
315}
316
317fn push_effect_frame(state: &Mutex<RecorderState>, node_id: &str, kind: &str, op: &str, args: &[Value]) {
318 let mut st = state.lock().unwrap();
319 st.open.push(OpenFrame {
320 node: TraceNode {
321 node_id: node_id.to_string(),
322 kind: TraceNodeKind::Effect,
323 target: format!("{kind}.{op}"),
324 input: values_to_json(args),
325 output: None,
326 error: None,
327 started_at: now_unix(),
328 ended_at: 0,
329 children: Vec::new(),
330 },
331 children: Vec::new(),
332 });
333}
334
335fn exit_ok_frame(state: &Mutex<RecorderState>, value: &Value) {
336 let mut st = state.lock().unwrap();
337 if let Some(mut frame) = st.open.pop() {
338 frame.node.ended_at = now_unix();
339 frame.node.output = Some(value_to_json(value));
340 frame.node.children = frame.children;
341 attach_completed(&mut st, frame.node);
342 }
343}
344
345fn exit_err_frame(state: &Mutex<RecorderState>, message: &str) {
346 let mut st = state.lock().unwrap();
347 if let Some(mut frame) = st.open.pop() {
348 frame.node.ended_at = now_unix();
349 frame.node.error = Some(message.to_string());
350 frame.node.children = frame.children;
351 attach_completed(&mut st, frame.node);
352 }
353}
354
355fn exit_tail_frame(state: &Mutex<RecorderState>) {
356 let mut st = state.lock().unwrap();
357 if let Some(mut frame) = st.open.pop() {
358 frame.node.ended_at = now_unix();
359 frame.node.output = Some(serde_json::Value::Null);
360 frame.node.children = frame.children;
361 attach_completed(&mut st, frame.node);
362 }
363}
364
365fn lookup_override(state: &Mutex<RecorderState>, node_id: &str) -> Option<Value> {
366 let st = state.lock().unwrap();
367 st.overrides.get(node_id).map(json_to_value)
368}
369
370fn attach_completed(st: &mut RecorderState, node: TraceNode) {
371 if let Some(parent) = st.open.last_mut() {
372 parent.children.push(node);
373 } else {
374 st.completed.push(node);
375 }
376}