clawhdf5/crates/clawhdf5/tests/partial_read_equivalence.rs

//! Selection reads must return exactly what a full read followed by element
//! extraction returns — for every layout, rank and selection shape — while
//! touching only what the selection needs.

use clawhdf5::{File, FileBuilder};
use clawhdf5_format::selection::Selection;

struct Rng(u64);
impl Rng {
    fn next(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^ (z >> 31)
    }
    fn below(&mut self, n: u64) -> u64 {
        self.next() % n.max(1)
    }
}

/// Row-major reference extraction from a full read.
fn reference(full: &[i32], dims: &[u64], selection: &Selection) -> Vec<i32> {
    let strides: Vec<u64> = (0..dims.len())
        .map(|d| dims[d + 1..].iter().product())
        .collect();
    let at =
        |coord: &[u64]| full[coord.iter().zip(&strides).map(|(c, s)| c * s).sum::<u64>() as usize];
    match selection {
        Selection::Points(points) => points.iter().map(|p| at(p)).collect(),
        Selection::Hyperslab {
            start,
            stride,
            count,
            block,
        } => {
            // Selected indices per dimension, then their cartesian product.
            let per_dim: Vec<Vec<u64>> = (0..dims.len())
                .map(|d| {
                    (0..count[d])
                        .flat_map(|c| (0..block[d]).map(move |b| (c, b)))
                        .map(|(c, b)| start[d] + c * stride[d] + b)
                        .collect()
                })
                .collect();
            let mut out = Vec::new();
            let mut idx = vec![0usize; dims.len()];
            loop {
                let coord: Vec<u64> = idx
                    .iter()
                    .enumerate()
                    .map(|(d, &i)| per_dim[d][i])
                    .collect();
                out.push(at(&coord));
                let mut d = dims.len();
                loop {
                    if d == 0 {
                        return out;
                    }
                    d -= 1;
                    idx[d] += 1;
                    if idx[d] < per_dim[d].len() {
                        break;
                    }
                    idx[d] = 0;
                }
            }
        }
        _ => unreachable!(),
    }
}

fn random_hyperslab(rng: &mut Rng, dims: &[u64]) -> Selection {
    let mut start = Vec::new();
    let mut stride = Vec::new();
    let mut count = Vec::new();
    let mut block = Vec::new();
    for &dim in dims {
        let b = 1 + rng.below(3);
        let st = b + rng.below(4); // stride >= block: no overlap
        let s = rng.below(dim - b + 1);
        let max_count = (dim - s - b) / st + 1;
        let c = 1 + rng.below(max_count.min(6));
        start.push(s);
        stride.push(st);
        count.push(c);
        block.push(b);
    }
    Selection::Hyperslab {
        start,
        stride,
        count,
        block,
    }
}

#[test]
fn selection_reads_match_full_reads_for_every_layout() {
    let dir = tempfile::tempdir().unwrap();
    let mut rng = Rng(7);
    // (dims, chunk dims)
    let shapes: [(&[u64], &[u64]); 3] = [
        (&[97], &[10]),
        (&[41, 53], &[8, 9]),
        (&[11, 13, 17], &[4, 5, 6]),
    ];
    for (dims, chunks) in shapes {
        let n: u64 = dims.iter().product();
        let data: Vec<i32> = (0..n as i32).map(|v| v * 3 - 7).collect();

        let path = dir.path().join(format!("r{}.h5", dims.len()));
        let mut builder = FileBuilder::new();
        builder
            .create_dataset("contiguous")
            .with_i32_data(&data)
            .with_shape(dims);
        builder
            .create_dataset("chunked")
            .with_i32_data(&data)
            .with_shape(dims)
            .with_chunks(chunks);
        builder
            .create_dataset("deflated")
            .with_i32_data(&data)
            .with_shape(dims)
            .with_chunks(chunks)
            .with_deflate(3);
        builder.write(&path).unwrap();

        let file = File::open(&path).unwrap();
        for name in ["contiguous", "chunked", "deflated"] {
            let ds = file.dataset(name).unwrap();
            let full = ds.read_i32().unwrap();
            assert_eq!(full, data, "{name} full read");

            for case in 0..60 {
                let selection = if case % 5 == 4 {
                    let points = (0..1 + rng.below(12))
                        .map(|_| dims.iter().map(|&d| rng.below(d)).collect())
                        .collect();
                    Selection::Points(points)
                } else {
                    random_hyperslab(&mut rng, dims)
                };
                assert_eq!(
                    ds.read_i32_selection(&selection).unwrap(),
                    reference(&full, dims, &selection),
                    "{name} rank {} case {case}: {selection:?}",
                    dims.len()
                );
            }
        }
    }
}

#[test]
fn out_of_bounds_selections_are_errors() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("oob.h5");
    let mut builder = FileBuilder::new();
    builder
        .create_dataset("d")
        .with_i32_data(&(0..100).collect::<Vec<i32>>())
        .with_shape(&[10, 10])
        .with_chunks(&[4, 4]);
    builder.write(&path).unwrap();
    let file = File::open(&path).unwrap();
    let ds = file.dataset("d").unwrap();
    let beyond = Selection::Hyperslab {
        start: vec![8, 8],
        stride: vec![1, 1],
        count: vec![5, 5],
        block: vec![1, 1],
    };
    use clawhdf5::Error;
    use clawhdf5_format::error::FormatError;
    let is_oob = |s: &Selection| {
        matches!(
            ds.read_i32_selection(s),
            Err(Error::Format(FormatError::SelectionOutOfBounds(_)))
        )
    };
    // Used to come back padded with zeros.
    assert!(is_oob(&beyond));
    // Row out of range.
    assert!(is_oob(&Selection::Points(vec![vec![10, 0]])));
    // Column out of range: used to wrap into the next row and return its value.
    assert!(is_oob(&Selection::Points(vec![vec![0, 12]])));
    // Wrong rank.
    assert!(is_oob(&Selection::Points(vec![vec![3]])));
    // In range is fine.
    assert_eq!(
        ds.read_i32_selection(&Selection::Points(vec![vec![9, 9]]))
            .unwrap(),
        [99]
    );
}