Skip to main content

marser_trace_schema/
session.rs

1//! AI assistance: this file was written with AI assistance. The maintainer reviewed it and did not find errors.
2
3use std::io::{self, Write};
4
5use serde_json::json;
6
7use crate::event::NodeTrace;
8use crate::version::SCHEMA_VERSION;
9
10#[derive(Clone, Debug, Default)]
11pub struct TraceSession {
12    nodes: Vec<NodeTrace>,
13    source_text: Option<String>,
14    dropped_events: usize,
15    max_events: Option<usize>,
16}
17
18impl TraceSession {
19    pub fn new() -> Self {
20        Self::default()
21    }
22
23    pub fn with_max_events(max_events: usize) -> Self {
24        Self {
25            nodes: Vec::new(),
26            source_text: None,
27            dropped_events: 0,
28            max_events: Some(max_events),
29        }
30    }
31
32    pub fn record(&mut self, node: NodeTrace) {
33        if let Some(max_events) = self.max_events
34            && self.nodes.len() >= max_events
35        {
36            self.dropped_events = self.dropped_events.saturating_add(1);
37            return;
38        }
39        self.nodes.push(node);
40    }
41
42    pub fn nodes(&self) -> &[NodeTrace] {
43        &self.nodes
44    }
45
46    pub fn events(&self) -> &[NodeTrace] {
47        self.nodes()
48    }
49
50    pub fn dropped_events(&self) -> usize {
51        self.dropped_events
52    }
53
54    pub fn source_text(&self) -> Option<&str> {
55        self.source_text.as_deref()
56    }
57
58    pub fn set_source_text(&mut self, source_text: impl Into<String>) {
59        self.source_text = Some(source_text.into());
60    }
61
62    pub fn write_json<W: Write>(&self, mut writer: W) -> io::Result<()> {
63        serde_json::to_writer(
64            &mut writer,
65            &json!({
66                "trace_version": SCHEMA_VERSION,
67                "nodes": self.nodes,
68                "source_text": self.source_text,
69            }),
70        )?;
71        Ok(())
72    }
73
74    pub fn write_jsonl<W: Write>(&self, mut writer: W) -> io::Result<()> {
75        for node in &self.nodes {
76            serde_json::to_writer(&mut writer, node)?;
77            writer.write_all(b"\n")?;
78        }
79        Ok(())
80    }
81
82    pub fn from_events(nodes: Vec<NodeTrace>) -> Self {
83        Self {
84            nodes,
85            source_text: None,
86            dropped_events: 0,
87            max_events: None,
88        }
89    }
90}