clawhdf5/crates/clawhdf5-format/src/object_header_writer.rs

//! Object header writer for v2 format.

#[cfg(not(feature = "std"))]
use alloc::{format, vec::Vec};

use crate::checksum::jenkins_lookup3;
use crate::error::FormatError;
use crate::message_type::MessageType;

/// Largest message payload a v2 object header can describe: the per-message
/// size field is 2 bytes. A bigger message cannot be encoded at all — writing
/// its size truncated to 16 bits produced files libhdf5 refuses.
pub const MAX_MESSAGE_SIZE: usize = u16::MAX as usize;

/// Object header flags: attribute creation order tracked (each message
/// then carries a 2-byte creation order) and indexed.
const OHDR_ATTR_CRT_ORDER_TRACKED: u8 = 0x04;
const OHDR_ATTR_CRT_ORDER_INDEXED: u8 = 0x08;

/// Writer for v2 object headers with proper checksums.
pub struct ObjectHeaderWriter {
    messages: Vec<(MessageType, Vec<u8>, u8, u16)>, // (type, data, msg_flags, creation order)
    /// Attribute creation order tracked and indexed.
    attr_order: bool,
}

impl ObjectHeaderWriter {
    /// Create a new empty object header writer.
    pub fn new() -> Self {
        Self {
            messages: Vec::new(),
            attr_order: false,
        }
    }

    /// Track and index attribute creation order, as libhdf5 does for an
    /// object created with `H5P_CRT_ORDER_TRACKED | H5P_CRT_ORDER_INDEXED`
    /// (h5py's `track_order=True`): the header's flags say so, and every
    /// message carries a creation order (an attribute's own; 0 for the
    /// others). libhdf5 reads the setting back from these flags.
    pub fn track_attr_order(&mut self) {
        self.attr_order = true;
    }

    /// Add a message to the header with default flags (0).
    pub fn add_message(&mut self, msg_type: MessageType, data: Vec<u8>) {
        self.messages.push((msg_type, data, 0, 0));
    }

    /// Add a message with specific flags.
    pub fn add_message_with_flags(&mut self, msg_type: MessageType, data: Vec<u8>, flags: u8) {
        self.messages.push((msg_type, data, flags, 0));
    }

    /// Add a message with its creation order, which is written only when
    /// attribute creation order is tracked ([`Self::track_attr_order`]).
    pub fn add_message_with_order(&mut self, msg_type: MessageType, data: Vec<u8>, order: u16) {
        self.messages.push((msg_type, data, 0, order));
    }

    /// Serialize the complete v2 object header (OHDR + messages + checksum).
    ///
    /// Fails with [`FormatError::SerializationError`] when a message is larger
    /// than [`MAX_MESSAGE_SIZE`] (e.g. an attribute over ~64 KiB, which would
    /// need dense attribute storage), rather than writing a corrupt header.
    pub fn serialize(&self) -> Result<Vec<u8>, FormatError> {
        if let Some((msg_type, data, _, _)) = self
            .messages
            .iter()
            .find(|(_, data, _, _)| data.len() > MAX_MESSAGE_SIZE)
        {
            return Err(FormatError::SerializationError(format!(
                "{msg_type:?} message is {} bytes; an object header message holds at most \
                 {MAX_MESSAGE_SIZE} bytes",
                data.len()
            )));
        }
        // Calculate total message bytes: each message has type(1) + size(2) +
        // flags(1) [+ creation order(2)] + data
        let msg_header = if self.attr_order { 6 } else { 4 };
        let msg_bytes_total: usize = self
            .messages
            .iter()
            .map(|(_, data, _, _)| msg_header + data.len())
            .sum();

        // Determine chunk size field width based on msg_bytes_total
        let (flags, chunk_size_width) = if msg_bytes_total <= 255 {
            (0x00u8, 1usize)
        } else if msg_bytes_total <= 65535 {
            (0x01u8, 2)
        } else {
            (0x02u8, 4)
        };

        let flags = if self.attr_order {
            flags | OHDR_ATTR_CRT_ORDER_TRACKED | OHDR_ATTR_CRT_ORDER_INDEXED
        } else {
            flags
        };

        let mut buf = Vec::new();

        // OHDR signature
        buf.extend_from_slice(b"OHDR");
        // version
        buf.push(2);
        // flags
        buf.push(flags);
        // chunk0 size
        match chunk_size_width {
            1 => buf.push(msg_bytes_total as u8),
            2 => buf.extend_from_slice(&(msg_bytes_total as u16).to_le_bytes()),
            4 => buf.extend_from_slice(&(msg_bytes_total as u32).to_le_bytes()),
            _ => unreachable!("unexpected chunk_size_width: {chunk_size_width}"),
        }

        // Messages
        for (msg_type, data, msg_flags, order) in &self.messages {
            let type_id = msg_type.to_u16();
            assert!(
                type_id <= 255,
                "v2 object header message type {type_id:#06x} exceeds u8 range (max 0xFF)"
            );
            buf.push(type_id as u8); // type (1 byte in v2)
            buf.extend_from_slice(&(data.len() as u16).to_le_bytes()); // size (2 bytes)
            buf.push(*msg_flags); // flags
            if self.attr_order {
                buf.extend_from_slice(&order.to_le_bytes()); // creation order
            }
            buf.extend_from_slice(data);
        }

        // Checksum
        let checksum = jenkins_lookup3(&buf);
        buf.extend_from_slice(&checksum.to_le_bytes());

        Ok(buf)
    }
}

