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

//! Chunked dataset reading: B-tree v1 type 1 traversal and chunk assembly.

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

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

use crate::addr::to_usize;
#[cfg(feature = "std")]
use crate::chunk_cache::{CacheAlignedBuffer, ChunkCache};
use crate::data_layout::DataLayout;
use crate::dataspace::Dataspace;
use crate::datatype::Datatype;
use crate::error::FormatError;
use crate::extensible_array::{ExtensibleArrayHeader, read_extensible_array_chunks};
use crate::filter_pipeline::FilterPipeline;
use crate::filters::{DecodeScratch, decompress_chunk_exact_with};
#[cfg(feature = "std")]
use crate::filters::{all_filters_skipped, decompress_chunk_exact};
use crate::fixed_array::{FixedArrayHeader, read_fixed_array_chunks};
#[cfg(feature = "std")]
use std::sync::Arc;

#[cfg(feature = "parallel")]
use crate::parallel_read;

#[cfg(feature = "parallel")]
use crate::lane_partition::PartitionStats;

/// Run `f` with this thread's chunk-decoding scratch buffers (see
/// [`DecodeScratch`]), kept between reads so decoding reuses memory instead
/// of faulting in fresh pages for every chunk. A re-entrant call (a
/// registered filter codec that itself reads a file) gets a fresh scratch.
pub(crate) fn with_scratch<R>(f: impl FnOnce(&mut DecodeScratch) -> R) -> R {
    #[cfg(feature = "std")]
    {
        use std::cell::RefCell;
        std::thread_local! {
            static SCRATCH: RefCell<DecodeScratch> = RefCell::new(DecodeScratch::new());
        }
        let mut f = Some(f);
        let kept = SCRATCH.try_with(|cell| {
            let mut scratch = cell.try_borrow_mut().ok()?;
            let f = f.take()?;
            let r = f(&mut scratch);
            scratch.trim();
            Some(r)
        });
        if let Ok(Some(r)) = kept {
            return r;
        }
        let f = f.expect("with_scratch: closure already run");
        f(&mut DecodeScratch::new())
    }
    #[cfg(not(feature = "std"))]
    f(&mut DecodeScratch::new())
}

/// A full read's output buffer, written through a raw pointer so that
/// several threads can place chunks into it at once.
struct OutBuf<'a> {
    ptr: *mut u8,
    len: usize,
    _borrow: core::marker::PhantomData<&'a mut [u8]>,
}

// SAFETY: `OutBuf` is a `&mut [u8]` that hands out writes; sharing it across
// threads is sound as long as concurrent writes do not overlap, which
// `write`'s contract requires.
unsafe impl Send for OutBuf<'_> {}
unsafe impl Sync for OutBuf<'_> {}

impl<'a> OutBuf<'a> {
    fn new(out: &'a mut [u8]) -> Self {
        Self {
            ptr: out.as_mut_ptr(),
            len: out.len(),
            _borrow: core::marker::PhantomData,
        }
    }

    fn len(&self) -> usize {
        self.len
    }

    /// Copy `src` to `[at, at + src.len())`. Out of range is a no-op (the
    /// callers check first).
    ///
    /// # Safety
    ///
    /// No other thread may be writing an overlapping range at the same time.
    unsafe fn write(&self, at: usize, src: &[u8]) {
        if at.checked_add(src.len()).is_none_or(|end| end > self.len) {
            debug_assert!(false, "OutBuf::write out of range");
            return;
        }
        // SAFETY: in range (checked above) of a live `&'a mut [u8]`; `src`
        // cannot overlap it (the output is exclusively borrowed); no
        // concurrent overlapping write (the caller's contract).
        unsafe { core::ptr::copy_nonoverlapping(src.as_ptr(), self.ptr.add(at), src.len()) }
    }
}

/// How a chunked dataset's chunks map into its row-major output.
struct ChunkPlacer {
    rank: usize,
    chunk_dims: Vec<usize>,
    ds_dims: Vec<usize>,
    ds_strides: Vec<usize>,
    chunk_strides: Vec<usize>,
    elem_size: usize,
}

impl ChunkPlacer {
    /// `chunk_dims` and `ds_dims` have one entry per dimension; the dataset
    /// must not be empty (so the stride products stay in range).
    fn new(chunk_dims: &[usize], ds_dims: &[usize], elem_size: usize) -> Self {
        let rank = chunk_dims.len();
        let mut ds_strides = vec![1usize; rank];
        let mut chunk_strides = vec![1usize; rank];
        for i in (0..rank.saturating_sub(1)).rev() {
            ds_strides[i] = ds_strides[i + 1] * ds_dims[i + 1];
            chunk_strides[i] = chunk_strides[i + 1] * chunk_dims[i + 1];
        }
        Self {
            rank,
            chunk_dims: chunk_dims.to_vec(),
            ds_dims: ds_dims.to_vec(),
            ds_strides,
            chunk_strides,
            elem_size,
        }
    }

    /// Copy the decoded chunk `data`, whose first element is at `offsets`
    /// (one per dimension), to its place in `out`.
    ///
    /// # Safety
    ///
    /// No other thread may be placing a chunk whose region overlaps this
    /// one's (see [`Self::regions_disjoint`]).
    unsafe fn place(&self, data: &[u8], offsets: &[u64], out: &OutBuf<'_>) {
        if self.rank == 0 {
            let copy_len = data.len().min(out.len());
            // SAFETY: in range; exclusivity is the caller's contract.
            unsafe { out.write(0, &data[..copy_len]) };
            return;
        }
        let mut short = [0usize; 8];
        let mut long = Vec::new();
        let chunk_offsets: &mut [usize] = if self.rank <= short.len() {
            &mut short[..self.rank]
        } else {
            long.resize(self.rank, 0);
            &mut long
        };
        for (o, &off) in chunk_offsets.iter_mut().zip(offsets) {
            // An offset past `usize` is past the dataset: the chunk writes
            // nothing (a plain cast would wrap it into the output).
            *o = usize::try_from(off).unwrap_or(usize::MAX);
        }
        // SAFETY: the caller's contract.
        unsafe {
            copy_chunk_into(
                data,
                out,
                chunk_offsets,
                &self.chunk_dims,
                &self.ds_dims,
                &self.ds_strides,
                &self.chunk_strides,
                self.elem_size,
                self.rank,
            )
        };
    }

    /// Whether placing `chunks` writes pairwise disjoint regions of the
    /// output, so they can be placed concurrently: every chunk starts on the
    /// chunk grid and no two start at the same place (a chunk wholly outside
    /// the dataset writes nothing and is ignored). A corrupt index that
    /// breaks this is read one chunk at a time instead.
    #[cfg(feature = "parallel")]
    fn regions_disjoint(&self, chunks: &[ChunkInfo]) -> bool {
        if self.rank == 0 {
            return chunks.len() <= 1;
        }
        let mut keys = Vec::with_capacity(chunks.len());
        'chunks: for c in chunks {
            if c.offsets.len() < self.rank {
                return false;
            }
            let mut key = 0u64;
            for d in 0..self.rank {
                let (off, cd, dd) = (
                    c.offsets[d],
                    self.chunk_dims[d] as u64,
                    self.ds_dims[d] as u64,
                );
                if off >= dd {
                    continue 'chunks;
                }
                if off % cd != 0 {
                    return false;
                }
                let Some(k) = key
                    .checked_mul(dd.div_ceil(cd))
                    .and_then(|k| k.checked_add(off / cd))
                else {
                    return false;
                };
                key = k;
            }
            keys.push(key);
        }
        keys.sort_unstable();
        keys.windows(2).all(|w| w[0] != w[1])
    }
}

/// The per-file chunk cache a full read uses, if any: the cache, this
/// dataset's key in it (its chunk-index address), and whether the dataset
/// fits in it — only then are its chunks looked up and inserted.
#[cfg(feature = "std")]
type CacheUse<'a> = Option<(&'a ChunkCache, u64, bool)>;
#[cfg(not(feature = "std"))]
type CacheUse<'a> = Option<&'a core::convert::Infallible>;

/// Decode every chunk in `chunks` and place it in `output` (the dataset's
/// whole row-major extent, zeroed).
///
/// Each chunk goes straight from its decoder to its place in the output:
/// decoded into this thread's reusable scratch (or, when the cache is to
/// keep it, into a buffer the cache takes), then copied. With the
/// `parallel` feature and a filter pipeline, the chunks are shared out
/// between the calling thread and idle rayon workers
/// ([`parallel_read::run_with_helpers`]), so the caller never waits on a
/// busy pool. The error returned is the first failing chunk's, in `chunks`
/// order.
#[allow(clippy::too_many_arguments)]
fn fill_from_chunks(
    file_data: &[u8],
    chunks: &[ChunkInfo],
    pipeline: Option<&FilterPipeline>,
    placer: &ChunkPlacer,
    chunk_total_bytes: usize,
    cache: CacheUse<'_>,
    output: &mut [u8],
) -> Result<(), FormatError> {
    let rank = placer.rank;
    let elem_size = placer.elem_size as u32;
    let out = OutBuf::new(output);
    #[cfg(not(feature = "std"))]
    let _ = cache;

    // Decode chunk `i` and place it. Its callers below run it either on one
    // thread, or on several for chunks whose regions are pairwise disjoint,
    // each chunk once: no two threads ever write the same bytes.
    let work = |i: usize, scratch: &mut DecodeScratch| -> Result<(), FormatError> {
        let c = &chunks[i];
        if c.offsets.len() < rank {
            return Err(FormatError::ChunkedReadError(format!(
                "chunk index entry has {} offsets for a rank-{rank} dataset",
                c.offsets.len()
            )));
        }
        let offsets = &c.offsets[..rank];
        let c_addr = to_usize(c.address)?;
        let size = c.chunk_size as usize;
        ensure_len(file_data, c_addr, size)?;
        let raw = &file_data[c_addr..c_addr + size];
        let Some(pl) = pipeline else {
            // SAFETY: see above.
            unsafe { placer.place(raw, offsets, &out) };
            return Ok(());
        };
        // A chunk stored as-is (every filter skipped) is checked and placed
        // straight from the file bytes, never cached.
        #[cfg(feature = "std")]
        if let Some((cache, key, true)) = cache
            && !all_filters_skipped(pl, c.filter_mask)
        {
            let cached = match cache.get_decompressed_in(key, offsets) {
                Some(hit) => hit,
                None => {
                    let data = decompress_chunk_exact(
                        raw,
                        pl,
                        chunk_total_bytes,
                        elem_size,
                        c.filter_mask,
                        &c.offsets,
                    )?;
                    cache.put_decompressed_in(key, offsets.to_vec(), data)
                }
            };
            // SAFETY: see above.
            unsafe { placer.place(&cached, offsets, &out) };
            return Ok(());
        }
        let data = decompress_chunk_exact_with(
            raw,
            pl,
            chunk_total_bytes,
            elem_size,
            c.filter_mask,
            &c.offsets,
            scratch,
        )?;
        // SAFETY: see above.
        unsafe { placer.place(data, offsets, &out) };
        Ok(())
    };

    #[cfg(feature = "parallel")]
    if pipeline.is_some()
        && parallel_read::should_use_parallel(chunks.len())
        && placer.regions_disjoint(chunks)
    {
        use core::sync::atomic::{AtomicUsize, Ordering};
        use std::sync::{Mutex, PoisonError};

        let n = chunks.len();
        let next = AtomicUsize::new(0);
        let failed: Mutex<Option<(usize, FormatError)>> = Mutex::new(None);
        let body = || {
            with_scratch(|scratch| {
                loop {
                    let i = next.fetch_add(1, Ordering::Relaxed);
                    if i >= n {
                        break;
                    }
                    if let Err(e) = work(i, scratch) {
                        let mut failed = failed.lock().unwrap_or_else(PoisonError::into_inner);
                        if failed.as_ref().is_none_or(|(at, _)| i < *at) {
                            *failed = Some((i, e));
                        }
                        // Stop handing out chunks. Every chunk before `i` was
                        // claimed already and finishes, so the error kept is
                        // the first in order.
                        next.store(n, Ordering::Relaxed);
                        break;
                    }
                }
            })
        };
        parallel_read::run_with_helpers(parallel_read::helper_count(n), &body);
        return match failed.into_inner().unwrap_or_else(PoisonError::into_inner) {
            Some((_, e)) => Err(e),
            None => Ok(()),
        };
    }

    with_scratch(|scratch| (0..chunks.len()).try_for_each(|i| work(i, scratch)))
}

/// Decompress all chunks with lane-partitioned parallelism and return
/// per-lane diagnostics.
///
/// This is the stats-returning variant for callers who want to inspect
/// the partition balance.  Only available with the `parallel` feature.
#[cfg(feature = "parallel")]
pub fn decompress_all_chunks_with_stats(
    file_data: &[u8],
    chunks: &[ChunkInfo],
    pipeline: &FilterPipeline,
    chunk_total_bytes: usize,
    element_size: u32,
    seed: u64,
    num_lanes: Option<usize>,
) -> Result<(Vec<Vec<u8>>, PartitionStats), FormatError> {
    parallel_read::decompress_chunks_lane_partitioned(
        file_data,
        chunks,
        pipeline,
        chunk_total_bytes,
        element_size,
        seed,
        num_lanes,
    )
}

/// Information about a single chunk in a chunked dataset.
#[derive(Debug, Clone)]
pub struct ChunkInfo {
    /// Size of chunk data in the file (after compression).
    pub chunk_size: u32,
    /// Bitmask of filters that were NOT applied (0 = all applied).
    pub filter_mask: u32,
    /// N-dimensional offset of this chunk in dataset space.
    pub offsets: Vec<u64>,
    /// File address of the chunk data.
    pub address: u64,
}

/// Checks that `[offset, offset + needed)` fits within `data`, guarding the
/// addition against `usize` overflow from a crafted near-`usize::MAX` offset.
fn ensure_len(data: &[u8], offset: usize, needed: usize) -> Result<(), FormatError> {
    if offset
        .checked_add(needed)
        .is_none_or(|end| end > data.len())
    {
        return Err(FormatError::UnexpectedEof {
            expected: offset.saturating_add(needed),
            available: data.len(),
        });
    }
    Ok(())
}

