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

//! Fill Value messages (0x0005, and the old 0x0004) and applying them on read.
//!
//! HDF5 allocates storage lazily: a chunk nobody wrote to does not exist in the
//! file, and a contiguous dataset nobody wrote to has no data address at all.
//! Reading such a region must yield the dataset's *fill value* (zeros unless
//! the creator chose otherwise). The readers in [`crate::chunked_read`] leave
//! those regions zeroed; [`apply_to_unallocated_chunks`] then overwrites exactly
//! the chunk-grid cells that are absent from the chunk index — so it can never
//! mistake a stored zero for a hole — and is skipped entirely in the common
//! case of a zero fill value.

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

use crate::addr::to_usize;
use crate::chunked_read::{alloc_output, checked_byte_len, list_chunks_in};
use crate::data_layout::DataLayout;
use crate::dataspace::Dataspace;
use crate::error::FormatError;
use crate::message_type::MessageType;
use crate::object_header::HeaderMessage;

/// Largest fill value accepted. A fill value is one element of the dataset's
/// datatype; this only bounds the allocation driven by the message's size field.
const MAX_FILL_VALUE_SIZE: usize = 1 << 20;

/// Parse a Fill Value message, returning the user-defined fill value bytes, or
/// `None` when the dataset uses the default (all zeros) or has the fill value
/// explicitly undefined.
pub fn parse_fill_value(msg: &HeaderMessage) -> Result<Option<Vec<u8>>, FormatError> {
    let data = msg.data.as_slice();
    let value_at = |pos: usize| -> Result<Option<Vec<u8>>, FormatError> {
        let size_bytes = data.get(pos..pos + 4).ok_or(FormatError::UnexpectedEof {
            expected: pos + 4,
            available: data.len(),
        })?;
        let size = u32::from_le_bytes([size_bytes[0], size_bytes[1], size_bytes[2], size_bytes[3]])
            as usize;
        if size == 0 {
            return Ok(None);
        }
        if size > MAX_FILL_VALUE_SIZE {
            return Err(FormatError::Overflow(format!(
                "fill value of {size} bytes exceeds the {MAX_FILL_VALUE_SIZE}-byte limit"
            )));
        }
        let start = pos + 4;
        let value =
            data.get(start..start.saturating_add(size))
                .ok_or(FormatError::UnexpectedEof {
                    expected: start.saturating_add(size),
                    available: data.len(),
                })?;
        Ok(Some(value.to_vec()))
    };

    match msg.msg_type {
        // Old fill value message: size(4), value.
        MessageType::FillValueOld => value_at(0),
        MessageType::FillValue => {
            let version = *data.first().ok_or(FormatError::UnexpectedEof {
                expected: 1,
                available: 0,
            })?;
            match version {
                // version, alloc time, write time, defined, [size, value]
                1 | 2 => {
                    let defined = *data.get(3).ok_or(FormatError::UnexpectedEof {
                        expected: 4,
                        available: data.len(),
                    })?;
                    if version == 2 && defined == 0 {
                        Ok(None)
                    } else if data.len() < 8 && version == 1 {
                        // v1 always carries a size, but tolerate its absence.
                        Ok(None)
                    } else {
                        value_at(4)
                    }
                }
                // version, flags (bit 4 = undefined, bit 5 = defined), [size, value]
                3 => {
                    let flags = *data.get(1).ok_or(FormatError::UnexpectedEof {
                        expected: 2,
                        available: data.len(),
                    })?;
                    if flags & 0x10 != 0 || flags & 0x20 == 0 {
                        Ok(None)
                    } else {
                        value_at(2)
                    }
                }
                v => Err(FormatError::UnsupportedVersion(v)),
            }
        }
        _ => Ok(None),
    }
}

