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

//! Chunk-index linearisation shared by the Fixed Array and Extensible Array
//! chunk indexes (reader and writer).
//!
//! Both indexes store one element per chunk at a *linear* index, and the
//! library derives that index from the chunk's scaled coordinates
//! (`offset / chunk_dim`) using the dataset's **maximum** dimensions, not its
//! current ones (`H5D__farray_idx_get_addr` / `H5D__earray_idx_get_addr`,
//! via `layout->max_down_chunks`). A dataset whose current shape is smaller
//! than its maxshape therefore has gaps in the index, and laying it out by the
//! current shape puts every chunk after the first row in the wrong place.
//!
//! The Extensible Array adds one more step: its one unlimited dimension has no
//! finite chunk count, so the library *swizzles* the coordinates to make that
//! dimension the slowest-varying one (`H5VM_swizzle_coords`, which moves
//! `coords[unlim_dim]` to the front and shifts the dimensions before it right
//! by one) before linearising with `swizzled_max_down_chunks`. When the
//! unlimited dimension is already dimension 0 no swizzle happens.

#[cfg(not(feature = "std"))]
extern crate alloc;

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

use crate::error::FormatError;

/// How a chunk index maps linear element indexes to chunk coordinates.
#[derive(Debug, Clone)]
pub(crate) struct ChunkGrid {
    /// Spatial chunk dimensions, in dataset order.
    chunk_dims: Vec<u64>,
    /// Chunks per dimension covering the *current* extent, in dataset order.
    cur_chunks: Vec<u64>,
    /// Dataset dimension stored at each linearisation position (slowest
    /// first). The identity except for a swizzled Extensible Array.
    order: Vec<usize>,
    /// Linear stride of each linearisation position.
    down: Vec<u64>,
}

impl ChunkGrid {
    /// Grid for a Fixed Array index: row-major over the chunk counts of the
    /// maximum dimensions (`max_dims`, falling back to the current dimensions
    /// when the dataspace records none).
    pub(crate) fn fixed_array(
        cur_dims: &[u64],
        max_dims: Option<&[u64]>,
        chunk_dims: &[u64],
    ) -> Result<Self, FormatError> {
        Self::build(cur_dims, max_dims, chunk_dims, None)
    }

    /// Grid for an Extensible Array index: like the Fixed Array, but the
    /// unlimited dimension (the one whose maximum is `H5S_UNLIMITED`) is moved
    /// to the slowest-varying position first.
    pub(crate) fn extensible_array(
        cur_dims: &[u64],
        max_dims: Option<&[u64]>,
        chunk_dims: &[u64],
    ) -> Result<Self, FormatError> {
        let unlim = max_dims.and_then(|m| m.iter().position(|&d| d == u64::MAX));
        Self::build(cur_dims, max_dims, chunk_dims, unlim)
    }

    fn build(
        cur_dims: &[u64],
        max_dims: Option<&[u64]>,
        chunk_dims: &[u64],
        unlim: Option<usize>,
    ) -> Result<Self, FormatError> {
        let rank = chunk_dims.len();
        if cur_dims.len() != rank || max_dims.is_some_and(|m| m.len() != rank) {
            return Err(FormatError::ChunkedReadError(
                "chunk index rank does not match the dataspace".into(),
            ));
        }
        if chunk_dims.contains(&0) {
            return Err(FormatError::ChunkedReadError(
                "chunk dimension is zero".into(),
            ));
        }
        let cur_chunks: Vec<u64> = cur_dims
            .iter()
            .zip(chunk_dims)
            .map(|(&d, &c)| d.div_ceil(c))
            .collect();
        // Chunk counts of the maximum extent. An unlimited dimension has no
        // finite count; it only ever sits in the slowest position, where its
        // count never enters a stride. A (corrupt) maximum smaller than the
        // current extent is widened so no allocated chunk becomes unreachable.
        let max_chunks: Vec<u64> = (0..rank)
            .map(|d| {
                let max = max_dims.map_or(cur_dims[d], |m| m[d]);
                if max == u64::MAX {
                    u64::MAX
                } else {
                    max.div_ceil(chunk_dims[d]).max(cur_chunks[d])
                }
            })
            .collect();

        let mut order: Vec<usize> = (0..rank).collect();
        if let Some(u) = unlim {
            order.remove(u);
            order.insert(0, u);
        }
        let mut down = vec![1u64; rank];
        for p in (0..rank.saturating_sub(1)).rev() {
            let next = max_chunks[order[p + 1]];
            if next == u64::MAX {
                // Only reachable with more than one unlimited dimension, which
                // neither index type can describe.
                return Err(FormatError::ChunkedReadError(
                    "array chunk index with more than one unlimited dimension".into(),
                ));
            }
            down[p] = down[p + 1].checked_mul(next).ok_or_else(|| {
                FormatError::Overflow("chunk index linear stride overflows u64".into())
            })?;
        }
        Ok(Self {
            chunk_dims: chunk_dims.to_vec(),
            cur_chunks,
            order,
            down,
        })
    }