/// `elements * elem_size` for sizes that come from the file. Dataspace and
/// chunk dimensions are untrusted 64-bit fields, so a crafted file can make
/// the plain product wrap to a small number (or to something enormous).
pub(crate) fn checked_byte_len(elements: u64, elem_size: usize) -> Result<usize, FormatError> {
    usize::try_from(elements)
        .ok()
        .and_then(|n| n.checked_mul(elem_size))
        .ok_or_else(|| {
            FormatError::Overflow(format!(
                "{elements} elements of {elem_size} bytes exceeds the addressable size"
            ))
        })
}

/// The spatial chunk dimensions of a chunked layout (`chunk_dimensions` is
/// the layout message's list: one per dataspace dimension, then the element
/// size), after the checks libhdf5 makes when it opens a chunked dataset
/// (`H5D__chunk_init` / `H5D__chunk_set_sizes`): the chunk rank must match
/// the dataspace's, no chunk dimension may be 0, and a chunk indexed by a
/// version-1 B-tree (`layout_version` below 4) may not be 4 GiB or more (the
/// B-tree records chunk sizes in 32 bits; libhdf5: "chunk size must be < 4GB
/// with v1 b-tree index"). The other chunk indexes allow larger chunks:
/// HDF5 2.0 writes them with layout version 5. A zero chunk dimension used
/// to read as all fill values, and a huge one to hang the reader.
pub(crate) fn chunk_geometry(
    chunk_dimensions: &[u32],
    layout_version: u8,
    dataspace: &Dataspace,
    elem_size: usize,
) -> Result<(usize, Vec<usize>), FormatError> {
    let rank = chunk_dimensions.len().checked_sub(1).ok_or_else(|| {
        FormatError::InvalidChunkDimensions("chunked layout has no dimensions".into())
    })?;
    if dataspace.dimensions.len() != rank {
        return Err(FormatError::InvalidChunkDimensions(format!(
            "dimensionality of chunks doesn't match the dataspace (chunk rank {rank}, \
             dataspace rank {})",
            dataspace.dimensions.len()
        )));
    }
    let spatial = &chunk_dimensions[..rank];
    if let Some(d) = spatial.iter().position(|&c| c == 0) {
        return Err(FormatError::InvalidChunkDimensions(format!(
            "chunk size must be > 0, dim = {d}"
        )));
    }
    let bytes = spatial
        .iter()
        .fold(elem_size as u128, |acc, &c| acc * u128::from(c));
    if layout_version < 4 && bytes > u128::from(u32::MAX) {
        return Err(FormatError::InvalidChunkDimensions(format!(
            "chunk size must be < 4GB with v1 b-tree index (chunk {spatial:?} of {elem_size}-byte elements)"
        )));
    }
    Ok((rank, spatial.iter().map(|&c| c as usize).collect()))
}

/// The size of one element of `dt` as stored in the file: a
/// variable-length element is its length (4), a global heap address
/// (`offset_size`) and an index (4), not the 16 of [`Datatype::type_size`].
fn stored_element_size(dt: &Datatype, offset_size: u8) -> u64 {
    match dt {
        Datatype::VariableLength { .. } => 8 + u64::from(offset_size),
        Datatype::Array {
            base_type,
            dimensions,
        } => dimensions
            .iter()
            .fold(stored_element_size(base_type, offset_size), |acc, &d| {
                acc.saturating_mul(u64::from(d))
            }),
        _ => u64::from(dt.type_size()),
    }
}

/// A chunked layout records the element size as its last dimension, and
/// libhdf5 refuses a dataset whose datatype has another size
/// (`H5D__chunk_set_sizes`: "stored datatype size in chunk layout does not
/// match datatype description"). Reading it anyway laid the chunks out with
/// the wrong element size.
pub(crate) fn check_chunk_element_size(
    layout: &DataLayout,
    datatype: &Datatype,
    offset_size: u8,
) -> Result<(), FormatError> {
    let DataLayout::Chunked {
        chunk_dimensions, ..
    } = layout
    else {
        return Ok(());
    };
    let Some(&stored) = chunk_dimensions.last() else {
        return Ok(());
    };
    let expected = stored_element_size(datatype, offset_size);
    if u64::from(stored) != expected {
        return Err(FormatError::InvalidChunkDimensions(format!(
            "stored datatype size in chunk layout does not match datatype description \
             (layout {stored} bytes, datatype {expected})"
        )));
    }
    Ok(())
}

/// Product of chunk dimensions times the element size, overflow-checked.
pub(crate) fn checked_chunk_byte_len(
    chunk_dims: &[usize],
    elem_size: usize,
) -> Result<usize, FormatError> {
    chunk_dims
        .iter()
        .try_fold(elem_size, |acc, &d| acc.checked_mul(d))
        .ok_or_else(|| {
            FormatError::Overflow(format!(
                "chunk dimensions {chunk_dims:?} x {elem_size} bytes exceeds the addressable size"
            ))
        })
}

/// A zero-filled output buffer of `len` bytes. `vec![0; len]` aborts the
/// process when the allocation fails; a size taken from the file must surface
/// as an error instead.
pub(crate) fn alloc_output(len: usize) -> Result<Vec<u8>, FormatError> {
    if len == 0 {
        return Ok(Vec::new());
    }
    let failed =
        || FormatError::Overflow(format!("cannot allocate {len} bytes for dataset output"));
    let layout = core::alloc::Layout::array::<u8>(len).map_err(|_| failed())?;
    // Ask the allocator for zeroed memory instead of reserving and then
    // writing zeros: for a large buffer the OS hands out already-zero pages
    // lazily, where an explicit fill touches every page up front — and most of
    // the buffer is about to be overwritten with chunk data anyway.
    //
    // SAFETY (both arms): `layout` has non-zero size (len > 0) and alignment 1.
    #[cfg(feature = "std")]
    let ptr = unsafe { std::alloc::alloc_zeroed(layout) };
    #[cfg(not(feature = "std"))]
    let ptr = unsafe { alloc::alloc::alloc_zeroed(layout) };
    if ptr.is_null() {
        return Err(failed());
    }
    // Before anything writes to it, so a large buffer faults in huge pages.
    crate::bulk_alloc::advise_huge_pages(ptr, len);
    // SAFETY: `ptr` came from the global allocator with the layout of
    // `[u8; len]`, which is exactly what `Vec<u8>` with capacity `len` frees;
    // all `len` bytes are initialised (zero).
    Ok(unsafe { Vec::from_raw_parts(ptr, len, len) })
}

fn read_offset(data: &[u8], pos: usize, size: u8) -> Result<u64, FormatError> {
    let s = size as usize;
    if pos.checked_add(s).is_none_or(|end| end > data.len()) {
        return Err(FormatError::UnexpectedEof {
            expected: pos.saturating_add(s),
            available: data.len(),
        });
    }
    let slice = &data[pos..pos + s];
    Ok(match size {
        2 => u16::from_le_bytes([slice[0], slice[1]]) as u64,
        4 => u32::from_le_bytes([slice[0], slice[1], slice[2], slice[3]]) as u64,
        8 => u64::from_le_bytes([
            slice[0], slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7],
        ]),
        _ => return Err(FormatError::InvalidOffsetSize(size)),
    })
}

/// Traverse B-tree v1 type 1 to collect all chunk locations.
///
/// `ndims` is the number of offset dimensions in each key, which equals
/// `chunk_dimensions.len()` from the DataLayout::Chunked message (rank+1).
pub fn collect_chunk_info(
    file_data: &[u8],
    btree_address: u64,
    ndims: usize,
    offset_size: u8,
    length_size: u8,
) -> Result<Vec<ChunkInfo>, FormatError> {
    let _ = length_size;
    let mut chunks = Vec::new();
    parse_chunk_node(
        file_data,
        btree_address,
        ndims,
        None,
        offset_size,
        0,
        &mut chunks,
    )?;
    Ok(chunks)
}

/// The chunks of a v1 B-tree chunk index as libhdf5 reads them, for a
/// layout with these `chunk_dimensions` (the layout message's list, element
/// size last).
///
/// Every key of the B-tree is checked as libhdf5 checks it
/// (`H5D__btree_decode_key`): each coordinate offset must be a multiple of
/// its chunk dimension. That includes the keys that only bound a node
/// (internal-node keys and each node's final key), which is where a
/// corrupt chunk dimension shows when the chunks themselves all start at
/// offset 0 in that dimension (`cve-2018-11205`). A key that fails ("bad
/// coordinate offset") means a corrupt index or chunk dimension.
///
/// libhdf5 does not read a chunk by walking the tree: it looks each chunk
/// up (`H5B_find` with `H5D__btree_cmp3` and `H5D__btree_found`), comparing
/// the element-size coordinate too, which it asks for as 0. So a chunk is
/// returned only where that lookup finds it: a key whose element-size
/// coordinate is not 0 is found in a 1-D dataset (the comparison looks at
/// that coordinate only against the next key) but not in a dataset of rank
/// 2 or more (`cve-2025-44905` `/Shuffle_float_data_le`), which then reads
/// as fill values, and a tree whose keys are out of order loses the chunks
/// libhdf5's binary search misses.
pub fn collect_chunk_info_checked(
    file_data: &[u8],
    btree_address: u64,
    chunk_dimensions: &[u32],
    offset_size: u8,
    length_size: u8,
) -> Result<Vec<ChunkInfo>, FormatError> {
    let _ = length_size;
    let ndims = chunk_dimensions.len();
    if ndims == 0 {
        return Err(FormatError::ChunkedReadError(
            "chunk layout has no dimensions".into(),
        ));
    }
    let mut stored = Vec::new();
    let mut root = parse_chunk_node(
        file_data,
        btree_address,
        ndims,
        Some(chunk_dimensions),
        offset_size,
        0,
        &mut stored,
    )?;
    // Keys were checked to be multiples of non-zero dimensions.
    root.scale_keys(chunk_dimensions);
    // Look every stored chunk's position up. A lookup finds a chunk only at
    // that chunk's own position, so each is returned at most once.
    let mut returned = vec![false; stored.len()];
    let mut wanted = vec![0u64; ndims];
    for chunk in &stored {
        for (w, (&o, &d)) in wanted
            .iter_mut()
            .zip(chunk.offsets.iter().zip(chunk_dimensions))
        {
            *w = o / u64::from(d);
        }
        wanted[ndims - 1] = 0;
        if let Some(i) = root.find(&wanted) {
            returned[i] = true;
        }
    }
    Ok(stored
        .into_iter()
        .zip(returned)
        .filter_map(|(c, r)| r.then_some(c))
        .collect())
}

/// A node of a v1 B-tree chunk index: its `n + 1` keys (`ndims`
/// coordinates each, flattened; byte offsets as stored, or scaled by the
/// chunk dimensions after [`ChunkNode::scale_keys`]) and its `n` children,
/// a leaf's as indices into the list of stored chunks.
struct ChunkNode {
    ndims: usize,
    keys: Vec<u64>,
    children: ChunkChildren,
}

enum ChunkChildren {
    Nodes(Vec<ChunkNode>),
    Chunks(Vec<usize>),
}

impl ChunkNode {
    fn key(&self, i: usize) -> &[u64] {
        &self.keys[i * self.ndims..(i + 1) * self.ndims]
    }

    fn len(&self) -> usize {
        match &self.children {
            ChunkChildren::Nodes(n) => n.len(),
            ChunkChildren::Chunks(c) => c.len(),
        }
    }

    fn scale_keys(&mut self, dims: &[u32]) {
        for (k, &d) in self.keys.iter_mut().zip(dims.iter().cycle()) {
            *k /= u64::from(d);
        }
        if let ChunkChildren::Nodes(nodes) = &mut self.children {
            for n in nodes {
                n.scale_keys(dims);
            }
        }
    }

    /// `H5B_find_helper` over scaled keys: binary search for the child
    /// whose keys bracket `scaled`, then `H5D__btree_found` at the leaf.
    /// Returns the index of the chunk found.
    fn find(&self, scaled: &[u64]) -> Option<usize> {
        let (mut lt, mut rt) = (0, self.len());
        let mut idx = 0;
        let mut cmp = core::cmp::Ordering::Greater;
        while lt < rt && cmp != core::cmp::Ordering::Equal {
            idx = (lt + rt) / 2;
            cmp = btree_cmp3(self.key(idx), scaled, self.key(idx + 1));
            if cmp == core::cmp::Ordering::Less {
                rt = idx;
            } else {
                lt = idx + 1;
            }
        }
        if cmp != core::cmp::Ordering::Equal {
            return None;
        }
        match &self.children {
            ChunkChildren::Nodes(nodes) => nodes[idx].find(scaled),
            ChunkChildren::Chunks(chunks) => {
                // "Is this *really* the requested chunk?"
                let lt_key = self.key(idx);
                let found = scaled
                    .iter()
                    .zip(lt_key)
                    .all(|(&s, &k)| s < k.wrapping_add(1));
                found.then_some(chunks[idx])
            }
        }
    }
}

/// `H5D__btree_cmp3`: where `scaled` falls against a child's left and
/// right keys. `Less` is left of the child, `Greater` right of it. With a
/// rank-1 dataset (two coordinates, element size last) libhdf5 compares
/// only the first coordinate, and the second against the right key.
fn btree_cmp3(lt: &[u64], scaled: &[u64], rt: &[u64]) -> core::cmp::Ordering {
    use core::cmp::Ordering;
    if scaled.len() == 2 {
        if scaled[0] > rt[0] || (scaled[0] == rt[0] && scaled[1] >= rt[1]) {
            Ordering::Greater
        } else if scaled[0] < lt[0] {
            Ordering::Less
        } else {
            Ordering::Equal
        }
    } else if scaled >= rt {
        Ordering::Greater
    } else if scaled < lt {
        Ordering::Less
    } else {
        Ordering::Equal
    }
}

