rust-hdf5 0.7.2

//! Local heap decode (for reading legacy HDF5 files).
//!
//! The local heap is used by v0/v1 groups to store link names as
//! null-terminated strings. The heap header lives at a known address and
//! points to a contiguous data block.
//!
//! Header layout:
//! ```text
//! "HEAP" (4 bytes)
//! version: 1 byte (0)
//! reserved: 3 bytes
//! data_size: sizeof_size bytes LE
//! free_list_offset: sizeof_size bytes LE (0xFFFFFFFFFFFFFFFF = none)
//! data_addr: sizeof_addr bytes LE
//! ```

use crate::format::bytes::read_le_uint as read_uint;
use crate::format::{FormatError, FormatResult};

/// The 4-byte local heap signature.
pub const LOCAL_HEAP_SIGNATURE: [u8; 4] = *b"HEAP";

/// Decoded local heap header.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LocalHeapHeader {
    /// Total size of the data segment.
    pub data_size: u64,
    /// Offset into the data segment of the first free block, or u64::MAX if none.
    pub free_list_offset: u64,
    /// File address of the data segment.
    pub data_addr: u64,
}

impl LocalHeapHeader {
    /// Decode a local heap header from `buf`.
    ///
    /// `sizeof_addr` and `sizeof_size` come from the superblock.
    pub fn decode(buf: &[u8], sizeof_addr: usize, sizeof_size: usize) -> FormatResult<Self> {
        let min_size = 4 + 1 + 3 + sizeof_size * 2 + sizeof_addr;
        if buf.len() < min_size {
            return Err(FormatError::BufferTooShort {
                needed: min_size,
                available: buf.len(),
            });
        }

        if buf[0..4] != LOCAL_HEAP_SIGNATURE {
            return Err(FormatError::InvalidSignature);
        }

        let version = buf[4];
        if version != 0 {
            return Err(FormatError::InvalidVersion(version));
        }

        // buf[5..8] reserved
        let mut pos = 8;

        let data_size = read_uint(&buf[pos..], sizeof_size);
        pos += sizeof_size;

        let free_list_offset = read_uint(&buf[pos..], sizeof_size);
        pos += sizeof_size;

        let data_addr = read_uint(&buf[pos..], sizeof_addr);

        Ok(LocalHeapHeader {
            data_size,
            free_list_offset,
            data_addr,
        })
    }
}

/// `H5HL_FREE_NULL` (H5HLpkg.h:54): the free-list head value that means "no
/// free block", written in the header's free-list field. It is 1, not an
/// undefined address, because 1 can never be a legal 8-aligned block offset.
pub const LOCAL_HEAP_FREE_NULL: u64 = 1;

/// `H5HL_ALIGN` (H5HLprivate.h:29): every heap object, and the data block
/// itself, is a multiple of eight bytes long.
pub fn local_heap_align(n: u64) -> u64 {
    (n + 7) & !7
}

/// `H5HL_SIZEOF_HDR`: the heap header's own on-disk length, aligned.
pub fn local_heap_header_size(sizeof_addr: usize, sizeof_size: usize) -> usize {
    local_heap_align((4 + 1 + 3 + 2 * sizeof_size + sizeof_addr) as u64) as usize
}

impl LocalHeapHeader {
    /// Encode the heap header (`H5HL__cache_prefix_serialize`).
    ///
    /// A `free_list_offset` of [`LOCAL_HEAP_FREE_NULL`] is the "no free space"
    /// value; the field is not an address, so it has no undefined form.
    pub fn encode(&self, sizeof_addr: usize, sizeof_size: usize) -> Vec<u8> {
        let size = local_heap_header_size(sizeof_addr, sizeof_size);
        let mut buf = Vec::with_capacity(size);
        buf.extend_from_slice(&LOCAL_HEAP_SIGNATURE);
        buf.push(0); // version
        buf.extend_from_slice(&[0u8; 3]); // reserved
        buf.extend_from_slice(&self.data_size.to_le_bytes()[..sizeof_size]);
        buf.extend_from_slice(&self.free_list_offset.to_le_bytes()[..sizeof_size]);
        buf.extend_from_slice(&self.data_addr.to_le_bytes()[..sizeof_addr]);
        buf.resize(size, 0); // the alignment tail of H5HL_SIZEOF_HDR
        buf
    }
}