/// The fill value that applies to a dataset given its header messages. The new
/// message wins over the old one when both are present.
///
/// A *shared* fill value message holds only a reference to the real message,
/// which cannot be followed without the file: this returns
/// [`FormatError::UnresolvedSharedMessage`] for one (it used to answer "zeros").
/// Use [`dataset_fill_value_in`] when the file bytes are at hand.
pub fn dataset_fill_value(messages: &[HeaderMessage]) -> Result<Option<Vec<u8>>, FormatError> {
    fill_value_from(messages, |_| Err(FormatError::UnresolvedSharedMessage))
}

/// [`dataset_fill_value`] for a dataset in `file_data`, following a shared
/// fill value message to where it lives: another object header, or the
/// file's shared-message (SOHM) heap, as libhdf5 writes it when the file has
/// a SOHM index for fill values.
pub fn dataset_fill_value_in(
    file_data: &[u8],
    messages: &[HeaderMessage],
    offset_size: u8,
    length_size: u8,
) -> Result<Option<Vec<u8>>, FormatError> {
    dataset_fill_value_from_storage(&file_data, messages, offset_size, length_size)
}

/// [`dataset_fill_value_in`] with the file behind any
/// [`Storage`](crate::storage::Storage) (a `&dyn Storage` too). (The trait
/// is not imported here: its `len` would shadow the slice method in this
/// module.)
pub fn dataset_fill_value_from_storage<S: crate::storage::Storage + ?Sized>(
    file: &S,
    messages: &[HeaderMessage],
    offset_size: u8,
    length_size: u8,
) -> Result<Option<Vec<u8>>, FormatError> {
    fill_value_from(messages, |msg| {
        crate::shared_message::message_data_with_sohm_in(file, msg, offset_size, length_size)
            .map(|data| data.into_owned())
    })
}

fn fill_value_from(
    messages: &[HeaderMessage],
    resolve_shared: impl Fn(&HeaderMessage) -> Result<Vec<u8>, FormatError>,
) -> Result<Option<Vec<u8>>, FormatError> {
    for wanted in [MessageType::FillValue, MessageType::FillValueOld] {
        if let Some(msg) = messages.iter().find(|m| m.msg_type == wanted) {
            let value = if crate::shared_message::is_shared(msg.flags) {
                let data = resolve_shared(msg)?;
                parse_fill_value(&HeaderMessage {
                    msg_type: msg.msg_type,
                    size: data.len(),
                    flags: msg.flags & !0x02,
                    creation_order: msg.creation_order,
                    data,
                })?
            } else {
                parse_fill_value(msg)?
            };
            if let Some(value) = value {
                return Ok(Some(value));
            }
        }
    }
    Ok(None)
}

/// `true` when a fill value is absent or all zeros, i.e. identical to what the
/// readers already produce for unallocated storage.
pub fn is_default(fill: Option<&[u8]>) -> bool {
    fill.is_none_or(|f| f.iter().all(|&b| b == 0))
}

/// A whole dataset's worth of fill value: what reading a dataset with no
/// allocated storage at all must return.
pub fn filled_dataset(
    dataspace: &Dataspace,
    elem_size: usize,
    fill: Option<&[u8]>,
) -> Result<Vec<u8>, FormatError> {
    let total = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?;
    let mut out = alloc_output(total)?;
    if let Some(fill) = fill.filter(|f| f.len() == elem_size && !is_default(Some(f))) {
        for element in out.chunks_exact_mut(elem_size) {
            element.copy_from_slice(fill);
        }
    }
    Ok(out)
}

/// Whether the layout has any storage in the file at all. A dataset that was
/// created but never written to has none.
pub fn has_storage(layout: &DataLayout) -> bool {
    !matches!(
        layout,
        DataLayout::Contiguous { address: None, .. }
            | DataLayout::Chunked {
                btree_address: None,
                ..
            }
    )
}