impl Default for ObjectHeaderWriter {
    fn default() -> Self {
        Self::new()
    }
}

/// Batch writer that collects multiple object headers in memory and flushes
/// them as a single contiguous I/O pass.
///
/// This reduces the number of serialization passes when creating many datasets
/// in parallel -- each thread builds its `ObjectHeaderWriter` independently,
/// then all headers are serialized together.
pub struct BatchObjectHeaderWriter {
    headers: Vec<ObjectHeaderWriter>,
}

impl BatchObjectHeaderWriter {
    /// Create a new empty batch writer.
    pub fn new() -> Self {
        Self {
            headers: Vec::new(),
        }
    }

    /// Add a pre-built ObjectHeaderWriter to the batch.
    pub fn add(&mut self, writer: ObjectHeaderWriter) {
        self.headers.push(writer);
    }

    /// Number of headers in the batch.
    pub fn len(&self) -> usize {
        self.headers.len()
    }

    /// Whether the batch is empty.
    pub fn is_empty(&self) -> bool {
        self.headers.is_empty()
    }

    /// Compute the serialized size of each header without actually serializing.
    /// Returns sizes in the same order as headers were added.
    pub fn compute_sizes(&self) -> Result<Vec<usize>, FormatError> {
        self.headers
            .iter()
            .map(|h| h.serialize().map(|b| b.len()))
            .collect()
    }

    /// Serialize all headers into a single contiguous buffer.
    /// Returns `(combined_bytes, offsets)` where `offsets[i]` is the byte
    /// offset of header `i` within the combined buffer.
    pub fn serialize_all(&self) -> Result<(Vec<u8>, Vec<usize>), FormatError> {
        let serialized: Vec<Vec<u8>> = self
            .headers
            .iter()
            .map(|h| h.serialize())
            .collect::<Result<_, _>>()?;
        let total: usize = serialized.iter().map(|s| s.len()).sum();
        let mut buf = Vec::with_capacity(total);
        let mut offsets = Vec::with_capacity(serialized.len());
        for s in &serialized {
            offsets.push(buf.len());
            buf.extend_from_slice(s);
        }
        Ok((buf, offsets))
    }
}

impl Default for BatchObjectHeaderWriter {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::object_header::ObjectHeader;

    #[test]
    fn empty_header_roundtrip() {
        let writer = ObjectHeaderWriter::new();
        let bytes = writer.serialize().unwrap();
        let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
        assert_eq!(hdr.version, 2);
        assert_eq!(hdr.messages.len(), 0);
    }