/// Check one v1 B-tree chunk key's offsets (see
/// [`collect_chunk_info_checked`]).
fn check_key_offsets(offsets: &[u64], chunk_dimensions: &[u32]) -> Result<(), FormatError> {
    for (&offset, &dim) in offsets.iter().zip(chunk_dimensions) {
        if dim == 0 || offset % u64::from(dim) != 0 {
            return Err(FormatError::ChunkedReadError(format!(
                "bad coordinate offset {offsets:?} for chunk dimensions {chunk_dimensions:?}"
            )));
        }
    }
    Ok(())
}

/// Read the `ndims` 8-byte offsets of the chunk key at `pos` (after its
/// chunk size and filter mask) into `out`, checking them when
/// `chunk_dimensions` is given.
fn read_key_offsets(
    file_data: &[u8],
    pos: usize,
    ndims: usize,
    chunk_dimensions: Option<&[u32]>,
    out: &mut Vec<u64>,
) -> Result<(), FormatError> {
    let start = out.len();
    let mut kp = pos + 8;
    for _ in 0..ndims {
        out.push(read_offset(file_data, kp, CHUNK_KEY_OFFSET_SIZE)?);
        kp += CHUNK_KEY_OFFSET_SIZE as usize;
    }
    if let Some(dims) = chunk_dimensions {
        check_key_offsets(&out[start..], dims)?;
    }
    Ok(())
}

/// Width of each chunk offset in a v1 chunk B-tree key, independent of the
/// file's size-of-offsets.
const CHUNK_KEY_OFFSET_SIZE: u8 = 8;

/// Maximum recursion depth for chunk B-tree traversal (malformed/cyclic data
/// protection), matching `btree_v1.rs`'s `MAX_BTREE_DEPTH`.
const MAX_CHUNK_BTREE_DEPTH: usize = 64;

/// Parse the v1 B-tree chunk index node at `btree_address` and its
/// subtree, appending its chunks to `stored` in tree order.
fn parse_chunk_node(
    file_data: &[u8],
    btree_address: u64,
    ndims: usize,
    chunk_dimensions: Option<&[u32]>,
    offset_size: u8,
    depth: usize,
    stored: &mut Vec<ChunkInfo>,
) -> Result<ChunkNode, FormatError> {
    if depth > MAX_CHUNK_BTREE_DEPTH {
        return Err(FormatError::NestingDepthExceeded);
    }

    let offset = to_usize(btree_address)?;
    let os = offset_size as usize;

    // Parse B-tree v1 header
    let header_size = 8 + os * 2;
    ensure_len(file_data, offset, header_size)?;

    if &file_data[offset..offset + 4] != b"TREE" {
        return Err(FormatError::InvalidBTreeSignature);
    }

    let node_type = file_data[offset + 4];
    if node_type != 1 {
        return Err(FormatError::InvalidBTreeNodeType(node_type));
    }

    let node_level = file_data[offset + 5];
    let entries_used = u16::from_le_bytes([file_data[offset + 6], file_data[offset + 7]]) as usize;

    let mut pos = offset + 8 + os * 2; // skip left/right sibling

    // Key: chunk_size(4) + filter_mask(4) + one offset per dimension. The
    // offsets are always 8 bytes each — they are dataset coordinates, not file
    // addresses, so they do not follow the superblock's size-of-offsets (only
    // the sibling and child addresses do).
    let key_size = ndims
        .checked_mul(CHUNK_KEY_OFFSET_SIZE as usize)
        .and_then(|n| n.checked_add(8))
        .ok_or_else(|| FormatError::ChunkedReadError("chunk key too large".into()))?;

    // key[0], child[0], key[1], child[1], ..., key[N-1], child[N-1], key[N]
    let needed = entries_used * (key_size + os) + key_size;
    ensure_len(file_data, pos, needed)?;

    let mut keys = Vec::with_capacity((entries_used + 1) * ndims);
    let mut chunks = Vec::new();
    let mut child_addrs = Vec::new();
    for _ in 0..entries_used {
        let chunk_size = u32::from_le_bytes([
            file_data[pos],
            file_data[pos + 1],
            file_data[pos + 2],
            file_data[pos + 3],
        ]);
        let filter_mask = u32::from_le_bytes([
            file_data[pos + 4],
            file_data[pos + 5],
            file_data[pos + 6],
            file_data[pos + 7],
        ]);
        let k = keys.len();
        read_key_offsets(file_data, pos, ndims, chunk_dimensions, &mut keys)?;
        pos += key_size;
        let address = read_offset(file_data, pos, offset_size)?;
        pos += os;
        if node_level == 0 {
            chunks.push(stored.len());
            stored.push(ChunkInfo {
                chunk_size,
                filter_mask,
                offsets: keys[k..].to_vec(),
                address,
            });
        } else {
            child_addrs.push(address);
        }
    }
    // The final key only bounds the node; libhdf5 still checks it.
    read_key_offsets(file_data, pos, ndims, chunk_dimensions, &mut keys)?;

    let children = if node_level == 0 {
        ChunkChildren::Chunks(chunks)
    } else {
        let mut nodes = Vec::with_capacity(child_addrs.len());
        for child_addr in child_addrs {
            nodes.push(parse_chunk_node(
                file_data,
                child_addr,
                ndims,
                chunk_dimensions,
                offset_size,
                depth + 1,
                stored,
            )?);
        }
        ChunkChildren::Nodes(nodes)
    };
    Ok(ChunkNode {
        ndims,
        keys,
        children,
    })
}

/// Generate ChunkInfo entries for an implicit index (v4 index type 2).
///
/// Chunks are stored contiguously starting at `base_address`. No stored index;
/// addresses are computed from the chunk position.
///
/// `chunk_dimensions` are the spatial chunk dimensions, one per entry of
/// `dataset_dims` — not the layout message's list, which carries the element
/// size as an extra last dimension.
pub fn generate_implicit_chunks(
    base_address: u64,
    dataset_dims: &[u64],
    chunk_dimensions: &[u32],
    element_size: u32,
) -> Vec<ChunkInfo> {
    let rank = chunk_dimensions.len();
    let chunk_byte_size: u64 =
        chunk_dimensions.iter().map(|&d| d as u64).product::<u64>() * element_size as u64;

    let mut num_chunks_per_dim = Vec::with_capacity(rank);
    for d in 0..rank {
        let ds = dataset_dims[d];
        let ch = chunk_dimensions[d] as u64;
        num_chunks_per_dim.push(ds.div_ceil(ch));
    }
    let total_chunks: u64 = num_chunks_per_dim.iter().product();

    // A capacity hint only (a count past `usize::MAX` could not be pushed).
    let mut chunks = Vec::with_capacity(usize::try_from(total_chunks).unwrap_or(0));
    for linear_idx in 0..total_chunks {
        let mut offsets = vec![0u64; rank];
        let mut remaining = linear_idx;
        for d in (0..rank).rev() {
            let nchunks = num_chunks_per_dim[d];
            let chunk_idx = remaining % nchunks;
            remaining /= nchunks;
            offsets[d] = chunk_idx * chunk_dimensions[d] as u64;
        }

        chunks.push(ChunkInfo {
            chunk_size: chunk_byte_size as u32,
            filter_mask: 0,
            offsets,
            address: base_address + linear_idx * chunk_byte_size,
        });
    }

    chunks
}

/// B-tree v2 record types used for chunk indexing.
const BT2_CHUNK_UNFILTERED: u8 = 10;
const BT2_CHUNK_FILTERED: u8 = 11;

/// Chunks indexed by a version-2 B-tree (layout v4, index type 5).
///
/// Record layouts (all little endian):
/// * type 10, unfiltered: address, then one 8-byte *scaled* offset per
///   dimension (offset / chunk dimension);
/// * type 11, filtered: address, stored chunk size (a variable number of
///   bytes), 4-byte filter mask, then the scaled offsets.
///
/// The width of the stored-size field depends on the largest possible chunk;
/// rather than re-derive the library's formula it is taken from the record
/// size the tree header declares, which is what actually governs the bytes.
fn read_btree_v2_chunks(
    file_data: &[u8],
    addr: u64,
    chunk_dims: &[usize],
    elem_size: usize,
    offset_size: u8,
    length_size: u8,
) -> Result<Vec<ChunkInfo>, FormatError> {
    use crate::btree_v2::{BTreeV2Header, collect_btree_v2_records};

    let bad = |what: &str| FormatError::ChunkedReadError(format!("B-tree v2 chunk index: {what}"));
    let header = BTreeV2Header::parse(file_data, to_usize(addr)?, offset_size, length_size)?;
    let rank = chunk_dims.len();
    let os = offset_size as usize;
    let record_size = header.record_size as usize;
    let size_len = match header.tree_type {
        BT2_CHUNK_UNFILTERED => {
            if record_size != os + 8 * rank {
                return Err(bad("unexpected record size for unfiltered chunks"));
            }
            0
        }
        BT2_CHUNK_FILTERED => {
            let fixed = os + 4 + 8 * rank;
            let size_len = record_size
                .checked_sub(fixed)
                .ok_or_else(|| bad("record too small"))?;
            if !(1..=8).contains(&size_len) {
                return Err(bad("implausible chunk-size field width"));
            }
            size_len
        }
        _ => return Err(bad("tree is not a chunk index")),
    };
    let unfiltered_bytes = checked_chunk_byte_len(chunk_dims, elem_size)?;
    let unfiltered_bytes =
        u32::try_from(unfiltered_bytes).map_err(|_| bad("chunk larger than 4 GiB"))?;

    let records = collect_btree_v2_records(file_data, &header, offset_size, length_size)?;
    let mut chunks = Vec::with_capacity(records.len());
    for record in &records {
        let data = record.data.as_slice();
        if data.len() < record_size {
            return Err(bad("truncated record"));
        }
        let address = read_offset(data, 0, offset_size)?;
        let mut pos = os;
        let (chunk_size, filter_mask) = if size_len == 0 {
            (unfiltered_bytes, 0)
        } else {
            let mut size = 0u64;
            for (i, &b) in data[pos..pos + size_len].iter().enumerate() {
                size |= u64::from(b) << (8 * i);
            }
            pos += size_len;
            let mask = u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]);
            pos += 4;
            (
                u32::try_from(size).map_err(|_| bad("stored chunk larger than 4 GiB"))?,
                mask,
            )
        };
        let mut offsets = Vec::with_capacity(rank);
        for &dim in chunk_dims {
            let scaled = u64::from_le_bytes([
                data[pos],
                data[pos + 1],
                data[pos + 2],
                data[pos + 3],
                data[pos + 4],
                data[pos + 5],
                data[pos + 6],
                data[pos + 7],
            ]);
            pos += 8;
            offsets.push(
                scaled
                    .checked_mul(dim as u64)
                    .ok_or_else(|| bad("chunk offset overflows"))?,
            );
        }
        chunks.push(ChunkInfo {
            chunk_size,
            filter_mask,
            offsets,
            address,
        });
    }
    Ok(chunks)
}

/// Every allocated chunk of a chunked dataset, for any supported chunk index,
/// plus the spatial chunk dimensions. Chunks the file never allocated (sparse
/// datasets) are simply absent from the list.
pub fn list_chunks(
    file_data: &[u8],
    layout: &DataLayout,
    dataspace: &Dataspace,
    elem_size: usize,
    offset_size: u8,
    length_size: u8,
) -> Result<(Vec<ChunkInfo>, Vec<usize>), FormatError> {
    let (
        chunk_dimensions,
        version,
        chunk_index_type,
        addr_opt,
        single_filtered_size,
        single_filter_mask,
        unfiltered_edges,
    ) = match layout {
        DataLayout::Chunked {
            chunk_dimensions,
            btree_address,
            version,
            chunk_index_type,
            single_chunk_filtered_size,
            single_chunk_filter_mask,
            dont_filter_partial_edge_chunks,
        } => (
            chunk_dimensions,
            *version,
            *chunk_index_type,
            *btree_address,
            *single_chunk_filtered_size,
            *single_chunk_filter_mask,
            *dont_filter_partial_edge_chunks,
        ),
        _ => {
            return Err(FormatError::ChunkedReadError(
                "expected chunked layout".into(),
            ));
        }
    };

    let addr = addr_opt
        .ok_or_else(|| FormatError::ChunkedReadError("no address for chunked layout".into()))?;

    // Both v3 and v4 include element size as last dim (rank+1)
    let (rank, chunk_dims) = chunk_geometry(chunk_dimensions, version, dataspace, elem_size)?;
    let ds_dims: Vec<usize> = dataspace
        .dimensions
        .iter()
        .map(|&d| to_usize(d))
        .collect::<Result<_, _>>()?;

    // Collect chunks based on version and index type
    let mut chunks = match (version, chunk_index_type) {
        (3, _) => {
            collect_chunk_info_checked(file_data, addr, chunk_dimensions, offset_size, length_size)?
        }
        (4, Some(1)) => {
            // Single chunk — one chunk covering the entire dataset
            let chunk_byte_size = checked_chunk_byte_len(&chunk_dims, elem_size)?;
            let (csize, fmask) = if let Some(fs) = single_filtered_size {
                (fs as u32, single_filter_mask.unwrap_or(0))
            } else {
                (chunk_byte_size as u32, 0)
            };
            vec![ChunkInfo {
                chunk_size: csize,
                filter_mask: fmask,
                offsets: vec![0u64; rank],
                address: addr,
            }]
        }
        (4, Some(2)) => {
            // Implicit index — use spatial chunk dims only
            let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
            generate_implicit_chunks(
                addr,
                &dataspace.dimensions,
                spatial_chunk_dims,
                elem_size as u32,
            )
        }
        (4, Some(3)) => {
            // Fixed Array — use spatial chunk dims only
            let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
            let header =
                FixedArrayHeader::parse(file_data, to_usize(addr)?, offset_size, length_size)?;
            read_fixed_array_chunks(
                file_data,
                &header,
                &dataspace.dimensions,
                dataspace.max_dimensions.as_deref(),
                spatial_chunk_dims,
                elem_size as u32,
                offset_size,
                length_size,
            )?
        }
        (4, Some(4)) => {
            // Extensible Array — use spatial chunk dims only
            let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
            let header =
                ExtensibleArrayHeader::parse(file_data, to_usize(addr)?, offset_size, length_size)?;
            read_extensible_array_chunks(
                file_data,
                &header,
                &dataspace.dimensions,
                dataspace.max_dimensions.as_deref(),
                spatial_chunk_dims,
                elem_size as u32,
                offset_size,
                length_size,
            )?
        }
        (4, Some(5)) => {
            // Version-2 B-tree: what the library uses for a dataset with two
            // or more unlimited dimensions.
            read_btree_v2_chunks(
                file_data,
                addr,
                &chunk_dims,
                elem_size,
                offset_size,
                length_size,
            )?
        }
        (v, idx) => {
            return Err(FormatError::ChunkedReadError(format!(
                "unsupported chunked layout version={v}, index_type={idx:?}"
            )));
        }
    };

    // With "don't filter partial edge chunks", a chunk that extends past the
    // dataset's extent is stored raw while its filter mask still reads 0.
    // Mark every filter skipped so all read paths copy it as-is.
    if unfiltered_edges {
        for chunk in &mut chunks {
            let partial = chunk
                .offsets
                .iter()
                .zip(&chunk_dims)
                .zip(&ds_dims)
                .any(|((&off, &cd), &dd)| off.saturating_add(cd as u64) > dd as u64);
            if partial {
                chunk.filter_mask = u32::MAX;
            }
        }
    }

    Ok((chunks, chunk_dims))
}