/// Run a full-dataset `read`, giving unallocated storage its fill value: a
/// dataset with no storage at all reads as entirely fill value (instead of
/// failing), and a chunked dataset has the fill value written into every
/// chunk the file never allocated.
#[allow(clippy::too_many_arguments)]
pub fn read_full_with_fill<E: From<FormatError>>(
    messages: &[HeaderMessage],
    file_data: &[u8],
    layout: &DataLayout,
    dataspace: &Dataspace,
    elem_size: usize,
    offset_size: u8,
    length_size: u8,
    read: impl FnOnce() -> Result<Vec<u8>, E>,
) -> Result<Vec<u8>, E> {
    read_full_with_fill_in(
        messages,
        file_data,
        layout,
        dataspace,
        elem_size,
        offset_size,
        length_size,
        read,
    )
}

/// [`read_full_with_fill`] over any [`Storage`](crate::storage::Storage).
#[allow(clippy::too_many_arguments)]
pub fn read_full_with_fill_in<E: From<FormatError>, S: crate::storage::Storage + ?Sized>(
    messages: &[HeaderMessage],
    file_data: &S,
    layout: &DataLayout,
    dataspace: &Dataspace,
    elem_size: usize,
    offset_size: u8,
    length_size: u8,
    read: impl FnOnce() -> Result<Vec<u8>, E>,
) -> Result<Vec<u8>, E> {
    // A dataset with external raw data also has no data address in this
    // file. It is NOT unallocated — its values live elsewhere — so it must
    // never be answered with the fill value.
    if messages
        .iter()
        .any(|m| m.msg_type == MessageType::ExternalDataFiles)
    {
        return Err(FormatError::ExternalDataFilesUnsupported.into());
    }
    let fill = dataset_fill_value_from_storage(file_data, messages, offset_size, length_size)?;
    if !has_storage(layout) {
        return Ok(filled_dataset(dataspace, elem_size, fill.as_deref())?);
    }
    let mut output = read()?;
    apply_to_unallocated_chunks_in(
        &mut output,
        file_data,
        layout,
        dataspace,
        elem_size,
        fill.as_deref(),
        offset_size,
        length_size,
    )?;
    Ok(output)
}

/// Overwrite, in a fully read chunked dataset `output`, every region whose
/// chunk was never allocated with `fill`. No-op for non-chunked layouts, a
/// default fill value, or a fill value whose size doesn't match the element.
#[allow(clippy::too_many_arguments)]
pub fn apply_to_unallocated_chunks(
    output: &mut [u8],
    file_data: &[u8],
    layout: &DataLayout,
    dataspace: &Dataspace,
    elem_size: usize,
    fill: Option<&[u8]>,
    offset_size: u8,
    length_size: u8,
) -> Result<(), FormatError> {
    apply_to_unallocated_chunks_in(
        output,
        file_data,
        layout,
        dataspace,
        elem_size,
        fill,
        offset_size,
        length_size,
    )
}

