Skip to main content

entrenar/monitor/inference/serialization/
serializer.rs

1//! Trace serializer implementation.
2
3use super::error::SerializationError;
4use super::format::{PathType, TraceFormat, APRT_MAGIC, APRT_VERSION};
5use crate::monitor::inference::path::DecisionPath;
6use crate::monitor::inference::trace::DecisionTrace;
7use serde::{Deserialize, Serialize};
8
9/// Trace serializer
10pub struct TraceSerializer {
11    format: TraceFormat,
12}
13
14impl TraceSerializer {
15    /// Create a new serializer
16    pub fn new(format: TraceFormat) -> Self {
17        Self { format }
18    }
19
20    /// Get the format
21    pub fn format(&self) -> TraceFormat {
22        self.format
23    }
24
25    /// Serialize a trace to bytes
26    pub fn serialize<P: DecisionPath + Serialize>(
27        &self,
28        trace: &DecisionTrace<P>,
29        path_type: PathType,
30    ) -> Result<Vec<u8>, SerializationError> {
31        match self.format {
32            TraceFormat::Binary => self.serialize_binary(trace, path_type),
33            TraceFormat::Json => self.serialize_json(trace),
34            TraceFormat::JsonLines => self.serialize_json_line(trace),
35        }
36    }
37
38    /// Serialize to APRT binary format
39    fn serialize_binary<P: DecisionPath + Serialize>(
40        &self,
41        trace: &DecisionTrace<P>,
42        path_type: PathType,
43    ) -> Result<Vec<u8>, SerializationError> {
44        let trace_bytes = trace.to_bytes();
45
46        let mut bytes = Vec::with_capacity(8 + trace_bytes.len());
47
48        // Header
49        bytes.extend_from_slice(&APRT_MAGIC);
50        bytes.push(APRT_VERSION);
51        bytes.push(path_type as u8);
52        bytes.extend_from_slice(&[0, 0]); // Reserved (alignment)
53
54        // Trace data
55        bytes.extend_from_slice(&trace_bytes);
56
57        Ok(bytes)
58    }
59
60    /// Serialize to JSON
61    fn serialize_json<P: DecisionPath + Serialize>(
62        &self,
63        trace: &DecisionTrace<P>,
64    ) -> Result<Vec<u8>, SerializationError> {
65        serde_json::to_vec_pretty(trace).map_err(SerializationError::Json)
66    }
67
68    /// Serialize to JSON line (no trailing newline)
69    fn serialize_json_line<P: DecisionPath + Serialize>(
70        &self,
71        trace: &DecisionTrace<P>,
72    ) -> Result<Vec<u8>, SerializationError> {
73        let mut bytes = serde_json::to_vec(trace).map_err(SerializationError::Json)?;
74        bytes.push(b'\n');
75        Ok(bytes)
76    }
77
78    /// Deserialize from bytes
79    pub fn deserialize<P: DecisionPath + for<'de> Deserialize<'de>>(
80        &self,
81        bytes: &[u8],
82    ) -> Result<DecisionTrace<P>, SerializationError> {
83        match self.format {
84            TraceFormat::Binary => self.deserialize_binary(bytes),
85            TraceFormat::Json | TraceFormat::JsonLines => self.deserialize_json(bytes),
86        }
87    }
88
89    /// Deserialize from APRT binary format
90    fn deserialize_binary<P: DecisionPath>(
91        &self,
92        bytes: &[u8],
93    ) -> Result<DecisionTrace<P>, SerializationError> {
94        if bytes.len() < 8 {
95            return Err(SerializationError::InvalidFormat("Insufficient header bytes".to_string()));
96        }
97
98        // Check magic
99        if bytes[0..4] != APRT_MAGIC {
100            return Err(SerializationError::InvalidFormat("Invalid APRT magic bytes".to_string()));
101        }
102
103        // Check version
104        let version = bytes[4];
105        if version != APRT_VERSION {
106            return Err(SerializationError::VersionMismatch {
107                expected: APRT_VERSION,
108                actual: version,
109            });
110        }
111
112        // Skip path type and reserved bytes (bytes[5], bytes[6], bytes[7])
113
114        // Deserialize trace
115        DecisionTrace::from_bytes(bytes.get(8..).unwrap_or_default()).map_err(|e| {
116            SerializationError::InvalidFormat(format!("Path deserialization failed: {e}"))
117        })
118    }
119
120    /// Deserialize from JSON
121    fn deserialize_json<P: DecisionPath + for<'de> Deserialize<'de>>(
122        &self,
123        bytes: &[u8],
124    ) -> Result<DecisionTrace<P>, SerializationError> {
125        serde_json::from_slice(bytes).map_err(SerializationError::Json)
126    }
127}
128
129impl Default for TraceSerializer {
130    fn default() -> Self {
131        Self::new(TraceFormat::Binary)
132    }
133}