Skip to main content

zon_lib/
writer.rs

1use crate::header::ZonHeader;
2
3pub struct ZonWriter {
4    pub(crate) buffer: Vec<u8>,
5}
6
7impl ZonWriter {
8    pub fn new() -> Self {
9        // start with a default header
10        let mut writer = Self {
11            buffer: Vec::with_capacity(4096),
12        };
13        
14        // write the header immediately
15        let header = ZonHeader::default();
16        
17        // safety: ZonHeader is POD and repr(C).
18        let header_slice = unsafe {
19            std::slice::from_raw_parts(
20                &header as *const ZonHeader as *const u8,
21                std::mem::size_of::<ZonHeader>(),
22            )
23        };
24        
25        writer.buffer.extend_from_slice(header_slice);
26        writer
27    }
28
29    pub fn len(&self) -> usize {
30        self.buffer.len()
31    }
32
33    pub fn is_empty(&self) -> bool {
34        self.buffer.is_empty()
35    }
36    
37    pub fn as_bytes(&self) -> &[u8] {
38        &self.buffer
39    }
40
41    /// appends the 4 bytes of val to the buffer.
42    /// returns the offset (index) where those bytes were written.
43    pub fn write_u32(&mut self, val: u32) -> u32 {
44        let offset = self.buffer.len() as u32;
45        self.buffer.extend_from_slice(&val.to_le_bytes());
46        offset
47    }
48
49    /// first, append a 4-byte length (u32).
50    /// then, append the raw string bytes.
51    /// crucial: append padding zeros until the buffer's total size is a multiple of 4 bytes.
52    /// returns the offset where the length was written.
53    pub fn write_string(&mut self, val: &str) -> u32 {
54        let start_offset = self.buffer.len() as u32;
55        let len = val.len() as u32;
56        
57        // write length
58        self.buffer.extend_from_slice(&len.to_le_bytes());
59        
60        // write string bytes
61        self.buffer.extend_from_slice(val.as_bytes());
62        
63        // add padding
64        let current_len = self.buffer.len();
65        let padding_needed = (4 - (current_len % 4)) % 4;
66        for _ in 0..padding_needed {
67            self.buffer.push(0);
68        }
69        
70        start_offset
71    }
72
73    /// updates the root offset in the header.
74    /// the header is always at the start of the buffer.
75    pub fn set_root(&mut self, offset: u32) {
76        // root is at offset 8 (magic=0..4, version=4..8, root=8..12)
77        if self.buffer.len() >= 12 {
78            let bytes = offset.to_le_bytes();
79            self.buffer[8] = bytes[0];
80            self.buffer[9] = bytes[1];
81            self.buffer[10] = bytes[2];
82            self.buffer[11] = bytes[3];
83        }
84    }
85}
86
87impl Default for ZonWriter {
88    fn default() -> Self {
89        Self::new()
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    #[test]
98    fn test_writer_initialization() {
99        let writer = ZonWriter::new();
100        assert_eq!(writer.len(), 64);
101        
102        let bytes = writer.as_bytes();
103        // check magic number at offset 0 (little endian)
104        assert_eq!(bytes[0], 0x21);
105        assert_eq!(bytes[1], 0x4E);
106        assert_eq!(bytes[2], 0x4F);
107        assert_eq!(bytes[3], 0x5A);
108    }
109
110    #[test]
111    fn test_write_primitive_and_string() {
112        let mut writer = ZonWriter::new();
113        
114        // write a u32
115        let u32_val = 0x12345678;
116        let u32_offset = writer.write_u32(u32_val);
117        
118        // expect header is 64 bytes
119        assert_eq!(u32_offset, 64);
120        
121        // write a string "hello" (5 bytes)
122        // length 5 (4 bytes) + "hello" (5 bytes) = 9 bytes.
123        // padding needed: (4 - (9 % 4)) % 4 = (4 - 1) = 3 bytes padding.
124        // total added: 4 + 5 + 3 = 12 bytes.
125        let string_val = "hello";
126        let str_offset = writer.write_string(string_val);
127        
128        assert_eq!(str_offset, 64 + 4); // 68
129        
130        // check buffer content
131        {
132            let bytes = writer.as_bytes();
133            
134            // check u32 at 64
135            let u32_slice = &bytes[64..68];
136            assert_eq!(u32::from_le_bytes(u32_slice.try_into().unwrap()), u32_val);
137            
138            // check string length at 68
139            let len_slice = &bytes[68..72];
140            assert_eq!(u32::from_le_bytes(len_slice.try_into().unwrap()), 5);
141            
142            // check string bytes at 72
143            let str_slice = &bytes[72..77];
144            assert_eq!(str_slice, b"hello");
145            
146            // check padding at 77..80
147            assert_eq!(bytes[77], 0);
148            assert_eq!(bytes[78], 0);
149            assert_eq!(bytes[79], 0);
150            
151            // check alignment
152            assert_eq!(bytes.len() % 4, 0);
153        }
154        
155        // set root and verify
156        writer.set_root(str_offset);
157        
158        {
159            let bytes = writer.as_bytes();
160            // header root is at offset 8
161            let root_slice = &bytes[8..12];
162            assert_eq!(u32::from_le_bytes(root_slice.try_into().unwrap()), str_offset);
163        }
164    }
165}