/// [`apply_to_unallocated_chunks`] over any [`Storage`](crate::storage::Storage).
#[allow(clippy::too_many_arguments)]
pub fn apply_to_unallocated_chunks_in<S: crate::storage::Storage + ?Sized>(
    output: &mut [u8],
    file_data: &S,
    layout: &DataLayout,
    dataspace: &Dataspace,
    elem_size: usize,
    fill: Option<&[u8]>,
    offset_size: u8,
    length_size: u8,
) -> Result<(), FormatError> {
    let Some(fill) = fill.filter(|f| f.len() == elem_size && !is_default(Some(f))) else {
        return Ok(());
    };
    if !matches!(layout, DataLayout::Chunked { .. }) || elem_size == 0 {
        return Ok(());
    }
    let (chunks, chunk_dims) = list_chunks_in(
        file_data,
        layout,
        dataspace,
        elem_size,
        offset_size,
        length_size,
    )?;
    let rank = chunk_dims.len();
    let ds_dims: Vec<usize> = dataspace
        .dimensions
        .iter()
        .map(|&d| to_usize(d))
        .collect::<Result<_, _>>()?;
    if rank == 0 || ds_dims.len() != rank || chunk_dims.contains(&0) {
        return Ok(());
    }

    // Row-major strides over the dataset and over the chunk grid.
    let mut ds_strides = vec![1usize; rank];
    for i in (0..rank - 1).rev() {
        ds_strides[i] = ds_strides[i + 1].saturating_mul(ds_dims[i + 1]);
    }
    let grid: Vec<usize> = ds_dims
        .iter()
        .zip(&chunk_dims)
        .map(|(&d, &c)| d.div_ceil(c))
        .collect();
    let cells = grid
        .iter()
        .try_fold(1usize, |acc, &g| acc.checked_mul(g))
        .ok_or_else(|| FormatError::Overflow("chunk grid size overflows".into()))?;
    if cells == 0 {
        return Ok(());
    }

    let mut allocated = vec![false; cells];
    for chunk in &chunks {
        // Undefined address: the index has a slot for the chunk but no storage.
        if chunk.address == u64::MAX || chunk.offsets.len() < rank {
            continue;
        }
        let mut cell = 0usize;
        let mut in_range = true;
        for d in 0..rank {
            let coord = to_usize(chunk.offsets[d])? / chunk_dims[d];
            if coord >= grid[d] {
                in_range = false;
                break;
            }
            cell = cell * grid[d] + coord;
        }
        if in_range {
            allocated[cell] = true;
        }
    }

    let mut coord = vec![0usize; rank];
    for (cell, is_allocated) in allocated.iter().enumerate() {
        if *is_allocated {
            continue;
        }
        // Decode the cell index into grid coordinates.
        let mut rem = cell;
        for d in (0..rank).rev() {
            coord[d] = rem % grid[d];
            rem /= grid[d];
        }
        fill_cell(
            output,
            &coord,
            &chunk_dims,
            &ds_dims,
            &ds_strides,
            elem_size,
            fill,
        );
    }
    Ok(())
}