/// [`list_chunks`] for reading the chunks through `pipeline`: a dataset
/// without filters stores every chunk at the chunk's full size, and a chunk
/// the index records at another size is refused, as libhdf5 refuses it
/// ("incorrect chunk size returned from index for unfiltered chunk"). Such
/// a chunk was read at its recorded size, with the rest of the chunk left
/// as zeros or fill values: `cve-2025-44904`'s `Scale_offset_float_data_le`
/// has chunks of 38 and 37 bytes for 48-byte chunks, where HDF5 2.0 reads
/// whatever its buffer held for the missing bytes.
pub fn list_chunks_for_read(
    file_data: &[u8],
    layout: &DataLayout,
    dataspace: &Dataspace,
    elem_size: usize,
    pipeline: Option<&FilterPipeline>,
    offset_size: u8,
    length_size: u8,
) -> Result<(Vec<ChunkInfo>, Vec<usize>), FormatError> {
    let (chunks, chunk_dims) = list_chunks(
        file_data,
        layout,
        dataspace,
        elem_size,
        offset_size,
        length_size,
    )?;
    if pipeline.is_none_or(|p| p.filters.is_empty()) {
        let chunk_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?;
        if let Some(c) = chunks
            .iter()
            .find(|c| c.address != u64::MAX && c.chunk_size as usize != chunk_bytes)
        {
            return Err(FormatError::ChunkedReadError(format!(
                "incorrect chunk size returned from index for unfiltered chunk at {:?}: \
                 {} bytes, expected {chunk_bytes}",
                c.offsets, c.chunk_size
            )));
        }
    }
    Ok((chunks, chunk_dims))
}

/// The chunk cache a full read may use (`None` without `std`).
#[cfg(feature = "std")]
pub(crate) type CacheRef<'a> = Option<&'a ChunkCache>;
/// The chunk cache a full read may use (`None` without `std`).
#[cfg(not(feature = "std"))]
pub(crate) type CacheRef<'a> = Option<&'a core::convert::Infallible>;

/// The body of every full chunked read: list the chunks (through the
/// cache's index for this dataset when there is a cache), allocate the
/// output with `alloc` (zeroed, `total_bytes` long, as bytes through
/// `bytes`), and decode every chunk straight into it.
#[allow(clippy::too_many_arguments)]
pub(crate) fn read_chunked_full<O>(
    file_data: &[u8],
    layout: &DataLayout,
    dataspace: &Dataspace,
    datatype: &Datatype,
    pipeline: Option<&FilterPipeline>,
    offset_size: u8,
    length_size: u8,
    cache: CacheRef<'_>,
    alloc: impl FnOnce(usize) -> Result<O, FormatError>,
    bytes: impl FnOnce(&mut O) -> &mut [u8],
) -> Result<O, FormatError> {
    check_chunk_element_size(layout, datatype, offset_size)?;
    let elem_size = datatype.type_size() as usize;
    let list = || {
        list_chunks_for_read(
            file_data,
            layout,
            dataspace,
            elem_size,
            pipeline,
            offset_size,
            length_size,
        )
    };

    #[cfg(feature = "std")]
    let (chunks, chunk_dims, cache_key) = match cache {
        Some(cache) => {
            let (chunk_dimensions, version, addr_opt) = match layout {
                DataLayout::Chunked {
                    chunk_dimensions,
                    version,
                    btree_address,
                    ..
                } => (chunk_dimensions, *version, *btree_address),
                _ => {
                    return Err(FormatError::ChunkedReadError(
                        "expected chunked layout".into(),
                    ));
                }
            };
            let addr = addr_opt.ok_or_else(|| {
                FormatError::ChunkedReadError("no address for chunked layout".into())
            })?;
            let (rank, chunk_dims) =
                chunk_geometry(chunk_dimensions, version, dataspace, elem_size)?;
            // The per-file cache is shared across datasets (and threads);
            // every lookup is keyed by this dataset's chunk-index address, so
            // another dataset's index or chunks are never used for this read.
            let chunks = cache.chunks_for(addr, rank, || list().map(|(chunks, _)| chunks))?;
            (chunks, chunk_dims, Some((cache, addr)))
        }
        None => {
            let (chunks, chunk_dims) = list()?;
            (chunks, chunk_dims, None)
        }
    };
    #[cfg(not(feature = "std"))]
    let (chunks, chunk_dims) = {
        let _ = cache;
        list()?
    };

    let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?;
    let mut output = alloc(total_bytes)?;
    if total_bytes == 0 {
        // Also keeps the stride products in range: with a zero-sized
        // dimension the total is 0 even if other dimensions are huge.
        return Ok(output);
    }
    let ds_dims: Vec<usize> = dataspace
        .dimensions
        .iter()
        .map(|&d| to_usize(d))
        .collect::<Result<_, _>>()?;
    let placer = ChunkPlacer::new(&chunk_dims, &ds_dims, elem_size);
    let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?;
    // Chunks are cached only when the whole dataset fits: pushing a larger
    // dataset through the cache just evicts each chunk moments after
    // inserting it.
    #[cfg(feature = "std")]
    let cache_use = cache_key.map(|(cache, key)| (cache, key, total_bytes <= cache.max_bytes()));
    #[cfg(not(feature = "std"))]
    let cache_use = None;
    fill_from_chunks(
        file_data,
        &chunks,
        pipeline,
        &placer,
        chunk_total_bytes,
        cache_use,
        bytes(&mut output),
    )?;
    Ok(output)
}

/// Read a chunked dataset, decompressing chunks as needed.
pub fn read_chunked_data(
    file_data: &[u8],
    layout: &DataLayout,
    dataspace: &Dataspace,
    datatype: &Datatype,
    pipeline: Option<&FilterPipeline>,
    offset_size: u8,
    length_size: u8,
) -> Result<Vec<u8>, FormatError> {
    read_chunked_full(
        file_data,
        layout,
        dataspace,
        datatype,
        pipeline,
        offset_size,
        length_size,
        None,
        alloc_output,
        |out| out.as_mut_slice(),
    )
}

/// Read a chunked dataset with caching support.
///
/// On the first call, scans the chunk index (B-tree / fixed array / etc.) once
/// and populates the cache's hash index.  Subsequent calls skip the index scan
/// entirely.  Decompressed chunk data is also cached with LRU eviction, when
/// the whole dataset fits in the cache.
#[cfg(feature = "std")]
#[allow(clippy::too_many_arguments)]
pub fn read_chunked_data_cached(
    file_data: &[u8],
    layout: &DataLayout,
    dataspace: &Dataspace,
    datatype: &Datatype,
    pipeline: Option<&FilterPipeline>,
    offset_size: u8,
    length_size: u8,
    cache: &ChunkCache,
) -> Result<Vec<u8>, FormatError> {
    read_chunked_full(
        file_data,
        layout,
        dataspace,
        datatype,
        pipeline,
        offset_size,
        length_size,
        Some(cache),
        alloc_output,
        |out| out.as_mut_slice(),
    )
}

/// Sweep context passed into `read_chunked_data_sweep` to enable adaptive
/// prefetching based on detected access patterns.
///
/// The caller is responsible for maintaining the `SweepContext` across
/// multiple reads on the same dataset. After each read, the context will
/// contain updated sweep detection state and any predicted next-chunk
/// coordinates.
pub struct SweepContext {
    /// Sliding window of recent chunk coordinates.
    pub history: Vec<Vec<u64>>,
    /// Maximum window size.
    pub window_size: usize,
    /// Currently detected sweep direction label.
    pub direction: &'static str,
    /// How many chunks ahead to predict.
    pub prefetch_count: usize,
    /// Predicted next chunk coordinates (populated after each read).
    pub predicted_next: Vec<Vec<u64>>,
}

impl SweepContext {
    /// Create a new sweep context with the given window size and prefetch count.
    pub fn new(window_size: usize, prefetch_count: usize) -> Self {
        Self {
            history: Vec::with_capacity(window_size),
            window_size,
            direction: "random",
            prefetch_count,
            predicted_next: Vec::new(),
        }
    }

    /// Create with default settings (window=12, prefetch=4).
    pub fn with_defaults() -> Self {
        Self::new(12, 4)
    }

    /// Record a chunk coordinate access and update predictions.
    fn record(&mut self, coord: Vec<u64>, ndims: usize) {
        if self.history.len() >= self.window_size {
            self.history.remove(0);
        }
        self.history.push(coord);

        if self.history.len() < 3 || ndims == 0 {
            self.direction = "random";
            self.predicted_next.clear();
            return;
        }

        // Inline sweep detection matching the algorithm in clawhdf5-io/sweep.rs
        let num_deltas = self.history.len() - 1;
        let mut changing = vec![0usize; ndims];
        for i in 0..num_deltas {
            let prev = &self.history[i];
            let curr = &self.history[i + 1];
            if prev.len() < ndims || curr.len() < ndims {
                self.direction = "random";
                self.predicted_next.clear();
                return;
            }
            for d in 0..ndims {
                if curr[d] != prev[d] {
                    changing[d] += 1;
                }
            }
        }

        let threshold = num_deltas.div_ceil(2);
        let (max_dim, max_changes) = changing.iter().enumerate().max_by_key(|(_, c)| *c).unwrap();

        if *max_changes < threshold {
            self.direction = "random";
            self.predicted_next.clear();
            return;
        }

        let others_max = changing
            .iter()
            .enumerate()
            .filter(|(d, _)| *d != max_dim)
            .map(|(_, c)| *c)
            .max()
            .unwrap_or(0);

        if others_max > 0 && *max_changes < others_max * 2 {
            self.direction = "random";
            self.predicted_next.clear();
            return;
        }

        self.direction = if max_dim == ndims - 1 {
            "row_major"
        } else if max_dim == 0 {
            "column_major"
        } else {
            "slice_major"
        };

        // Predict next chunks
        let sweep_dim = max_dim;
        let mut total_step: i64 = 0;
        let mut step_count: usize = 0;
        for i in 1..self.history.len() {
            let prev = self.history[i - 1][sweep_dim] as i64;
            let curr = self.history[i][sweep_dim] as i64;
            let diff = curr - prev;
            if diff != 0 {
                total_step += diff;
                step_count += 1;
            }
        }

        if step_count == 0 {
            self.predicted_next.clear();
            return;
        }

        let avg_step = total_step / step_count as i64;
        if avg_step == 0 {
            self.predicted_next.clear();
            return;
        }

        let last = self.history.last().unwrap();
        self.predicted_next.clear();
        for i in 1..=self.prefetch_count {
            let mut pred = last.clone();
            let new_val = last[sweep_dim] as i64 + avg_step * i as i64;
            if new_val < 0 {
                break;
            }
            pred[sweep_dim] = new_val as u64;
            self.predicted_next.push(pred);
        }
    }
}

