rustyhdf5_format/
object_header_writer.rs

1//! Object header writer for v2 format.
2
3#[cfg(not(feature = "std"))]
4use alloc::vec::Vec;
5
6use crate::checksum::jenkins_lookup3;
7use crate::message_type::MessageType;
8
9/// Writer for v2 object headers with proper checksums.
10pub struct ObjectHeaderWriter {
11    messages: Vec<(MessageType, Vec<u8>, u8)>,  // (type, data, msg_flags)
12}
13
14impl ObjectHeaderWriter {
15    /// Create a new empty object header writer.
16    pub fn new() -> Self {
17        Self {
18            messages: Vec::new(),
19        }
20    }
21
22    /// Add a message to the header with default flags (0).
23    pub fn add_message(&mut self, msg_type: MessageType, data: Vec<u8>) {
24        self.messages.push((msg_type, data, 0));
25    }
26
27    /// Add a message with specific flags.
28    pub fn add_message_with_flags(&mut self, msg_type: MessageType, data: Vec<u8>, flags: u8) {
29        self.messages.push((msg_type, data, flags));
30    }
31
32    /// Serialize the complete v2 object header (OHDR + messages + checksum).
33    pub fn serialize(&self) -> Vec<u8> {
34        // Calculate total message bytes: each message has type(1) + size(2) + flags(1) + data
35        let msg_bytes_total: usize = self.messages.iter()
36            .map(|(_, data, _)| 4 + data.len())
37            .sum();
38
39        // Determine chunk size field width based on msg_bytes_total
40        let (flags, chunk_size_width) = if msg_bytes_total <= 255 {
41            (0x00u8, 1usize)
42        } else if msg_bytes_total <= 65535 {
43            (0x01u8, 2)
44        } else {
45            (0x02u8, 4)
46        };
47
48        let mut buf = Vec::new();
49
50        // OHDR signature
51        buf.extend_from_slice(b"OHDR");
52        // version
53        buf.push(2);
54        // flags
55        buf.push(flags);
56        // chunk0 size
57        match chunk_size_width {
58            1 => buf.push(msg_bytes_total as u8),
59            2 => buf.extend_from_slice(&(msg_bytes_total as u16).to_le_bytes()),
60            4 => buf.extend_from_slice(&(msg_bytes_total as u32).to_le_bytes()),
61            _ => {}
62        }
63
64        // Messages
65        for (msg_type, data, msg_flags) in &self.messages {
66            buf.push(msg_type.to_u16() as u8); // type (1 byte in v2)
67            buf.extend_from_slice(&(data.len() as u16).to_le_bytes()); // size (2 bytes)
68            buf.push(*msg_flags); // flags
69            buf.extend_from_slice(data);
70        }
71
72        // Checksum
73        let checksum = jenkins_lookup3(&buf);
74        buf.extend_from_slice(&checksum.to_le_bytes());
75
76        buf
77    }
78}
79
80impl Default for ObjectHeaderWriter {
81    fn default() -> Self {
82        Self::new()
83    }
84}
85
86/// A deferred header entry for batch writing.
87struct DeferredHeader {
88    writer: ObjectHeaderWriter,
89}
90
91/// Batch writer that collects multiple object headers in memory and flushes
92/// them as a single contiguous I/O pass.
93///
94/// This reduces the number of serialization passes when creating many datasets
95/// in parallel — each thread builds its `ObjectHeaderWriter` independently,
96/// then all headers are serialized together.
97pub struct BatchObjectHeaderWriter {
98    headers: Vec<DeferredHeader>,
99}
100
101impl BatchObjectHeaderWriter {
102    /// Create a new empty batch writer.
103    pub fn new() -> Self {
104        Self {
105            headers: Vec::new(),
106        }
107    }
108
109    /// Add a pre-built ObjectHeaderWriter to the batch.
110    pub fn add(&mut self, writer: ObjectHeaderWriter) {
111        self.headers.push(DeferredHeader { writer });
112    }
113
114    /// Number of headers in the batch.
115    pub fn len(&self) -> usize {
116        self.headers.len()
117    }
118
119    /// Whether the batch is empty.
120    pub fn is_empty(&self) -> bool {
121        self.headers.is_empty()
122    }
123
124    /// Compute the serialized size of each header without actually serializing.
125    /// Returns sizes in the same order as headers were added.
126    pub fn compute_sizes(&self) -> Vec<usize> {
127        self.headers
128            .iter()
129            .map(|h| h.writer.serialize().len())
130            .collect()
131    }
132
133    /// Serialize all headers into a single contiguous buffer.
134    /// Returns `(combined_bytes, offsets)` where `offsets[i]` is the byte
135    /// offset of header `i` within the combined buffer.
136    pub fn serialize_all(&self) -> (Vec<u8>, Vec<usize>) {
137        let serialized: Vec<Vec<u8>> = self.headers.iter().map(|h| h.writer.serialize()).collect();
138        let total: usize = serialized.iter().map(|s| s.len()).sum();
139        let mut buf = Vec::with_capacity(total);
140        let mut offsets = Vec::with_capacity(serialized.len());
141        for s in &serialized {
142            offsets.push(buf.len());
143            buf.extend_from_slice(s);
144        }
145        (buf, offsets)
146    }
147}
148
149impl Default for BatchObjectHeaderWriter {
150    fn default() -> Self {
151        Self::new()
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158    use crate::object_header::ObjectHeader;
159
160    #[test]
161    fn empty_header_roundtrip() {
162        let writer = ObjectHeaderWriter::new();
163        let bytes = writer.serialize();
164        let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
165        assert_eq!(hdr.version, 2);
166        assert_eq!(hdr.messages.len(), 0);
167    }
168
169    #[test]
170    fn two_messages_roundtrip() {
171        let mut writer = ObjectHeaderWriter::new();
172        writer.add_message(MessageType::Dataspace, vec![1, 2, 3, 4]);
173        writer.add_message(MessageType::Datatype, vec![5, 6]);
174        let bytes = writer.serialize();
175        let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
176        assert_eq!(hdr.messages.len(), 2);
177        assert_eq!(hdr.messages[0].msg_type, MessageType::Dataspace);
178        assert_eq!(hdr.messages[0].data, vec![1, 2, 3, 4]);
179        assert_eq!(hdr.messages[1].msg_type, MessageType::Datatype);
180        assert_eq!(hdr.messages[1].data, vec![5, 6]);
181    }
182
183    #[test]
184    fn large_header_uses_2byte_chunk_size() {
185        let mut writer = ObjectHeaderWriter::new();
186        // Add a message with >255 bytes of payload
187        writer.add_message(MessageType::Datatype, vec![0xAA; 300]);
188        let bytes = writer.serialize();
189        let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
190        assert_eq!(hdr.messages.len(), 1);
191        assert_eq!(hdr.messages[0].data.len(), 300);
192    }
193
194    #[test]
195    fn batch_writer_serialize_all() {
196        let mut batch = BatchObjectHeaderWriter::new();
197
198        let mut w1 = ObjectHeaderWriter::new();
199        w1.add_message(MessageType::Dataspace, vec![1, 2, 3]);
200
201        let mut w2 = ObjectHeaderWriter::new();
202        w2.add_message(MessageType::Datatype, vec![4, 5]);
203
204        batch.add(w1);
205        batch.add(w2);
206        assert_eq!(batch.len(), 2);
207
208        let (buf, offsets) = batch.serialize_all();
209        assert_eq!(offsets.len(), 2);
210        assert_eq!(offsets[0], 0);
211
212        // Parse each header from the combined buffer
213        let h1 = ObjectHeader::parse(&buf, offsets[0], 8, 8).unwrap();
214        assert_eq!(h1.messages.len(), 1);
215        assert_eq!(h1.messages[0].msg_type, MessageType::Dataspace);
216
217        let h2 = ObjectHeader::parse(&buf, offsets[1], 8, 8).unwrap();
218        assert_eq!(h2.messages.len(), 1);
219        assert_eq!(h2.messages[0].msg_type, MessageType::Datatype);
220    }
221
222    #[test]
223    fn batch_writer_empty() {
224        let batch = BatchObjectHeaderWriter::new();
225        assert!(batch.is_empty());
226        let (buf, offsets) = batch.serialize_all();
227        assert!(buf.is_empty());
228        assert!(offsets.is_empty());
229    }
230}
rustyhdf5_format/object_header_writer.rs

rustyhdf5_format/
object_header_writer.rs