    /// Dataset-space offsets of the chunk stored at linear `index`, or `None`
    /// when that chunk lies outside the current extent (the index still has a
    /// slot for it; the library ignores such chunks on read).
    pub(crate) fn offsets(&self, index: u64) -> Option<Vec<u64>> {
        let rank = self.chunk_dims.len();
        let mut offsets = vec![0u64; rank];
        let mut rem = index;
        for p in 0..rank {
            let d = self.order[p];
            // A zero stride: a later dimension has no chunks (its maximum,
            // or with none recorded its current extent, is 0), so no slot of
            // the index is a chunk of the dataset.
            if self.down[p] == 0 {
                return None;
            }
            let scaled = rem / self.down[p];
            rem %= self.down[p];
            if scaled >= self.cur_chunks[d] {
                return None;
            }
            offsets[d] = scaled * self.chunk_dims[d];
        }
        Some(offsets)
    }

    /// Linear index of the chunk with scaled coordinates `scaled`
    /// (`offset / chunk_dim` per dimension, in dataset order).
    pub(crate) fn linear_index(&self, scaled: &[u64]) -> u64 {
        self.order
            .iter()
            .zip(&self.down)
            .map(|(&d, &stride)| scaled[d] * stride)
            .sum()
    }
}

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

    #[test]
    fn fixed_array_uses_max_dims() {
        // shape (4, 6), chunks (2, 3), maxshape (20, 10): 10 x 4 chunk grid.
        let g = ChunkGrid::fixed_array(&[4, 6], Some(&[20, 10]), &[2, 3]).unwrap();
        assert_eq!(g.offsets(0), Some(vec![0, 0]));
        assert_eq!(g.offsets(1), Some(vec![0, 3]));
        assert_eq!(g.offsets(2), None); // column chunk 2 is beyond the extent
        assert_eq!(g.offsets(4), Some(vec![2, 0]));
        assert_eq!(g.offsets(5), Some(vec![2, 3]));
        assert_eq!(g.offsets(8), None); // row chunk 2 is beyond the extent
        assert_eq!(g.linear_index(&[1, 1]), 5);
    }

    #[test]
    fn extensible_array_swizzles_unlimited_dim() {
        // maxshape (10, None): dim 1 is unlimited and becomes slowest.
        let g = ChunkGrid::extensible_array(&[4, 6], Some(&[10, u64::MAX]), &[2, 3]).unwrap();
        // max chunks of dim 0 = 5, so index = c1 * 5 + c0.
        assert_eq!(g.linear_index(&[1, 0]), 1);
        assert_eq!(g.linear_index(&[0, 1]), 5);
        assert_eq!(g.offsets(5), Some(vec![0, 3]));
        assert_eq!(g.offsets(6), Some(vec![2, 3]));
        assert_eq!(g.offsets(2), None);
    }

    #[test]
    fn extensible_array_unlimited_first_is_row_major() {
        let g = ChunkGrid::extensible_array(&[4, 6], Some(&[u64::MAX, 30]), &[2, 3]).unwrap();
        // max chunks of dim 1 = 10.
        assert_eq!(g.linear_index(&[1, 1]), 11);
        assert_eq!(g.offsets(11), Some(vec![2, 3]));
    }

    #[test]
    fn zero_extent_has_no_chunks() {
        // No maximum recorded and a zero current dimension: every stride
        // before it is 0 (this divided by zero).
        let g = ChunkGrid::fixed_array(&[1, 0], None, &[6, 6]).unwrap();
        for i in 0..16 {
            assert_eq!(g.offsets(i), None);
        }
        let g = ChunkGrid::fixed_array(&[0, 0, 3], Some(&[4, 0, 3]), &[2, 2, 3]).unwrap();
        for i in 0..16 {
            assert_eq!(g.offsets(i), None);
        }
        let g = ChunkGrid::extensible_array(&[0, 5], Some(&[u64::MAX, 0]), &[2, 2]).unwrap();
        for i in 0..16 {
            assert_eq!(g.offsets(i), None);
        }
        // A zero last dimension leaves the other strides alone.
        let g = ChunkGrid::fixed_array(&[4, 0], Some(&[4, 6]), &[2, 3]).unwrap();
        assert_eq!(g.offsets(0), None);
        assert_eq!(g.linear_index(&[1, 1]), 3);
    }

    #[test]
    fn rejects_two_unlimited_dims_after_the_first() {
        assert!(ChunkGrid::fixed_array(&[4, 6], Some(&[u64::MAX, u64::MAX]), &[2, 3]).is_err());
    }
}