/// Fill the part of chunk-grid cell `coord` that lies inside the dataset.
fn fill_cell(
    output: &mut [u8],
    coord: &[usize],
    chunk_dims: &[usize],
    ds_dims: &[usize],
    ds_strides: &[usize],
    elem_size: usize,
    fill: &[u8],
) {
    let rank = coord.len();
    let start: Vec<usize> = (0..rank).map(|d| coord[d] * chunk_dims[d]).collect();
    let end: Vec<usize> = (0..rank)
        .map(|d| (start[d] + chunk_dims[d]).min(ds_dims[d]))
        .collect();
    if (0..rank).any(|d| start[d] >= end[d]) {
        return;
    }
    // Walk every row (all dims but the last) and fill the run along the last.
    let run = end[rank - 1] - start[rank - 1];
    let mut idx = start.clone();
    loop {
        let first: usize = (0..rank).map(|d| idx[d] * ds_strides[d]).sum();
        let from = first * elem_size;
        let to = from + run * elem_size;
        if let Some(region) = output.get_mut(from..to) {
            for element in region.chunks_exact_mut(elem_size) {
                element.copy_from_slice(fill);
            }
        }
        // Advance the odometer over dims 0..rank-1.
        let mut d = rank - 1;
        loop {
            if d == 0 {
                return;
            }
            d -= 1;
            idx[d] += 1;
            if idx[d] < end[d] {
                break;
            }
            idx[d] = start[d];
        }
    }
}

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

    fn msg(msg_type: MessageType, data: &[u8]) -> HeaderMessage {
        HeaderMessage {
            msg_type,
            size: data.len(),
            flags: 0,
            creation_order: None,
            data: data.to_vec(),
        }
    }

    #[test]
    fn parses_v3_defined_undefined_and_default() {
        // Real message for h5py `fillvalue=-1` on an i4 dataset (HDF5 2.0).
        let defined = msg(
            MessageType::FillValue,
            &[3, 0x2b, 4, 0, 0, 0, 0xff, 0xff, 0xff, 0xff],
        );
        assert_eq!(parse_fill_value(&defined).unwrap(), Some(vec![0xff; 4]));
        let default = msg(MessageType::FillValue, &[3, 0x0a]);
        assert_eq!(parse_fill_value(&default).unwrap(), None);
        let undefined = msg(MessageType::FillValue, &[3, 0x19]);
        assert_eq!(parse_fill_value(&undefined).unwrap(), None);
    }

    #[test]
    fn parses_v2_and_old_messages() {
        let v2 = msg(MessageType::FillValue, &[2, 2, 2, 1, 2, 0, 0, 0, 7, 0]);
        assert_eq!(parse_fill_value(&v2).unwrap(), Some(vec![7, 0]));
        let v2_undefined = msg(MessageType::FillValue, &[2, 2, 2, 0]);
        assert_eq!(parse_fill_value(&v2_undefined).unwrap(), None);
        let old = msg(MessageType::FillValueOld, &[2, 0, 0, 0, 9, 9]);
        assert_eq!(parse_fill_value(&old).unwrap(), Some(vec![9, 9]));
    }

    #[test]
    fn truncated_or_oversized_fill_is_an_error() {
        let short = msg(MessageType::FillValue, &[3, 0x29, 4, 0, 0, 0, 0xff]);
        assert!(parse_fill_value(&short).is_err());
        let huge = msg(MessageType::FillValue, &[3, 0x29, 0xff, 0xff, 0xff, 0x7f]);
        assert!(matches!(
            parse_fill_value(&huge),
            Err(FormatError::Overflow(_))
        ));
    }

    #[test]
    fn fill_cell_clips_edge_chunks_in_2d() {
        // 3x5 dataset, 2x2 chunks; fill grid cell (1, 2): rows 2..3, cols 4..5.
        let mut out = vec![0u8; 15];
        fill_cell(&mut out, &[1, 2], &[2, 2], &[3, 5], &[5, 1], 1, &[9]);
        let mut expected = vec![0u8; 15];
        expected[2 * 5 + 4] = 9;
        assert_eq!(out, expected);

        // Interior cell (0, 1): rows 0..2, cols 2..4.
        let mut out = vec![0u8; 15];
        fill_cell(&mut out, &[0, 1], &[2, 2], &[3, 5], &[5, 1], 1, &[7]);
        let filled: Vec<usize> = out
            .iter()
            .enumerate()
            .filter(|(_, b)| **b == 7)
            .map(|(i, _)| i)
            .collect();
        assert_eq!(filled, [2, 3, 7, 8]);
    }

    /// Fill values, shared ones in the SOHM heap included, resolve
    /// identically through a read_at-only CountingStorage.
    #[test]
    fn storage_reads_match_slice_reads() {
        use crate::object_header::ObjectHeader;
        use crate::storage::CountingStorage;
        let file: &[u8] = include_bytes!("../tests/fixtures/shared_fill_value.h5");
        let sb = crate::superblock::Superblock::parse(file, 0).unwrap();
        let (os, ls) = (sb.offset_size, sb.length_size);
        let storage = CountingStorage::new(file.to_vec());
        let mut shared = 0;
        let children =
            crate::group_v2::resolve_group_children(file, &sb, sb.root_group_address).unwrap();
        assert!(children.len() >= 3);
        for child in children {
            let h =
                ObjectHeader::parse(file, child.object_header_address as usize, os, ls).unwrap();
            shared += h
                .messages
                .iter()
                .filter(|m| {
                    m.msg_type == MessageType::FillValue
                        && crate::shared_message::is_shared(m.flags)
                })
                .count();
            let want = dataset_fill_value_in(file, &h.messages, os, ls);
            assert_eq!(want, Ok(Some((-7i32).to_le_bytes().to_vec())));
            let got = dataset_fill_value_from_storage(&storage, &h.messages, os, ls);
            assert_eq!(got, want, "{}", child.name);
        }
        assert!(shared >= 2);
    }
}