/// A local heap's data segment, built up object by object.
///
/// This is the bulk-load half of `H5HL_insert`: every object is placed at the
/// end of the block and the block grows to fit, which is what that function
/// does once its free list is empty. The free list itself is not modelled —
/// the writer rebuilds a group's heap whole rather than patching the one on
/// disk, so there is never a hole to record.
#[derive(Debug, Default)]
pub struct LocalHeapImage {
    data: Vec<u8>,
}

impl LocalHeapImage {
    /// A heap holding just the empty string at offset 0.
    ///
    /// Not an optimisation: `H5G__stab_create_components` asserts that offset,
    /// because a symbol-table B-tree's leftmost key is the empty string and
    /// every comparison against it reads heap offset 0.
    pub fn with_empty_string() -> Self {
        let mut image = Self::default();
        let offset = image.insert(b"\0");
        debug_assert_eq!(offset, 0);
        image
    }

    /// Append `bytes` (null terminator included) and return its offset.
    pub fn insert(&mut self, bytes: &[u8]) -> u64 {
        let offset = self.data.len() as u64;
        self.data.extend_from_slice(bytes);
        let padded = local_heap_align(self.data.len() as u64) as usize;
        self.data.resize(padded, 0);
        offset
    }

    /// Append a name as a null-terminated string, returning its offset.
    pub fn insert_str(&mut self, s: &str) -> u64 {
        let offset = self.data.len() as u64;
        self.data.extend_from_slice(s.as_bytes());
        self.data.push(0);
        let padded = local_heap_align(self.data.len() as u64) as usize;
        self.data.resize(padded, 0);
        offset
    }

    /// The data segment's bytes, whose length is the heap's `data_size`.
    pub fn as_bytes(&self) -> &[u8] {
        &self.data
    }
}

/// Look up a null-terminated string in the heap data block by offset.
///
/// `heap_data` is the raw bytes of the local heap data segment.
/// `offset` is the byte offset into that segment where the string starts.
pub fn local_heap_get_string(heap_data: &[u8], offset: u64) -> FormatResult<String> {
    let start = offset as usize;
    if start >= heap_data.len() {
        return Err(FormatError::InvalidData(format!(
            "local heap offset {} out of range (heap size {})",
            offset,
            heap_data.len()
        )));
    }

    // Find the null terminator
    let end = heap_data[start..]
        .iter()
        .position(|&b| b == 0)
        .map(|p| start + p)
        .unwrap_or(heap_data.len());

    String::from_utf8(heap_data[start..end].to_vec())
        .map_err(|e| FormatError::InvalidData(format!("invalid UTF-8 in local heap string: {}", e)))
}

