serializer/llm/
machine_zerocopy.rs

1//! Zero-copy machine format implementation
2//!
3//! Optimized binary format with minimal deserialization overhead
4
5use crate::llm::types::{DxDocument, DxLlmValue, DxSection};
6use std::collections::HashMap;
7use thiserror::Error;
8
9/// Errors returned by the zero-copy machine format reader and writer.
10#[derive(Debug, Error)]
11pub enum ZeroCopyError {
12    /// Magic bytes did not match the zero-copy format marker.
13    #[error("Invalid magic number")]
14    InvalidMagic,
15    /// Format version is newer or older than this implementation supports.
16    #[error("Unsupported version: {0}")]
17    UnsupportedVersion(u8),
18    /// Payload bytes are structurally invalid.
19    #[error("Invalid data: {0}")]
20    InvalidData(String),
21    /// A read attempted to move outside the payload boundary.
22    #[error("Out of bounds access")]
23    OutOfBounds,
24}
25
26const MAGIC: &[u8; 4] = b"DXZC";
27const VERSION: u8 = 1;
28
29/// Zero-copy machine format
30pub struct ZeroCopyMachine {
31    data: Vec<u8>,
32}
33
34impl ZeroCopyMachine {
35    /// Encode a document into the zero-copy machine format.
36    pub fn from_document(doc: &DxDocument) -> Self {
37        let mut data = Vec::new();
38
39        // Header
40        data.extend_from_slice(MAGIC);
41        data.push(VERSION);
42
43        // String table
44        let mut string_table: HashMap<&str, u32> = HashMap::new();
45        let mut string_data = Vec::new();
46
47        // Collect all strings
48        for key in doc.context.keys() {
49            if !string_table.contains_key(key.as_str()) {
50                let id = string_table.len() as u32;
51                string_table.insert(key, id);
52                string_data.push(key.as_bytes());
53            }
54        }
55        for value in doc.context.values() {
56            collect_strings_from_value(value, &mut string_table, &mut string_data);
57        }
58        for section in doc.sections.values() {
59            for col in &section.schema {
60                if !string_table.contains_key(col.as_str()) {
61                    let id = string_table.len() as u32;
62                    string_table.insert(col, id);
63                    string_data.push(col.as_bytes());
64                }
65            }
66            for row in &section.rows {
67                for value in row {
68                    collect_strings_from_value(value, &mut string_table, &mut string_data);
69                }
70            }
71        }
72
73        // Write string table
74        data.extend_from_slice(&(string_table.len() as u32).to_le_bytes());
75        for s in &string_data {
76            data.extend_from_slice(&(s.len() as u32).to_le_bytes());
77            data.extend_from_slice(s);
78        }
79
80        // Write context
81        data.extend_from_slice(&(doc.context.len() as u32).to_le_bytes());
82        for (key, value) in &doc.context {
83            let key_id = string_table[key.as_str()];
84            data.extend_from_slice(&key_id.to_le_bytes());
85            write_value(&mut data, value, &string_table);
86        }
87
88        // Write sections
89        data.extend_from_slice(&(doc.sections.len() as u32).to_le_bytes());
90        for (id, section) in &doc.sections {
91            data.push(*id as u8);
92            data.extend_from_slice(&(section.schema.len() as u32).to_le_bytes());
93            for col in &section.schema {
94                let col_id = string_table[col.as_str()];
95                data.extend_from_slice(&col_id.to_le_bytes());
96            }
97            data.extend_from_slice(&(section.rows.len() as u32).to_le_bytes());
98            for row in &section.rows {
99                for value in row {
100                    write_value(&mut data, value, &string_table);
101                }
102            }
103        }
104
105        Self { data }
106    }
107
108    /// Borrow the encoded machine-format bytes.
109    pub fn as_bytes(&self) -> &[u8] {
110        &self.data
111    }
112
113    /// Decode the machine-format bytes into a document.
114    pub fn to_document(&self) -> Result<DxDocument, ZeroCopyError> {
115        let mut pos = 0;
116        let data = &self.data;
117
118        // Check magic
119        if data.len() < 5 || &data[0..4] != MAGIC {
120            return Err(ZeroCopyError::InvalidMagic);
121        }
122        pos += 4;
123
124        let version = data[pos];
125        if version != VERSION {
126            return Err(ZeroCopyError::UnsupportedVersion(version));
127        }
128        pos += 1;
129
130        // Read string table
131        let string_count = read_u32(data, &mut pos)?;
132        let mut strings = Vec::with_capacity(string_count as usize);
133        for _ in 0..string_count {
134            let len = read_u32(data, &mut pos)? as usize;
135            if pos + len > data.len() {
136                return Err(ZeroCopyError::OutOfBounds);
137            }
138            let s = std::str::from_utf8(&data[pos..pos + len])
139                .map_err(|_| ZeroCopyError::InvalidData("Invalid UTF-8".to_string()))?
140                .to_string();
141            strings.push(s);
142            pos += len;
143        }
144
145        let mut doc = DxDocument::new();
146
147        // Read context
148        let context_count = read_u32(data, &mut pos)?;
149        for _ in 0..context_count {
150            let key_id = read_u32(data, &mut pos)? as usize;
151            if key_id >= strings.len() {
152                return Err(ZeroCopyError::OutOfBounds);
153            }
154            let key = strings[key_id].clone();
155            let value = read_value(data, &mut pos, &strings)?;
156            doc.context.insert(key, value);
157        }
158
159        // Read sections
160        let sections_count = read_u32(data, &mut pos)?;
161        for _ in 0..sections_count {
162            if pos >= data.len() {
163                return Err(ZeroCopyError::OutOfBounds);
164            }
165            let id = data[pos] as char;
166            pos += 1;
167
168            let schema_count = read_u32(data, &mut pos)? as usize;
169            let mut schema = Vec::with_capacity(schema_count);
170            for _ in 0..schema_count {
171                let col_id = read_u32(data, &mut pos)? as usize;
172                if col_id >= strings.len() {
173                    return Err(ZeroCopyError::OutOfBounds);
174                }
175                schema.push(strings[col_id].clone());
176            }
177
178            let mut section = DxSection::new(schema.clone());
179            let rows_count = read_u32(data, &mut pos)?;
180            for _ in 0..rows_count {
181                let mut row = Vec::with_capacity(schema.len());
182                for _ in 0..schema.len() {
183                    row.push(read_value(data, &mut pos, &strings)?);
184                }
185                section.rows.push(row);
186            }
187
188            doc.sections.insert(id, section);
189        }
190
191        Ok(doc)
192    }
193
194    /// Validate and borrow the encoded bytes through a document view.
195    pub fn access(&self) -> Result<ZeroCopyDocument<'_>, ZeroCopyError> {
196        ZeroCopyDocument::new(&self.data)
197    }
198}
199
200/// Borrowed view over a validated zero-copy machine-format payload.
201pub struct ZeroCopyDocument<'a> {
202    _data: &'a [u8],
203}
204
205impl<'a> ZeroCopyDocument<'a> {
206    fn new(data: &'a [u8]) -> Result<Self, ZeroCopyError> {
207        if data.len() < 5 || &data[0..4] != MAGIC {
208            return Err(ZeroCopyError::InvalidMagic);
209        }
210        if data[4] != VERSION {
211            return Err(ZeroCopyError::UnsupportedVersion(data[4]));
212        }
213        Ok(Self { _data: data })
214    }
215
216    /// Decode this borrowed payload into an owned document.
217    pub fn to_document(&self) -> Result<DxDocument, ZeroCopyError> {
218        ZeroCopyMachine {
219            data: self._data.to_vec(),
220        }
221        .to_document()
222    }
223}
224
225fn collect_strings_from_value<'a>(
226    value: &'a DxLlmValue,
227    table: &mut HashMap<&'a str, u32>,
228    data: &mut Vec<&'a [u8]>,
229) {
230    match value {
231        DxLlmValue::Str(s) => {
232            if !table.contains_key(s.as_str()) {
233                let id = table.len() as u32;
234                table.insert(s, id);
235                data.push(s.as_bytes());
236            }
237        }
238        DxLlmValue::Ref(s) => {
239            if !table.contains_key(s.as_str()) {
240                let id = table.len() as u32;
241                table.insert(s, id);
242                data.push(s.as_bytes());
243            }
244        }
245        DxLlmValue::Arr(items) => {
246            for item in items {
247                collect_strings_from_value(item, table, data);
248            }
249        }
250        DxLlmValue::Obj(fields) => {
251            for (k, v) in fields {
252                if !table.contains_key(k.as_str()) {
253                    let id = table.len() as u32;
254                    table.insert(k, id);
255                    data.push(k.as_bytes());
256                }
257                collect_strings_from_value(v, table, data);
258            }
259        }
260        _ => {}
261    }
262}
263
264fn write_value(data: &mut Vec<u8>, value: &DxLlmValue, string_table: &HashMap<&str, u32>) {
265    match value {
266        DxLlmValue::Str(s) => {
267            data.push(0);
268            let id = string_table[s.as_str()];
269            data.extend_from_slice(&id.to_le_bytes());
270        }
271        DxLlmValue::Num(n) => {
272            data.push(1);
273            data.extend_from_slice(&n.to_le_bytes());
274        }
275        DxLlmValue::Bool(b) => {
276            data.push(2);
277            data.push(if *b { 1 } else { 0 });
278        }
279        DxLlmValue::Null => {
280            data.push(3);
281        }
282        _ => {
283            data.push(3); // Treat complex types as null for now
284        }
285    }
286}
287
288fn read_u32(data: &[u8], pos: &mut usize) -> Result<u32, ZeroCopyError> {
289    if *pos + 4 > data.len() {
290        return Err(ZeroCopyError::OutOfBounds);
291    }
292    let bytes = [data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3]];
293    *pos += 4;
294    Ok(u32::from_le_bytes(bytes))
295}
296
297fn read_value(
298    data: &[u8],
299    pos: &mut usize,
300    strings: &[String],
301) -> Result<DxLlmValue, ZeroCopyError> {
302    if *pos >= data.len() {
303        return Err(ZeroCopyError::OutOfBounds);
304    }
305
306    let type_tag = data[*pos];
307    *pos += 1;
308
309    match type_tag {
310        0 => {
311            let id = read_u32(data, pos)? as usize;
312            if id >= strings.len() {
313                return Err(ZeroCopyError::OutOfBounds);
314            }
315            Ok(DxLlmValue::Str(strings[id].clone()))
316        }
317        1 => {
318            if *pos + 8 > data.len() {
319                return Err(ZeroCopyError::OutOfBounds);
320            }
321            let mut bytes = [0u8; 8];
322            bytes.copy_from_slice(&data[*pos..*pos + 8]);
323            *pos += 8;
324            Ok(DxLlmValue::Num(f64::from_le_bytes(bytes)))
325        }
326        2 => {
327            if *pos >= data.len() {
328                return Err(ZeroCopyError::OutOfBounds);
329            }
330            let b = data[*pos] != 0;
331            *pos += 1;
332            Ok(DxLlmValue::Bool(b))
333        }
334        3 => Ok(DxLlmValue::Null),
335        _ => Err(ZeroCopyError::InvalidData(format!(
336            "Unknown type tag: {}",
337            type_tag
338        ))),
339    }
340}
serializer/llm/machine_zerocopy.rs

serializer/llm/
machine_zerocopy.rs