/// Read a chunked dataset with caching and sweep-aware prefetching.
///
/// Extends `read_chunked_data_cached` by feeding each chunk coordinate to a
/// [`SweepContext`]. When a sweep pattern is detected, predicted next-chunk
/// coordinates are pre-populated in the cache index via `prefetch_hint`.
#[cfg(feature = "std")]
#[allow(clippy::too_many_arguments)]
pub fn read_chunked_data_sweep(
    file_data: &[u8],
    layout: &DataLayout,
    dataspace: &Dataspace,
    datatype: &Datatype,
    pipeline: Option<&FilterPipeline>,
    offset_size: u8,
    length_size: u8,
    cache: &ChunkCache,
    sweep: &mut SweepContext,
) -> Result<Vec<u8>, FormatError> {
    let (chunk_dimensions, version, addr_opt) = match layout {
        DataLayout::Chunked {
            chunk_dimensions,
            version,
            btree_address,
            ..
        } => (chunk_dimensions, *version, *btree_address),
        _ => {
            return Err(FormatError::ChunkedReadError(
                "expected chunked layout".into(),
            ));
        }
    };

    let addr = addr_opt
        .ok_or_else(|| FormatError::ChunkedReadError("no address for chunked layout".into()))?;

    check_chunk_element_size(layout, datatype, offset_size)?;
    let elem_size = datatype.type_size() as usize;
    let (rank, chunk_dims) = chunk_geometry(chunk_dimensions, version, dataspace, elem_size)?;
    let ds_dims: Vec<usize> = dataspace
        .dimensions
        .iter()
        .map(|&d| to_usize(d))
        .collect::<Result<_, _>>()?;

    // The per-file cache is shared across datasets (and threads); every
    // lookup is keyed by this dataset's chunk-index address, so another
    // dataset's index or chunks are never used for this read.
    let chunks = cache.chunks_for(addr, rank, || {
        list_chunks_for_read(
            file_data,
            layout,
            dataspace,
            elem_size,
            pipeline,
            offset_size,
            length_size,
        )
        .map(|(chunks, _)| chunks)
    })?;

    // Assemble output
    let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?;
    if total_bytes == 0 {
        // Also keeps the stride products below in range: with a zero-sized
        // dimension the total is 0 even if other dimensions are huge.
        return Ok(Vec::new());
    }
    let mut output = alloc_output(total_bytes)?;

    let mut ds_strides = vec![1usize; rank];
    for i in (0..rank.saturating_sub(1)).rev() {
        ds_strides[i] = ds_strides[i + 1] * ds_dims[i + 1];
    }

    let mut chunk_strides = vec![1usize; rank];
    for i in (0..rank.saturating_sub(1)).rev() {
        chunk_strides[i] = chunk_strides[i + 1] * chunk_dims[i + 1];
    }

    let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?;

    for chunk_info in &chunks {
        let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();

        // Feed coordinate to sweep detector
        sweep.record(coord.clone(), rank);

        // Issue prefetch hint for predicted next chunks
        if !sweep.predicted_next.is_empty() {
            cache.prefetch_hint_in(addr, &sweep.predicted_next);
            cache.set_sweep_direction(sweep.direction);
        }

        // Try decompressed cache first
        let decompressed = if let Some(cached) = cache.get_decompressed_in(addr, &coord) {
            cached
        } else {
            // Decompress from file
            let c_addr = to_usize(chunk_info.address)?;
            let size = chunk_info.chunk_size as usize;
            ensure_len(file_data, c_addr, size)?;
            let raw_chunk = &file_data[c_addr..c_addr + size];
            let dec = if let Some(pl) = pipeline {
                decompress_chunk_exact(
                    raw_chunk,
                    pl,
                    chunk_total_bytes,
                    elem_size as u32,
                    chunk_info.filter_mask,
                    &coord,
                )?
            } else {
                raw_chunk.to_vec()
            };
            cache.put_decompressed_in(addr, coord, dec)
        };

        let chunk_offsets: Vec<usize> = chunk_info
            .offsets
            .iter()
            .take(rank)
            .map(|&o| to_usize(o))
            .collect::<Result<_, _>>()?;

        if rank == 0 {
            let copy_len = decompressed.len().min(output.len());
            output[..copy_len].copy_from_slice(&decompressed[..copy_len]);
        } else {
            copy_chunk_to_output(
                &decompressed,
                &mut output,
                &chunk_offsets,
                &chunk_dims,
                &ds_dims,
                &ds_strides,
                &chunk_strides,
                elem_size,
                rank,
            );
        }
    }

    Ok(output)
}

/// Read chunked data using pre-computed chunk layout for fast assembly.
///
/// This path builds a `ChunkIndex` and `ChunkLayout` on first access (cached
/// in the `ChunkCache`), then uses the pre-computed row-copy plan for assembly,
/// avoiding per-element N-D coordinate math on repeated reads.
#[cfg(feature = "std")]
#[allow(clippy::too_many_arguments)]
pub fn read_chunked_data_indexed(
    file_data: &[u8],
    layout: &DataLayout,
    dataspace: &Dataspace,
    datatype: &Datatype,
    pipeline: Option<&FilterPipeline>,
    offset_size: u8,
    length_size: u8,
    cache: &ChunkCache,
) -> Result<Vec<u8>, FormatError> {
    let (chunk_dimensions, version, addr_opt) = match layout {
        DataLayout::Chunked {
            chunk_dimensions,
            version,
            btree_address,
            ..
        } => (chunk_dimensions, *version, *btree_address),
        _ => {
            return Err(FormatError::ChunkedReadError(
                "expected chunked layout".into(),
            ));
        }
    };

    let addr = addr_opt
        .ok_or_else(|| FormatError::ChunkedReadError("no address for chunked layout".into()))?;

    check_chunk_element_size(layout, datatype, offset_size)?;
    let elem_size = datatype.type_size() as usize;
    let (rank, chunk_dims) = chunk_geometry(chunk_dimensions, version, dataspace, elem_size)?;
    let ds_dims: Vec<usize> = dataspace
        .dimensions
        .iter()
        .map(|&d| to_usize(d))
        .collect::<Result<_, _>>()?;

    // Chunk index and assembly plan for this dataset, built on first access
    // and kept per dataset (keyed by chunk-index address) in the shared cache.
    let plan = cache.chunk_layout_for(
        addr,
        rank,
        || {
            list_chunks_for_read(
                file_data,
                layout,
                dataspace,
                elem_size,
                pipeline,
                offset_size,
                length_size,
            )
            .map(|(chunks, _)| chunks)
        },
        &ds_dims,
        &chunk_dims,
        elem_size,
    )?;
    let chunk_total_bytes = plan.chunk_total_bytes;

    // Decompress chunks (using LRU cache where possible)
    let mut chunk_buffers: Vec<Arc<CacheAlignedBuffer>> = Vec::with_capacity(plan.mappings.len());
    for m in &plan.mappings {
        let (coord, file_offset, file_size, filter_mask) =
            (&m.coord, &m.file_offset, &m.file_size, &m.filter_mask);
        if let Some(cached) = cache.get_decompressed_in(addr, coord) {
            chunk_buffers.push(cached);
        } else {
            let c_addr = to_usize(*file_offset)?;
            let size = *file_size as usize;
            ensure_len(file_data, c_addr, size)?;
            let raw_chunk = &file_data[c_addr..c_addr + size];
            let decompressed = if let Some(pl) = pipeline {
                decompress_chunk_exact(
                    raw_chunk,
                    pl,
                    chunk_total_bytes,
                    elem_size as u32,
                    *filter_mask,
                    coord,
                )?
            } else {
                raw_chunk.to_vec()
            };
            let aligned = CacheAlignedBuffer::from_vec(decompressed);
            let arc = cache.put_decompressed_aligned_in(addr, coord.clone(), aligned);
            chunk_buffers.push(arc);
        }
    }

    // Assemble using pre-computed layout
    let mut output = vec![0u8; plan.output_bytes];
    let data_refs: Vec<&[u8]> = chunk_buffers.iter().map(|b| b.as_slice()).collect();
    plan.assemble(&data_refs, &mut output);

    Ok(output)
}

/// Copy chunk data into the output buffer at the correct N-D position.
#[cfg(any(feature = "std", test))]
#[allow(clippy::too_many_arguments)]
fn copy_chunk_to_output(
    chunk_data: &[u8],
    output: &mut [u8],
    chunk_offsets: &[usize],
    chunk_dims: &[usize],
    ds_dims: &[usize],
    ds_strides: &[usize],
    chunk_strides: &[usize],
    elem_size: usize,
    rank: usize,
) {
    // SAFETY: `output` is exclusively borrowed, so no other write can
    // overlap this one.
    unsafe {
        copy_chunk_into(
            chunk_data,
            &OutBuf::new(output),
            chunk_offsets,
            chunk_dims,
            ds_dims,
            ds_strides,
            chunk_strides,
            elem_size,
            rank,
        )
    }
}