    #[test]
    fn tracked_attribute_order_is_in_the_flags_and_every_message() {
        let mut writer = ObjectHeaderWriter::new();
        writer.track_attr_order();
        writer.add_message(MessageType::Dataspace, vec![1, 2, 3, 4]);
        writer.add_message_with_order(MessageType::Attribute, vec![5, 6], 7);
        let bytes = writer.serialize().unwrap();
        assert_eq!(bytes[5] & 0x0C, 0x0C);
        let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
        assert_eq!(hdr.messages.len(), 2);
        assert_eq!(hdr.messages[0].creation_order, Some(0));
        assert_eq!(hdr.messages[1].creation_order, Some(7));
        assert_eq!(hdr.messages[1].data, vec![5, 6]);
    }

    #[test]
    fn two_messages_roundtrip() {
        let mut writer = ObjectHeaderWriter::new();
        writer.add_message(MessageType::Dataspace, vec![1, 2, 3, 4]);
        writer.add_message(MessageType::Datatype, vec![5, 6]);
        let bytes = writer.serialize().unwrap();
        let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
        assert_eq!(hdr.messages.len(), 2);
        assert_eq!(hdr.messages[0].msg_type, MessageType::Dataspace);
        assert_eq!(hdr.messages[0].data, vec![1, 2, 3, 4]);
        assert_eq!(hdr.messages[1].msg_type, MessageType::Datatype);
        assert_eq!(hdr.messages[1].data, vec![5, 6]);
    }

    #[test]
    fn large_header_uses_2byte_chunk_size() {
        let mut writer = ObjectHeaderWriter::new();
        // Add a message with >255 bytes of payload
        writer.add_message(MessageType::Datatype, vec![0xAA; 300]);
        let bytes = writer.serialize().unwrap();
        let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
        assert_eq!(hdr.messages.len(), 1);
        assert_eq!(hdr.messages[0].data.len(), 300);
    }

    #[test]
    fn oversized_message_is_an_error_not_a_truncated_size() {
        // 65535 bytes is the largest encodable payload.
        let mut writer = ObjectHeaderWriter::new();
        writer.add_message(MessageType::Attribute, vec![0; MAX_MESSAGE_SIZE]);
        let bytes = writer.serialize().unwrap();
        let hdr = ObjectHeader::parse(&bytes, 0, 8, 8).unwrap();
        assert_eq!(hdr.messages[0].data.len(), MAX_MESSAGE_SIZE);

        // One byte more used to be written with its size wrapped to 0.
        let mut writer = ObjectHeaderWriter::new();
        writer.add_message(MessageType::Attribute, vec![0; MAX_MESSAGE_SIZE + 1]);
        assert!(matches!(
            writer.serialize(),
            Err(FormatError::SerializationError(_))
        ));
    }

    #[test]
    fn batch_writer_serialize_all() {
        let mut batch = BatchObjectHeaderWriter::new();

        let mut w1 = ObjectHeaderWriter::new();
        w1.add_message(MessageType::Dataspace, vec![1, 2, 3]);

        let mut w2 = ObjectHeaderWriter::new();
        w2.add_message(MessageType::Datatype, vec![4, 5]);

        batch.add(w1);
        batch.add(w2);
        assert_eq!(batch.len(), 2);

        let (buf, offsets) = batch.serialize_all().unwrap();
        assert_eq!(offsets.len(), 2);
        assert_eq!(offsets[0], 0);

        // Parse each header from the combined buffer
        let h1 = ObjectHeader::parse(&buf, offsets[0], 8, 8).unwrap();
        assert_eq!(h1.messages.len(), 1);
        assert_eq!(h1.messages[0].msg_type, MessageType::Dataspace);

        let h2 = ObjectHeader::parse(&buf, offsets[1], 8, 8).unwrap();
        assert_eq!(h2.messages.len(), 1);
        assert_eq!(h2.messages[0].msg_type, MessageType::Datatype);
    }

    #[test]
    fn batch_writer_empty() {
        let batch = BatchObjectHeaderWriter::new();
        assert!(batch.is_empty());
        let (buf, offsets) = batch.serialize_all().unwrap();
        assert!(buf.is_empty());
        assert!(offsets.is_empty());
    }
}