// ======================================================================= tests

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

    fn build_heap_header(
        data_size: u64,
        free_list_offset: u64,
        data_addr: u64,
        sizeof_addr: usize,
        sizeof_size: usize,
    ) -> Vec<u8> {
        let mut buf = Vec::new();
        buf.extend_from_slice(&LOCAL_HEAP_SIGNATURE);
        buf.push(0); // version
        buf.extend_from_slice(&[0u8; 3]); // reserved
        buf.extend_from_slice(&data_size.to_le_bytes()[..sizeof_size]);
        buf.extend_from_slice(&free_list_offset.to_le_bytes()[..sizeof_size]);
        buf.extend_from_slice(&data_addr.to_le_bytes()[..sizeof_addr]);
        buf
    }

    #[test]
    fn decode_basic() {
        let buf = build_heap_header(128, u64::MAX, 0x1000, 8, 8);
        let hdr = LocalHeapHeader::decode(&buf, 8, 8).unwrap();
        assert_eq!(hdr.data_size, 128);
        assert_eq!(hdr.free_list_offset, u64::MAX);
        assert_eq!(hdr.data_addr, 0x1000);
    }

    #[test]
    fn decode_4byte() {
        let buf = build_heap_header(64, 0xFFFFFFFF, 0x800, 4, 4);
        let hdr = LocalHeapHeader::decode(&buf, 4, 4).unwrap();
        assert_eq!(hdr.data_size, 64);
        assert_eq!(hdr.free_list_offset, 0xFFFFFFFF);
        assert_eq!(hdr.data_addr, 0x800);
    }

    #[test]
    fn decode_bad_sig() {
        let mut buf = build_heap_header(64, 0, 0x800, 8, 8);
        buf[0] = b'X';
        assert!(matches!(
            LocalHeapHeader::decode(&buf, 8, 8).unwrap_err(),
            FormatError::InvalidSignature
        ));
    }

    #[test]
    fn decode_bad_version() {
        let mut buf = build_heap_header(64, 0, 0x800, 8, 8);
        buf[4] = 1;
        assert!(matches!(
            LocalHeapHeader::decode(&buf, 8, 8).unwrap_err(),
            FormatError::InvalidVersion(1)
        ));
    }

    #[test]
    fn decode_too_short() {
        let buf = [0u8; 4];
        assert!(matches!(
            LocalHeapHeader::decode(&buf, 8, 8).unwrap_err(),
            FormatError::BufferTooShort { .. }
        ));
    }

    #[test]
    fn get_string_basic() {
        let mut data = Vec::new();
        data.extend_from_slice(b"\0"); // offset 0: empty
        data.extend_from_slice(b"hello\0");
        data.extend_from_slice(b"world\0");

        assert_eq!(local_heap_get_string(&data, 0).unwrap(), "");
        assert_eq!(local_heap_get_string(&data, 1).unwrap(), "hello");
        assert_eq!(local_heap_get_string(&data, 7).unwrap(), "world");
    }

    #[test]
    fn get_string_out_of_range() {
        let data = b"hello\0";
        assert!(matches!(
            local_heap_get_string(data, 100).unwrap_err(),
            FormatError::InvalidData(_)
        ));
    }

    /// The 32-byte header of the root group's heap in a file h5py wrote with
    /// no `libver` argument (three datasets `alpha`/`beta`/`gamma`).
    #[test]
    fn a_local_heap_header_matches_the_bytes_libhdf5_wrote() {
        let hdr = LocalHeapHeader {
            data_size: 88,
            free_list_offset: 0x20,
            data_addr: 0x2c8,
        };
        let expected = [
            b'H', b'E', b'A', b'P', 0, 0, 0, 0, // signature, version 0, reserved
            0x58, 0, 0, 0, 0, 0, 0, 0, // data_size
            0x20, 0, 0, 0, 0, 0, 0, 0, // free list head
            0xc8, 0x02, 0, 0, 0, 0, 0, 0, // data segment address
        ];
        assert_eq!(hdr.encode(8, 8), expected);
        assert_eq!(local_heap_header_size(8, 8), 32);
        assert_eq!(LocalHeapHeader::decode(&expected, 8, 8).unwrap(), hdr);
    }

    /// The header is 8-aligned, so the 4/4 form pads rather than shrinking to
    /// its 20 significant bytes.
    #[test]
    fn a_four_byte_local_heap_header_pads_to_its_alignment() {
        assert_eq!(local_heap_header_size(4, 4), 24);
        let hdr = LocalHeapHeader {
            data_size: 64,
            free_list_offset: LOCAL_HEAP_FREE_NULL,
            data_addr: 0x800,
        };
        let buf = hdr.encode(4, 4);
        assert_eq!(buf.len(), 24);
        assert_eq!(&buf[20..], &[0, 0, 0, 0]);
        assert_eq!(LocalHeapHeader::decode(&buf, 4, 4).unwrap(), hdr);
    }

    /// The same three names, laid into a data segment: the empty string at
    /// offset 0 and every object padded to eight.
    #[test]
    fn a_heap_image_lays_names_out_the_way_libhdf5_does() {
        let mut heap = LocalHeapImage::with_empty_string();
        assert_eq!(heap.insert_str("alpha"), 8);
        assert_eq!(heap.insert_str("beta"), 16);
        assert_eq!(heap.insert_str("gamma"), 24);
        assert_eq!(
            heap.as_bytes(),
            b"\0\0\0\0\0\0\0\0alpha\0\0\0beta\0\0\0\0gamma\0\0\0"
        );
        for (offset, name) in [(0, ""), (8, "alpha"), (16, "beta"), (24, "gamma")] {
            assert_eq!(
                local_heap_get_string(heap.as_bytes(), offset).unwrap(),
                name
            );
        }
    }

    /// A name whose length is already a multiple of eight still gets its
    /// terminator, so it costs a whole extra eight bytes.
    #[test]
    fn a_heap_object_is_padded_after_its_terminator_not_before() {
        let mut heap = LocalHeapImage::with_empty_string();
        assert_eq!(heap.insert_str("12345678"), 8);
        assert_eq!(heap.insert_str("next"), 24);
        assert_eq!(heap.as_bytes().len(), 32);
        assert_eq!(
            local_heap_get_string(heap.as_bytes(), 8).unwrap(),
            "12345678"
        );
    }
}