/// [`copy_chunk_to_output`] through an [`OutBuf`]: only the bytes of the
/// chunk's region (clipped to the dataset) are written.
///
/// # Safety
///
/// No other thread may be writing an overlapping region of `output`.
#[allow(clippy::too_many_arguments)]
unsafe fn copy_chunk_into(
    chunk_data: &[u8],
    output: &OutBuf<'_>,
    chunk_offsets: &[usize],
    chunk_dims: &[usize],
    ds_dims: &[usize],
    ds_strides: &[usize],
    chunk_strides: &[usize],
    elem_size: usize,
    rank: usize,
) {
    // Row-copy approach: iterate over outer dimensions, memcpy the innermost
    // dimension in bulk. For 1-D data this is a single memcpy per chunk.
    // For N-D data this is one memcpy per "row" (innermost dim slice).

    if rank == 1 {
        // Fast path for 1-D: single contiguous copy per chunk
        let global_start = chunk_offsets[0];
        let copy_len = chunk_dims[0].min(ds_dims[0].saturating_sub(global_start));
        let (Some(src_bytes), Some(dst_start)) = (
            copy_len.checked_mul(elem_size),
            global_start.checked_mul(elem_size),
        ) else {
            return;
        };
        if src_bytes > 0
            && dst_start
                .checked_add(src_bytes)
                .is_some_and(|end| end <= output.len())
            && src_bytes <= chunk_data.len()
        {
            // SAFETY: in range (checked); exclusivity is the caller's contract.
            unsafe { output.write(dst_start, &chunk_data[..src_bytes]) };
        }
        return;
    }

    // General N-D: iterate over outer dimensions, copy innermost rows
    let inner_dim = rank - 1;
    let inner_chunk_len =
        chunk_dims[inner_dim].min(ds_dims[inner_dim].saturating_sub(chunk_offsets[inner_dim]));
    let Some(row_bytes) = inner_chunk_len.checked_mul(elem_size) else {
        return;
    };

    if row_bytes == 0 {
        return;
    }

    // Number of rows = product of all outer chunk dimensions
    let Some(outer_count) = chunk_dims[..inner_dim]
        .iter()
        .try_fold(1usize, |acc, &d| acc.checked_mul(d))
    else {
        return;
    };

    // Outer strides for iterating chunk-local coordinates
    let mut outer_strides = vec![1usize; inner_dim];
    for i in (0..inner_dim.saturating_sub(1)).rev() {
        let Some(stride) = outer_strides[i + 1].checked_mul(chunk_dims[i + 1]) else {
            return;
        };
        outer_strides[i] = stride;
    }

    for outer_idx in 0..outer_count {
        // Convert outer flat index to N-D chunk-local coords for dims 0..inner_dim
        let mut remaining = outer_idx;
        let mut ds_flat = 0usize;
        let mut src_flat = 0usize;
        let mut out_of_bounds = false;

        for d in 0..inner_dim {
            let coord_in_chunk = if inner_dim > 1 {
                remaining / outer_strides[d]
            } else {
                remaining
            };
            if inner_dim > 1 {
                remaining %= outer_strides[d];
            }

            let Some(global_coord) = chunk_offsets[d].checked_add(coord_in_chunk) else {
                out_of_bounds = true;
                break;
            };
            if global_coord >= ds_dims[d] {
                out_of_bounds = true;
                break;
            }
            let (Some(ds_term), Some(src_term)) = (
                global_coord.checked_mul(ds_strides[d]),
                coord_in_chunk.checked_mul(chunk_strides[d]),
            ) else {
                out_of_bounds = true;
                break;
            };
            let (Some(new_ds_flat), Some(new_src_flat)) =
                (ds_flat.checked_add(ds_term), src_flat.checked_add(src_term))
            else {
                out_of_bounds = true;
                break;
            };
            ds_flat = new_ds_flat;
            src_flat = new_src_flat;
        }

        if out_of_bounds {
            continue;
        }

        // Add innermost dimension offset
        let Some(inner_term) = chunk_offsets[inner_dim].checked_mul(ds_strides[inner_dim]) else {
            continue;
        };
        let Some(ds_flat) = ds_flat.checked_add(inner_term) else {
            continue;
        };

        let (Some(src_start), Some(dst_start)) = (
            src_flat.checked_mul(elem_size),
            ds_flat.checked_mul(elem_size),
        ) else {
            continue;
        };

        let fits = src_start
            .checked_add(row_bytes)
            .is_some_and(|end| end <= chunk_data.len())
            && dst_start
                .checked_add(row_bytes)
                .is_some_and(|end| end <= output.len());
        if fits {
            // SAFETY: in range (checked); exclusivity is the caller's contract.
            unsafe { output.write(dst_start, &chunk_data[src_start..src_start + row_bytes]) };
        }
    }
}

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

    fn simple_space(dimensions: Vec<u64>) -> Dataspace {
        Dataspace {
            space_type: crate::dataspace::DataspaceType::Simple,
            rank: dimensions.len() as u8,
            dimensions,
            max_dimensions: None,
        }
    }

    #[test]
    fn crafted_dimensions_are_errors_not_wraparound() {
        // 2^63 * 2 wraps to 0 with a plain product; 2^40 * 2^40 wraps too.
        for dims in [
            vec![1u64 << 63, 2],
            vec![1 << 40, 1 << 40],
            vec![u64::MAX, u64::MAX],
        ] {
            let space = simple_space(dims.clone());
            assert!(
                matches!(space.checked_num_elements(), Err(FormatError::Overflow(_))),
                "{dims:?}"
            );
            // The infallible accessor saturates instead of wrapping.
            assert_eq!(space.num_elements(), u64::MAX, "{dims:?}");
        }
        assert_eq!(simple_space(vec![3, 4]).checked_num_elements().unwrap(), 12);
        // A zero-sized dimension makes the whole product 0, not an overflow.
        assert_eq!(
            simple_space(vec![0, 1 << 40, 1 << 40])
                .checked_num_elements()
                .unwrap(),
            0
        );
    }

    #[test]
    fn byte_length_helpers_check_overflow() {
        assert_eq!(checked_byte_len(10, 8).unwrap(), 80);
        assert!(matches!(
            checked_byte_len(u64::MAX, 8),
            Err(FormatError::Overflow(_))
        ));
        assert_eq!(checked_chunk_byte_len(&[10, 10], 4).unwrap(), 400);
        assert!(matches!(
            checked_chunk_byte_len(&[usize::MAX, 2], 4),
            Err(FormatError::Overflow(_))
        ));
    }

    #[test]
    fn unallocatable_output_is_an_error_not_an_abort() {
        assert_eq!(alloc_output(16).unwrap(), vec![0u8; 16]);
        assert!(matches!(
            alloc_output(usize::MAX / 2),
            Err(FormatError::Overflow(_))
        ));
    }

    fn write_offset(buf: &mut Vec<u8>, val: u64, size: u8) {
        match size {
            4 => buf.extend_from_slice(&(val as u32).to_le_bytes()),
            8 => buf.extend_from_slice(&val.to_le_bytes()),
            _ => panic!("unsupported offset size in test"),
        }
    }

    /// Build a B-tree v1 type 1 leaf node with given chunk infos.
    fn build_chunk_btree_leaf(chunks: &[ChunkInfo], ndims: usize, offset_size: u8) -> Vec<u8> {
        build_chunk_btree_leaf_to(chunks, &vec![0; ndims], offset_size)
    }

    /// A leaf whose final key is the one libhdf5 writes past the last chunk:
    /// each coordinate of the last chunk plus its chunk dimension (the
    /// element size last).
    fn build_chunk_btree_leaf_dims(chunks: &[ChunkInfo], dims: &[u32], offset_size: u8) -> Vec<u8> {
        let last = &chunks.last().expect("a chunk").offsets;
        let end: Vec<u64> = dims
            .iter()
            .enumerate()
            .map(|(d, &c)| last.get(d).copied().unwrap_or(0) + u64::from(c))
            .collect();
        build_chunk_btree_leaf_to(chunks, &end, offset_size)
    }

    /// A leaf holding `chunks`, with final key `end`.
    fn build_chunk_btree_leaf_to(chunks: &[ChunkInfo], end: &[u64], offset_size: u8) -> Vec<u8> {
        let ndims = end.len();
        let _os = offset_size as usize;
        let entries_used = chunks.len() as u16;
        let mut buf = Vec::new();

        // Header
        buf.extend_from_slice(b"TREE");
        buf.push(1); // node_type = 1 (raw data chunks)
        buf.push(0); // node_level = 0 (leaf)
        buf.extend_from_slice(&entries_used.to_le_bytes());

        // Left/right sibling = undefined
        let undef: u64 = if offset_size == 4 {
            0xFFFFFFFF
        } else {
            0xFFFFFFFFFFFFFFFF
        };
        write_offset(&mut buf, undef, offset_size);
        write_offset(&mut buf, undef, offset_size);

        // Entries: key[i], child[i] pairs, then final key
        for chunk in chunks {
            // Key: chunk_size(4) + filter_mask(4) + ndims offsets
            buf.extend_from_slice(&chunk.chunk_size.to_le_bytes());
            buf.extend_from_slice(&chunk.filter_mask.to_le_bytes());
            for d in 0..ndims {
                let off = if d < chunk.offsets.len() {
                    chunk.offsets[d]
                } else {
                    0
                };
                // Key offsets are always 8 bytes (they are coordinates).
                write_offset(&mut buf, off, 8);
            }
            // Child: address
            write_offset(&mut buf, chunk.address, offset_size);
        }

        // Final key (its offsets must be on the chunk grid, as libhdf5
        // checks; 0 always is)
        buf.extend_from_slice(&0u32.to_le_bytes()); // chunk_size
        buf.extend_from_slice(&0u32.to_le_bytes()); // filter_mask
        for &e in end {
            write_offset(&mut buf, e, 8);
        }

        buf
    }

    // --- ChunkInfo collection tests ---

    #[test]
    fn checked_collection_refuses_keys_off_the_chunk_grid() {
        let chunk = |offsets: Vec<u64>, address| ChunkInfo {
            chunk_size: 80,
            filter_mask: 0,
            offsets,
            address,
        };
        let good = build_chunk_btree_leaf_dims(
            &[chunk(vec![0, 0], 0x100), chunk(vec![10, 0], 0x200)],
            &[10, 8],
            8,
        );
        assert_eq!(
            collect_chunk_info_checked(&good, 0, &[10, 8], 8, 8)
                .unwrap()
                .len(),
            2
        );
        // A chunk key off the grid.
        let bad =
            build_chunk_btree_leaf(&[chunk(vec![0, 0], 0x100), chunk(vec![7, 0], 0x200)], 2, 8);
        assert!(collect_chunk_info(&bad, 0, 2, 8, 8).is_ok());
        assert!(matches!(
            collect_chunk_info_checked(&bad, 0, &[10, 8], 8, 8),
            Err(FormatError::ChunkedReadError(m)) if m.starts_with("bad coordinate offset")
        ));
        // cve-2018-11205: the chunks all start at 0 in dimension 1, and only
        // the node's final key shows the chunk dimension is wrong.
        let mut two_d = build_chunk_btree_leaf(
            &[chunk(vec![0, 0, 0], 0x100), chunk(vec![10, 0, 0], 0x200)],
            3,
            8,
        );
        // Final key: (20, 20, 0), the end of a 20 x 20 dataset.
        let final_key = two_d.len() - 24;
        two_d[final_key..final_key + 8].copy_from_slice(&20u64.to_le_bytes());
        two_d[final_key + 8..final_key + 16].copy_from_slice(&20u64.to_le_bytes());
        assert!(collect_chunk_info_checked(&two_d, 0, &[10, 20, 4], 8, 8).is_ok());
        assert!(matches!(
            collect_chunk_info_checked(&two_d, 0, &[10, 32788, 4], 8, 8),
            Err(FormatError::ChunkedReadError(m)) if m.starts_with("bad coordinate offset [20, 20, 0]")
        ));
    }

    #[test]
    fn collect_two_chunks_from_leaf() {
        let ndims = 2; // rank+1 for 1D dataset
        let os: u8 = 8;

        let chunks = vec![
            ChunkInfo {
                chunk_size: 80,
                filter_mask: 0,
                offsets: vec![0, 0],
                address: 0x1000,
            },
            ChunkInfo {
                chunk_size: 80,
                filter_mask: 0,
                offsets: vec![10, 0],
                address: 0x2000,
            },
        ];

        let btree = build_chunk_btree_leaf(&chunks, ndims, os);
        let mut file_data = vec![0u8; 0x3000];
        file_data[..btree.len()].copy_from_slice(&btree);

        let result = collect_chunk_info(&file_data, 0, ndims, os, os).unwrap();
        assert_eq!(result.len(), 2);
        assert_eq!(result[0].address, 0x1000);
        assert_eq!(result[0].offsets, vec![0, 0]);
        assert_eq!(result[0].chunk_size, 80);
        assert_eq!(result[1].address, 0x2000);
        assert_eq!(result[1].offsets, vec![10, 0]);
    }

    #[test]
    fn collect_three_chunks() {
        let ndims = 2;
        let os: u8 = 8;

        let chunks = vec![
            ChunkInfo {
                chunk_size: 40,
                filter_mask: 0,
                offsets: vec![0, 0],
                address: 0x100,
            },
            ChunkInfo {
                chunk_size: 40,
                filter_mask: 0,
                offsets: vec![5, 0],
                address: 0x200,
            },
            ChunkInfo {
                chunk_size: 40,
                filter_mask: 0,
                offsets: vec![10, 0],
                address: 0x300,
            },
        ];

        let btree = build_chunk_btree_leaf(&chunks, ndims, os);
        let mut file_data = vec![0u8; 0x1000];
        file_data[..btree.len()].copy_from_slice(&btree);

        let result = collect_chunk_info(&file_data, 0, ndims, os, os).unwrap();
        assert_eq!(result.len(), 3);
        assert_eq!(result[0].address, 0x100);
        assert_eq!(result[1].address, 0x200);
        assert_eq!(result[2].address, 0x300);
    }

    #[test]
    fn collect_chunks_with_four_byte_addresses() {
        // Sibling and child addresses are 4 bytes; the key offsets stay 8.
        let ndims = 3;
        let os: u8 = 4;
        let chunks = vec![
            ChunkInfo {
                chunk_size: 80,
                filter_mask: 2,
                offsets: vec![0, 5, 0],
                address: 0x1000,
            },
            ChunkInfo {
                chunk_size: 96,
                filter_mask: 0,
                offsets: vec![8, 10, 0],
                address: 0x2000,
            },
        ];
        let btree = build_chunk_btree_leaf(&chunks, ndims, os);
        assert_eq!(btree.len(), 8 + 2 * 4 + 2 * (8 + 3 * 8 + 4) + (8 + 3 * 8));
        let result = collect_chunk_info(&btree, 0, ndims, os, os).unwrap();
        assert_eq!(result.len(), 2);
        for (got, want) in result.iter().zip(&chunks) {
            assert_eq!(got.offsets, want.offsets);
            assert_eq!(got.address, want.address);
            assert_eq!(got.chunk_size, want.chunk_size);
            assert_eq!(got.filter_mask, want.filter_mask);
        }
    }

    #[test]
    fn collect_empty_btree() {
        let ndims = 2;
        let os: u8 = 8;
        let btree = build_chunk_btree_leaf(&[], ndims, os);
        let mut file_data = vec![0u8; 0x1000];
        file_data[..btree.len()].copy_from_slice(&btree);

        let result = collect_chunk_info(&file_data, 0, ndims, os, os).unwrap();
        assert_eq!(result.len(), 0);
    }

    // --- Chunked read tests (synthetic) ---

    use crate::dataspace::{Dataspace, DataspaceType};
    use crate::datatype::{Datatype, DatatypeByteOrder};

    #[test]
    fn chunk_geometry_matches_libhdf5_open_checks() {
        let space = |dims: &[u64]| Dataspace {
            space_type: DataspaceType::Simple,
            rank: dims.len() as u8,
            dimensions: dims.to_vec(),
            max_dimensions: None,
        };
        for v in [3, 4] {
            assert_eq!(
                chunk_geometry(&[4, 5, 8], v, &space(&[10, 10]), 8).unwrap(),
                (2, vec![4, 5])
            );
            // Rank mismatch.
            assert!(matches!(
                chunk_geometry(&[4, 8], v, &space(&[10, 10]), 8),
                Err(FormatError::InvalidChunkDimensions(m)) if m.contains("doesn't match")
            ));
            // Zero dimension (a layout built in memory, bypassing the parser).
            assert!(matches!(
                chunk_geometry(&[4, 0, 8], v, &space(&[10, 10]), 8),
                Err(FormatError::InvalidChunkDimensions(m)) if m.contains("must be > 0")
            ));
            assert!(chunk_geometry(&[0xFFFF_FFFF, 1], v, &space(&[10]), 1).is_ok());
        }
        // With a v1 B-tree index (layout version 3) the largest chunk is
        // 4 GiB - 1 bytes: 0x80000000 x 4-byte elements (8 GiB) is refused.
        // These dims used to hang the reader.
        assert!(matches!(
            chunk_geometry(&[0x8000_0000, 4], 3, &space(&[10]), 4),
            Err(FormatError::InvalidChunkDimensions(m)) if m.contains("4GB with v1 b-tree")
        ));
        assert!(matches!(
            chunk_geometry(&[0xFFFF_FFFF, 0xFFFF_FFFF, 1], 3, &space(&[10, 10]), 1),
            Err(FormatError::InvalidChunkDimensions(m)) if m.contains("4GB with v1 b-tree")
        ));
        // The other chunk indexes (layout version 4, and 5, which is read as
        // 4) allow chunks of 4 GiB and more; HDF5 2.0 writes them.
        assert_eq!(
            chunk_geometry(&[0x2000_0001, 8], 4, &space(&[10]), 8).unwrap(),
            (1, vec![0x2000_0001])
        );
        assert!(chunk_geometry(&[0xFFFF_FFFF, 0xFFFF_FFFF, 1], 4, &space(&[10, 10]), 1).is_ok());
    }

    fn make_f64_type() -> Datatype {
        Datatype::FloatingPoint {
            size: 8,
            byte_order: DatatypeByteOrder::LittleEndian,
            bit_offset: 0,
            bit_precision: 64,
            exponent_location: 52,
            exponent_size: 11,
            mantissa_location: 0,
            mantissa_size: 52,
            exponent_bias: 1023,
        }
    }

    fn make_f32_type() -> Datatype {
        Datatype::FloatingPoint {
            size: 4,
            byte_order: DatatypeByteOrder::LittleEndian,
            bit_offset: 0,
            bit_precision: 32,
            exponent_location: 23,
            exponent_size: 8,
            mantissa_location: 0,
            mantissa_size: 23,
            exponent_bias: 127,
        }
    }

    /// Build a synthetic file with a B-tree and chunk data for a 1D uncompressed dataset.
    fn build_1d_chunked_file(
        values: &[f64],
        chunk_size_elems: usize,
    ) -> (Vec<u8>, DataLayout, Dataspace) {
        let os: u8 = 8;
        let elem_size = 8usize;
        let total = values.len();

        // Place chunk data starting at offset 0x2000
        let mut file_data = vec![0u8; 0x10000];
        let mut chunk_infos = Vec::new();
        let mut data_offset = 0x2000usize;

        let mut start = 0;
        while start < total {
            let end = (start + chunk_size_elems).min(total);
            let chunk_bytes = chunk_size_elems * elem_size; // full chunk allocation

            // Write chunk data (full chunk size, padding with zeros)
            for (i, value) in values.iter().enumerate().take(end).skip(start) {
                let byte_offset = data_offset + (i - start) * elem_size;
                file_data[byte_offset..byte_offset + 8].copy_from_slice(&value.to_le_bytes());
            }

            chunk_infos.push(ChunkInfo {
                chunk_size: chunk_bytes as u32,
                filter_mask: 0,
                offsets: vec![start as u64, 0],
                address: data_offset as u64,
            });

            data_offset += chunk_bytes;
            start += chunk_size_elems;
        }

        // Build B-tree at offset 0x100
        let dims = [chunk_size_elems as u32, elem_size as u32];
        let btree = build_chunk_btree_leaf_dims(&chunk_infos, &dims, os);
        let btree_addr = 0x100usize;
        file_data[btree_addr..btree_addr + btree.len()].copy_from_slice(&btree);

        let layout = DataLayout::Chunked {
            chunk_dimensions: dims.to_vec(),
            btree_address: Some(btree_addr as u64),
            version: 3,
            chunk_index_type: None,
            single_chunk_filtered_size: None,
            single_chunk_filter_mask: None,
            dont_filter_partial_edge_chunks: false,
        };

        let dataspace = Dataspace {
            space_type: DataspaceType::Simple,
            rank: 1,
            dimensions: vec![total as u64],
            max_dimensions: None,
        };

        (file_data, layout, dataspace)
    }

    #[test]
    fn read_chunked_data_rejects_zero_dim_chunk_layout() {
        // Found by fuzzing: chunk_dimensions.len() == 0 caused `ndims - 1` to
        // underflow. A malformed/degenerate chunked layout must error cleanly.
        let layout = DataLayout::Chunked {
            chunk_dimensions: vec![],
            btree_address: Some(0),
            version: 3,
            chunk_index_type: None,
            single_chunk_filtered_size: None,
            single_chunk_filter_mask: None,
            dont_filter_partial_edge_chunks: false,
        };
        let dataspace = Dataspace {
            space_type: DataspaceType::Simple,
            rank: 1,
            dimensions: vec![10],
            max_dimensions: None,
        };
        let datatype = make_f64_type();
        let file_data = vec![0u8; 64];
        let result = read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8);
        assert!(
            matches!(result, Err(FormatError::InvalidChunkDimensions(_))),
            "expected a clean InvalidChunkDimensions, got {result:?}"
        );
    }

    /// A 2-D `f32` dataset of `rows x cols` in `cr x cc` chunks, deflated
    /// behind shuffle: (file bytes, chunk list, pipeline, expected output).
    #[cfg(feature = "deflate")]
    fn deflated_grid(
        rows: usize,
        cols: usize,
        cr: usize,
        cc: usize,
    ) -> (Vec<u8>, Vec<ChunkInfo>, FilterPipeline, Vec<u8>) {
        use crate::filter_pipeline::{FILTER_DEFLATE, FILTER_SHUFFLE, FilterDescription};
        let pipeline = FilterPipeline {
            version: 2,
            filters: vec![
                FilterDescription {
                    filter_id: FILTER_SHUFFLE,
                    name: None,
                    flags: 0,
                    client_data: vec![4],
                },
                FilterDescription {
                    filter_id: FILTER_DEFLATE,
                    name: None,
                    flags: 0,
                    client_data: vec![1],
                },
            ],
        };
        let value = |r: usize, c: usize| (r * cols + c) as f32 * 0.5;
        let expected: Vec<u8> = (0..rows * cols)
            .flat_map(|i| value(i / cols, i % cols).to_le_bytes())
            .collect();
        let (mut file, mut chunks) = (vec![0u8; 8], Vec::new());
        for r0 in (0..rows).step_by(cr) {
            for c0 in (0..cols).step_by(cc) {
                // Full-size chunks; the edge padding is zeros.
                let chunk: Vec<u8> = (0..cr * cc)
                    .flat_map(|i| {
                        let (r, c) = (r0 + i / cc, c0 + i % cc);
                        let v = if r < rows && c < cols {
                            value(r, c)
                        } else {
                            0.0
                        };
                        v.to_le_bytes()
                    })
                    .collect();
                let stored = crate::filters::compress_chunk(&chunk, &pipeline, 4).unwrap();
                chunks.push(ChunkInfo {
                    chunk_size: stored.len() as u32,
                    filter_mask: 0,
                    offsets: vec![r0 as u64, c0 as u64, 0],
                    address: file.len() as u64,
                });
                file.extend(stored);
            }
        }
        (file, chunks, pipeline, expected)
    }

    #[cfg(feature = "deflate")]
    fn fill(
        file: &[u8],
        chunks: &[ChunkInfo],
        pipeline: &FilterPipeline,
        dims: [usize; 2],
        chunk: [usize; 2],
    ) -> Result<Vec<u8>, FormatError> {
        let placer = ChunkPlacer::new(&chunk, &dims, 4);
        let mut out = vec![0u8; dims[0] * dims[1] * 4];
        fill_from_chunks(
            file,
            chunks,
            Some(pipeline),
            &placer,
            chunk[0] * chunk[1] * 4,
            None,
            &mut out,
        )?;
        Ok(out)
    }

    /// Chunks decoded straight into the output (in parallel with the
    /// `parallel` feature, partial edge chunks included) land where the
    /// row-by-row reference puts them.
    #[test]
    #[cfg(feature = "deflate")]
    fn chunks_fill_the_output_in_place() {
        for (dims, chunk) in [
            ([100, 70], [16, 32]),
            ([64, 64], [16, 16]),
            ([5, 3], [8, 8]),
        ] {
            let (file, chunks, pipeline, expected) =
                deflated_grid(dims[0], dims[1], chunk[0], chunk[1]);
            for _ in 0..4 {
                assert_eq!(
                    fill(&file, &chunks, &pipeline, dims, chunk).unwrap(),
                    expected
                );
            }
            // Any chunk order gives the same output.
            let mut reversed = chunks.clone();
            reversed.reverse();
            assert_eq!(
                fill(&file, &reversed, &pipeline, dims, chunk).unwrap(),
                expected
            );
        }
    }

    /// Of several corrupt chunks, the error names the first in chunk order,
    /// however the chunks were shared out between threads.
    #[test]
    #[cfg(feature = "deflate")]
    fn first_corrupt_chunk_is_the_error() {
        let (dims, chunk) = ([128, 64], [8, 16]);
        let (mut file, mut chunks, pipeline, _) =
            deflated_grid(dims[0], dims[1], chunk[0], chunk[1]);
        // Valid streams that decode short: an error naming the chunk.
        let short = crate::filters::compress_chunk(&[1u8; 64], &pipeline, 4).unwrap();
        for bad in [5usize, 11, 40] {
            chunks[bad].address = file.len() as u64;
            chunks[bad].chunk_size = short.len() as u32;
            file.extend_from_slice(&short);
        }
        for _ in 0..20 {
            let err = fill(&file, &chunks, &pipeline, dims, chunk).unwrap_err();
            let want = format!("{:?}", chunks[5].offsets);
            assert!(err.to_string().contains(&want), "{err} (want {want})");
        }
    }

    /// Only chunks on the grid, each at a distinct place, may be placed
    /// concurrently; anything else is read one chunk at a time.
    #[test]
    #[cfg(feature = "parallel")]
    fn regions_disjoint_only_for_distinct_grid_chunks() {
        let placer = ChunkPlacer::new(&[10, 10], &[25, 30], 4);
        let chunk = |r: u64, c: u64| ChunkInfo {
            chunk_size: 0,
            filter_mask: 0,
            offsets: vec![r, c, 0],
            address: 0,
        };
        let grid: Vec<ChunkInfo> = (0..3)
            .flat_map(|r| (0..3).map(move |c| chunk(r * 10, c * 10)))
            .collect();
        assert!(placer.regions_disjoint(&grid));
        // Chunks wholly outside the dataset write nothing.
        let mut outside = grid.clone();
        outside.push(chunk(30, 0));
        outside.push(chunk(30, 0));
        assert!(placer.regions_disjoint(&outside));
        let mut duplicate = grid.clone();
        duplicate.push(chunk(10, 20));
        assert!(!placer.regions_disjoint(&duplicate));
        let mut off_grid = grid.clone();
        off_grid[4] = chunk(15, 10);
        assert!(!placer.regions_disjoint(&off_grid));
        let mut short = grid;
        short[0].offsets.truncate(1);
        assert!(!placer.regions_disjoint(&short));
    }

    /// A duplicated chunk in a corrupt index is not placed from two threads
    /// at once: the read goes one chunk at a time, as before.
    #[test]
    #[cfg(feature = "deflate")]
    fn duplicate_chunks_are_read_in_order() {
        let (dims, chunk) = ([64, 64], [16, 16]);
        let (file, mut chunks, pipeline, expected) =
            deflated_grid(dims[0], dims[1], chunk[0], chunk[1]);
        // A second entry at chunk 3's place, holding chunk 7's data: the
        // later entry wins, as a sequential read has it.
        let mut dup = chunks[7].clone();
        dup.offsets = chunks[3].offsets.clone();
        chunks.push(dup);
        let got = fill(&file, &chunks, &pipeline, dims, chunk).unwrap();
        let chunk7_in_3: Vec<u8> = {
            let mut e = expected.clone();
            let (r3, c3) = (chunks[3].offsets[0] as usize, chunks[3].offsets[1] as usize);
            let (r7, c7) = (chunks[7].offsets[0] as usize, chunks[7].offsets[1] as usize);
            for r in 0..16 {
                let src = ((r7 + r) * 64 + c7) * 4;
                let dst = ((r3 + r) * 64 + c3) * 4;
                e[dst..dst + 64].copy_from_slice(&expected[src..src + 64]);
            }
            e
        };
        assert_eq!(got, chunk7_in_3);
    }

    #[test]
    fn copy_chunk_to_output_1d_rejects_overflowing_offset_without_panicking() {
        // Found by fuzzing: `global_start * elem_size` overflowed for a
        // crafted large chunk offset.
        let chunk_data = vec![1u8; 16];
        let mut output = vec![0u8; 16];
        let chunk_offsets = [usize::MAX - 1];
        let chunk_dims = [1usize];
        let ds_dims = [usize::MAX];
        let ds_strides = [1usize];
        let chunk_strides = [1usize];
        copy_chunk_to_output(
            &chunk_data,
            &mut output,
            &chunk_offsets,
            &chunk_dims,
            &ds_dims,
            &ds_strides,
            &chunk_strides,
            8,
            1,
        );
        // No panic; the out-of-range write was skipped, output left untouched.
        assert_eq!(output, vec![0u8; 16]);
    }

    #[test]
    fn copy_chunk_to_output_nd_rejects_overflowing_offset_without_panicking() {
        let chunk_data = vec![1u8; 16];
        let mut output = vec![0u8; 16];
        let chunk_offsets = [usize::MAX - 1, 0];
        let chunk_dims = [1usize, 1usize];
        let ds_dims = [usize::MAX, usize::MAX];
        let ds_strides = [1usize, 1usize];
        let chunk_strides = [1usize, 1usize];
        copy_chunk_to_output(
            &chunk_data,
            &mut output,
            &chunk_offsets,
            &chunk_dims,
            &ds_dims,
            &ds_strides,
            &chunk_strides,
            8,
            2,
        );
        assert_eq!(output, vec![0u8; 16]);
    }

    #[test]
    fn read_1d_two_chunks_no_compression() {
        let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
        let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
        let datatype = make_f64_type();

        let raw =
            read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8).unwrap();
        assert_eq!(raw.len(), 20 * 8);

        // Verify values
        for i in 0..20 {
            let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
            assert_eq!(val, i as f64);
        }
    }

    #[test]
    fn read_1d_three_chunks_partial_last() {
        // 25 elements, chunk size 10 => 3 chunks, last has only 5 valid
        let values: Vec<f64> = (0..25).map(|i| i as f64).collect();
        let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
        let datatype = make_f64_type();

        let raw =
            read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8).unwrap();
        assert_eq!(raw.len(), 25 * 8);

        for i in 0..25 {
            let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
            assert_eq!(val, i as f64, "mismatch at index {i}");
        }
    }

    #[cfg(feature = "deflate")]
    #[test]
    fn read_1d_two_chunks_with_deflate() {
        use crate::filter_pipeline::{FILTER_DEFLATE, FilterDescription, FilterPipeline};
        use crate::filters::compress_chunk;

        let os: u8 = 8;
        let elem_size = 8usize;
        let chunk_elems = 10usize;
        let total = 20usize;

        let pipeline = FilterPipeline {
            version: 2,
            filters: vec![FilterDescription {
                filter_id: FILTER_DEFLATE,
                name: None,
                flags: 0,
                client_data: vec![6],
            }],
        };

        let values: Vec<f64> = (0..total).map(|i| i as f64).collect();
        let mut file_data = vec![0u8; 0x10000];
        let mut chunk_infos = Vec::new();
        let mut data_offset = 0x2000usize;

        for chunk_idx in 0..2 {
            let start = chunk_idx * chunk_elems;
            let mut chunk_bytes = Vec::new();
            for value in values.iter().skip(start).take(chunk_elems) {
                chunk_bytes.extend_from_slice(&value.to_le_bytes());
            }
            let compressed = compress_chunk(&chunk_bytes, &pipeline, elem_size as u32).unwrap();

            file_data[data_offset..data_offset + compressed.len()].copy_from_slice(&compressed);

            chunk_infos.push(ChunkInfo {
                chunk_size: compressed.len() as u32,
                filter_mask: 0,
                offsets: vec![start as u64, 0],
                address: data_offset as u64,
            });

            data_offset += compressed.len() + 16; // some padding
        }

        let dims = [chunk_elems as u32, elem_size as u32];
        let btree = build_chunk_btree_leaf_dims(&chunk_infos, &dims, os);
        let btree_addr = 0x100usize;
        file_data[btree_addr..btree_addr + btree.len()].copy_from_slice(&btree);

        let layout = DataLayout::Chunked {
            chunk_dimensions: dims.to_vec(),
            btree_address: Some(btree_addr as u64),
            version: 3,
            chunk_index_type: None,
            single_chunk_filtered_size: None,
            single_chunk_filter_mask: None,
            dont_filter_partial_edge_chunks: false,
        };
        let dataspace = Dataspace {
            space_type: DataspaceType::Simple,
            rank: 1,
            dimensions: vec![total as u64],
            max_dimensions: None,
        };
        let datatype = make_f64_type();

        let raw = read_chunked_data(
            &file_data,
            &layout,
            &dataspace,
            &datatype,
            Some(&pipeline),
            8,
            8,
        )
        .unwrap();

        for i in 0..total {
            let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
            assert_eq!(val, i as f64, "mismatch at index {i}");
        }
    }

    #[test]
    fn read_2d_four_chunks() {
        // 4x6 dataset with chunk size 2x3 => 4 chunks
        let os: u8 = 8;
        let elem_size = 4usize; // f32
        let ds_dims = [4usize, 6];
        let chunk_dims = [2usize, 3];

        let values: Vec<f32> = (0..24).map(|i| i as f32).collect();
        let mut file_data = vec![0u8; 0x10000];
        let mut chunk_infos = Vec::new();
        let mut data_offset = 0x2000usize;

        // Generate chunks: (0,0), (0,3), (2,0), (2,3)
        for row_start in (0..ds_dims[0]).step_by(chunk_dims[0]) {
            for col_start in (0..ds_dims[1]).step_by(chunk_dims[1]) {
                let mut chunk_bytes = Vec::new();
                for r in 0..chunk_dims[0] {
                    for c in 0..chunk_dims[1] {
                        let gr = row_start + r;
                        let gc = col_start + c;
                        let val = if gr < ds_dims[0] && gc < ds_dims[1] {
                            values[gr * ds_dims[1] + gc]
                        } else {
                            0.0
                        };
                        chunk_bytes.extend_from_slice(&val.to_le_bytes());
                    }
                }

                let chunk_size = chunk_bytes.len();
                file_data[data_offset..data_offset + chunk_size].copy_from_slice(&chunk_bytes);

                chunk_infos.push(ChunkInfo {
                    chunk_size: chunk_size as u32,
                    filter_mask: 0,
                    offsets: vec![row_start as u64, col_start as u64, 0],
                    address: data_offset as u64,
                });

                data_offset += chunk_size + 8;
            }
        }

        let dims = [chunk_dims[0] as u32, chunk_dims[1] as u32, elem_size as u32];
        let btree = build_chunk_btree_leaf_dims(&chunk_infos, &dims, os);
        let btree_addr = 0x100usize;
        file_data[btree_addr..btree_addr + btree.len()].copy_from_slice(&btree);

        let layout = DataLayout::Chunked {
            chunk_dimensions: dims.to_vec(),
            btree_address: Some(btree_addr as u64),
            version: 3,
            chunk_index_type: None,
            single_chunk_filtered_size: None,
            single_chunk_filter_mask: None,
            dont_filter_partial_edge_chunks: false,
        };
        let dataspace = Dataspace {
            space_type: DataspaceType::Simple,
            rank: 2,
            dimensions: vec![ds_dims[0] as u64, ds_dims[1] as u64],
            max_dimensions: None,
        };
        let datatype = make_f32_type();

        let raw =
            read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8).unwrap();
        assert_eq!(raw.len(), 24 * 4);

        for i in 0..24 {
            let val = f32::from_le_bytes(raw[i * 4..(i + 1) * 4].try_into().unwrap());
            assert_eq!(val, i as f32, "mismatch at element {i}");
        }
    }

    #[test]
    fn wrong_node_type_error() {
        // Build a type-0 B-tree and try to collect chunk info
        let mut buf = Vec::new();
        buf.extend_from_slice(b"TREE");
        buf.push(0); // type 0, not 1
        buf.push(0);
        buf.extend_from_slice(&0u16.to_le_bytes());
        buf.extend_from_slice(&[0xFF; 16]); // siblings
        // final key
        buf.extend_from_slice(&[0u8; 24]);

        let mut file_data = vec![0u8; 512];
        file_data[..buf.len()].copy_from_slice(&buf);

        let err = collect_chunk_info(&file_data, 0, 2, 8, 8).unwrap_err();
        assert_eq!(err, FormatError::InvalidBTreeNodeType(0));
    }

    #[test]
    fn collect_chunk_info_rejects_near_usize_max_offset() {
        let file_data = vec![0u8; 64];
        let result = collect_chunk_info(&file_data, u64::MAX - 4, 2, 8, 8);
        assert!(
            matches!(result, Err(FormatError::UnexpectedEof { .. })),
            "expected a clean UnexpectedEof, got {result:?}"
        );
    }

    #[test]
    fn collect_chunk_info_rejects_self_referencing_internal_node() {
        // A type-1 internal node (level 1) whose single child address points
        // back to itself: an infinite-recursion / cyclic B-tree attack.
        let ndims = 2;
        let os: u8 = 8;
        let mut buf = Vec::new();
        buf.extend_from_slice(b"TREE");
        buf.push(1); // node_type = 1 (raw data chunks)
        buf.push(1); // node_level = 1 (internal)
        buf.extend_from_slice(&1u16.to_le_bytes()); // entries_used = 1
        write_offset(&mut buf, u64::MAX, os); // left sibling undefined
        write_offset(&mut buf, u64::MAX, os); // right sibling undefined
        // key[0]: chunk_size(4) + filter_mask(4) + ndims offsets
        buf.extend_from_slice(&0u32.to_le_bytes());
        buf.extend_from_slice(&0u32.to_le_bytes());
        for _ in 0..ndims {
            write_offset(&mut buf, 0, os);
        }
        // child[0]: points back to offset 0 (this same node) — cyclic.
        write_offset(&mut buf, 0, os);
        // final key
        buf.extend_from_slice(&0u32.to_le_bytes());
        buf.extend_from_slice(&0u32.to_le_bytes());
        for _ in 0..ndims {
            write_offset(&mut buf, u64::MAX, os);
        }

        let mut file_data = vec![0u8; 256];
        file_data[..buf.len()].copy_from_slice(&buf);

        let result = collect_chunk_info(&file_data, 0, ndims, os, os);
        assert!(
            matches!(result, Err(FormatError::NestingDepthExceeded)),
            "expected a clean NestingDepthExceeded, got {result:?}"
        );
    }

    // --- Implicit chunk generation tests ---

    #[test]
    fn implicit_chunks_1d_five_chunks() {
        let chunks = generate_implicit_chunks(
            0x1000,
            &[100],
            &[20],
            8, // f64
        );
        assert_eq!(chunks.len(), 5);
        let chunk_byte_size = 20 * 8;
        for (i, c) in chunks.iter().enumerate() {
            assert_eq!(c.address, 0x1000 + i as u64 * chunk_byte_size as u64);
            assert_eq!(c.offsets, vec![i as u64 * 20]);
            assert_eq!(c.filter_mask, 0);
            assert_eq!(c.chunk_size, chunk_byte_size as u32);
        }
    }

    #[test]
    fn implicit_chunks_2d() {
        // 10x6 dataset, 4x3 chunks => ceil(10/4)=3, ceil(6/3)=2 => 6 chunks
        let chunks = generate_implicit_chunks(
            0x2000,
            &[10, 6],
            &[4, 3],
            4, // f32
        );
        assert_eq!(chunks.len(), 6);
        let chunk_byte_size = 4 * 3 * 4;
        // Row-major: (0,0), (0,3), (4,0), (4,3), (8,0), (8,3)
        assert_eq!(chunks[0].offsets, vec![0, 0]);
        assert_eq!(chunks[1].offsets, vec![0, 3]);
        assert_eq!(chunks[2].offsets, vec![4, 0]);
        assert_eq!(chunks[3].offsets, vec![4, 3]);
        assert_eq!(chunks[4].offsets, vec![8, 0]);
        assert_eq!(chunks[5].offsets, vec![8, 3]);
        for (i, c) in chunks.iter().enumerate() {
            assert_eq!(c.address, 0x2000 + i as u64 * chunk_byte_size as u64);
        }
    }

    #[test]
    fn implicit_chunks_partial_last() {
        // 25 elements, chunk size 10 => 3 chunks (last partial)
        let chunks = generate_implicit_chunks(0x0, &[25], &[10], 8);
        assert_eq!(chunks.len(), 3);
        assert_eq!(chunks[0].offsets, vec![0]);
        assert_eq!(chunks[1].offsets, vec![10]);
        assert_eq!(chunks[2].offsets, vec![20]);
    }

    // --- V4 single chunk synthetic test ---

    #[test]
    fn read_v4_single_chunk_synthetic() {
        // Build a synthetic v4 single chunk dataset (no filters)
        let values: Vec<f64> = vec![10.0, 20.0, 30.0];
        let elem_size = 8usize;
        let chunk_elems = 3usize;

        let mut file_data = vec![0u8; 0x2000];
        let data_addr = 0x1000usize;
        for (i, &v) in values.iter().enumerate() {
            file_data[data_addr + i * elem_size..data_addr + (i + 1) * elem_size]
                .copy_from_slice(&v.to_le_bytes());
        }

        let layout = DataLayout::Chunked {
            chunk_dimensions: vec![chunk_elems as u32, elem_size as u32],
            btree_address: Some(data_addr as u64),
            version: 4,
            chunk_index_type: Some(1),
            single_chunk_filtered_size: None,
            single_chunk_filter_mask: None,
            dont_filter_partial_edge_chunks: false,
        };
        let dataspace = Dataspace {
            space_type: DataspaceType::Simple,
            rank: 1,
            dimensions: vec![3],
            max_dimensions: None,
        };
        let datatype = make_f64_type();

        let raw =
            read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8).unwrap();
        assert_eq!(raw.len(), 24);
        for i in 0..3 {
            let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
            assert_eq!(val, values[i]);
        }
    }

    // --- Cached read tests ---

    use crate::chunk_cache::ChunkCache;

    /// A chunk stored unfiltered in a filtered dataset (every filter-mask
    /// bit set) must still hold the whole chunk. The cached reader placed a
    /// short one and read its missing rows as zeros; both readers refuse it.
    #[test]
    fn short_unfiltered_chunk_of_a_filtered_dataset_is_an_error() {
        use crate::filter_pipeline::{FILTER_SHUFFLE, FilterDescription};
        let mut file_data = vec![0u8; 0x2000];
        let mut infos = Vec::new();
        for k in 0..3u64 {
            let address = 0x1000 + k as usize * 80;
            for i in 0..10u64 {
                let at = address + i as usize * 8;
                file_data[at..at + 8].copy_from_slice(&((k * 10 + i) as f64).to_le_bytes());
            }
            infos.push(ChunkInfo {
                chunk_size: if k == 1 { 40 } else { 80 },
                filter_mask: 1,
                offsets: vec![k * 10, 0],
                address: address as u64,
            });
        }
        let btree = build_chunk_btree_leaf(&infos, 2, 8);
        file_data[0x100..0x100 + btree.len()].copy_from_slice(&btree);
        let layout = DataLayout::Chunked {
            chunk_dimensions: vec![10, 8],
            btree_address: Some(0x100),
            version: 3,
            chunk_index_type: None,
            single_chunk_filtered_size: None,
            single_chunk_filter_mask: None,
            dont_filter_partial_edge_chunks: false,
        };
        let dataspace = simple_space(vec![30]);
        let pipeline = FilterPipeline {
            version: 2,
            filters: vec![FilterDescription {
                filter_id: FILTER_SHUFFLE,
                name: None,
                flags: 0,
                client_data: vec![8],
            }],
        };
        let dt = make_f64_type();
        let cache = ChunkCache::new();
        for result in [
            read_chunked_data(&file_data, &layout, &dataspace, &dt, Some(&pipeline), 8, 8),
            read_chunked_data_cached(
                &file_data,
                &layout,
                &dataspace,
                &dt,
                Some(&pipeline),
                8,
                8,
                &cache,
            ),
        ] {
            let err = result.unwrap_err().to_string();
            assert!(err.contains("[10, 0]") && err.contains("40"), "{err}");
        }
    }

    #[test]
    fn cached_read_populates_index_and_returns_correct_data() {
        let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
        let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
        let datatype = make_f64_type();
        let cache = ChunkCache::new();

        assert_eq!(cache.indexed_dataset_count(), 0);
        let raw = read_chunked_data_cached(
            &file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
        )
        .unwrap();
        assert_eq!(cache.indexed_dataset_count(), 1);
        assert_eq!(raw.len(), 20 * 8);
        for i in 0..20 {
            let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
            assert_eq!(val, i as f64);
        }
    }

    #[test]
    fn cached_read_second_call_reuses_the_index() {
        let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
        let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
        let datatype = make_f64_type();
        let cache = ChunkCache::new();

        // First read — populates the chunk index. These chunks are stored
        // unfiltered, so they are copied straight from the file bytes and the
        // decompressed-chunk cache is (deliberately) not involved.
        let raw1 = read_chunked_data_cached(
            &file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
        )
        .unwrap();
        assert_eq!(cache.indexed_dataset_count(), 1);
        assert_eq!(cache.cached_chunk_count(), 0);

        // Second read — reuses the cached index
        let raw2 = read_chunked_data_cached(
            &file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
        )
        .unwrap();
        assert_eq!(raw1, raw2);
        assert_eq!(cache.indexed_dataset_count(), 1);
    }

    #[test]
    fn cached_read_with_partial_last_chunk() {
        let values: Vec<f64> = (0..25).map(|i| i as f64).collect();
        let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
        let datatype = make_f64_type();
        let cache = ChunkCache::new();

        let raw = read_chunked_data_cached(
            &file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
        )
        .unwrap();
        assert_eq!(raw.len(), 25 * 8);
        for i in 0..25 {
            let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
            assert_eq!(val, i as f64, "mismatch at index {i}");
        }
    }

    // --- Sweep-aware read tests ---

    #[test]
    fn sweep_read_returns_correct_data() {
        let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
        let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
        let datatype = make_f64_type();
        let cache = ChunkCache::new();
        let mut sweep = SweepContext::with_defaults();

        let raw = read_chunked_data_sweep(
            &file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache, &mut sweep,
        )
        .unwrap();
        assert_eq!(raw.len(), 20 * 8);
        for i in 0..20 {
            let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
            assert_eq!(val, i as f64);
        }
    }

    #[test]
    fn sweep_read_populates_sweep_context() {
        let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
        let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
        let datatype = make_f64_type();
        let cache = ChunkCache::new();
        let mut sweep = SweepContext::with_defaults();

        read_chunked_data_sweep(
            &file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache, &mut sweep,
        )
        .unwrap();

        // After reading 2 chunks (offsets [0] and [10]), history should be populated
        assert!(!sweep.history.is_empty());
    }

    #[test]
    fn sweep_context_unit_test() {
        let mut ctx = SweepContext::with_defaults();
        ctx.record(vec![0, 0], 2);
        ctx.record(vec![0, 10], 2);
        ctx.record(vec![0, 20], 2);
        assert_eq!(ctx.direction, "row_major");
        assert!(!ctx.predicted_next.is_empty());
        assert_eq!(ctx.predicted_next[0], vec![0, 30]);
    }

    #[test]
    fn sweep_context_random() {
        let mut ctx = SweepContext::with_defaults();
        ctx.record(vec![0, 0], 2);
        ctx.record(vec![30, 20], 2);
        ctx.record(vec![10, 0], 2);
        assert_eq!(ctx.direction, "random");
        assert!(ctx.predicted_next.is_empty());
    }

    #[test]
    fn sweep_read_access_stats() {
        let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
        let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
        let datatype = make_f64_type();
        let cache = ChunkCache::new();
        let mut sweep = SweepContext::with_defaults();

        read_chunked_data_sweep(
            &file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache, &mut sweep,
        )
        .unwrap();

        let stats = cache.access_stats();
        // We accessed 2 chunks; the second should be sequential to the first
        assert!(stats.sequential_count > 0 || stats.random_count > 0);
    }
}