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

//! Where the parsers read the file from: the [`Storage`] trait.
//!
//! Every parser used to take the whole file as one `&[u8]`. [`Storage`] is
//! the abstraction that replaces it (see `docs/design/range-reads.md`,
//! option (a)): a parser asks for the bytes it needs, `[offset, offset +
//! len)`, with 64-bit offsets, and gets them back as a [`Cow`] — borrowed
//! when the backend holds the file in memory (a `Vec`, an mmap), owned when
//! it had to fetch them (a range request, a block cache).
//!
//! `impl Storage for [u8]` serves the in-memory case with no copy, and
//! [`Storage::as_contiguous`] lets a hot loop borrow the whole file at once
//! when the backend has it. Modules are converted one at a time: a converted
//! parser has an `*_in(file: &dyn Storage, ..)` core and keeps its old
//! `&[u8]` signature as a thin wrapper, so callers do not change.
//!
//! The trait is synchronous and `no_std`: parsing is CPU work, and a remote
//! backend bridges to its own I/O.

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

use core::ops::Range;

use crate::error::FormatError;

/// A random-access source of file bytes.
///
/// Offsets are relative to the start of the HDF5 data (the superblock), like
/// every address in the file.
pub trait Storage {
    /// Bytes `[offset, offset + len)`.
    ///
    /// The result is shorter than `len` only when the range runs past the
    /// end of the storage (and empty when `offset` is at or past the end);
    /// a backend that cannot serve a range returns an error instead of a
    /// short read.
    fn read_at(&self, offset: u64, len: usize) -> Result<Cow<'_, [u8]>, FormatError>;

    /// Current length of the storage in bytes.
    fn len(&self) -> u64;

    /// Whether the storage holds no bytes.
    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Several reads at once, in the order given. Backends that talk to a
    /// remote store coalesce and parallelise these; the default reads them
    /// one by one with [`Storage::read_at`].
    fn read_ranges(&self, ranges: &[Range<u64>]) -> Result<Vec<Cow<'_, [u8]>>, FormatError> {
        ranges
            .iter()
            .map(|r| {
                let len = usize::try_from(r.end.saturating_sub(r.start)).map_err(|_| {
                    FormatError::Overflow("read range longer than the address space".into())
                })?;
                self.read_at(r.start, len)
            })
            .collect()
    }

    /// The whole storage as one slice, when the backend has it in memory
    /// (a `Vec`, an mmap). Hot loops use this to keep their zero-copy path;
    /// `None` means every byte has to go through [`Storage::read_at`].
    fn as_contiguous(&self) -> Option<&[u8]> {
        None
    }
}

impl Storage for [u8] {
    #[inline]
    fn read_at(&self, offset: u64, len: usize) -> Result<Cow<'_, [u8]>, FormatError> {
        let n = self.len();
        let start = usize::try_from(offset).map_or(n, |o| o.min(n));
        let end = start.saturating_add(len).min(n);
        Ok(Cow::Borrowed(&self[start..end]))
    }

    #[inline]
    fn len(&self) -> u64 {
        <[u8]>::len(self) as u64
    }

    #[inline]
    fn as_contiguous(&self) -> Option<&[u8]> {
        Some(self)
    }
}

impl Storage for Vec<u8> {
    #[inline]
    fn read_at(&self, offset: u64, len: usize) -> Result<Cow<'_, [u8]>, FormatError> {
        self.as_slice().read_at(offset, len)
    }

    #[inline]
    fn len(&self) -> u64 {
        Vec::len(self) as u64
    }

    #[inline]
    fn as_contiguous(&self) -> Option<&[u8]> {
        Some(self.as_slice())
    }
}

impl<T: Storage + ?Sized> Storage for &T {
    #[inline]
    fn read_at(&self, offset: u64, len: usize) -> Result<Cow<'_, [u8]>, FormatError> {
        (**self).read_at(offset, len)
    }

    #[inline]
    fn len(&self) -> u64 {
        (**self).len()
    }

    #[inline]
    fn read_ranges(&self, ranges: &[Range<u64>]) -> Result<Vec<Cow<'_, [u8]>>, FormatError> {
        (**self).read_ranges(ranges)
    }

    #[inline]
    fn as_contiguous(&self) -> Option<&[u8]> {
        (**self).as_contiguous()
    }
}

impl<T: Storage + ?Sized> Storage for Box<T> {
    #[inline]
    fn read_at(&self, offset: u64, len: usize) -> Result<Cow<'_, [u8]>, FormatError> {
        (**self).read_at(offset, len)
    }

    #[inline]
    fn len(&self) -> u64 {
        (**self).len()
    }

    #[inline]
    fn read_ranges(&self, ranges: &[Range<u64>]) -> Result<Vec<Cow<'_, [u8]>>, FormatError> {
        (**self).read_ranges(ranges)
    }

    #[inline]
    fn as_contiguous(&self) -> Option<&[u8]> {
        (**self).as_contiguous()
    }
}

#[cfg(feature = "std")]
impl<T: Storage + ?Sized> Storage for std::sync::Arc<T> {
    #[inline]
    fn read_at(&self, offset: u64, len: usize) -> Result<Cow<'_, [u8]>, FormatError> {
        (**self).read_at(offset, len)
    }

    #[inline]
    fn len(&self) -> u64 {
        (**self).len()
    }

    #[inline]
    fn read_ranges(&self, ranges: &[Range<u64>]) -> Result<Vec<Cow<'_, [u8]>>, FormatError> {
        (**self).read_ranges(ranges)
    }

    #[inline]
    fn as_contiguous(&self) -> Option<&[u8]> {
        (**self).as_contiguous()
    }
}

/// `storage.len()` as the `usize` the parsers' end-of-file errors report
/// (saturating on targets where the file is larger than the address space).
#[inline]
pub(crate) fn len_usize(file: &dyn Storage) -> usize {
    usize::try_from(file.len()).unwrap_or(usize::MAX)
}

/// Bytes `[offset, offset + len)`, all of them.
///
/// A range that runs past the end of the storage is
/// [`FormatError::UnexpectedEof`] with `expected = offset + len` and
/// `available = storage length` — the error the `&[u8]` parsers give for
/// the same bounds check (`offset + len > file_data.len()`).
#[inline]
pub fn read_exact_at(
    file: &dyn Storage,
    offset: u64,
    len: usize,
) -> Result<Cow<'_, [u8]>, FormatError> {
    let eof = || FormatError::UnexpectedEof {
        expected: usize::try_from(offset)
            .unwrap_or(usize::MAX)
            .saturating_add(len),
        available: len_usize(file),
    };
    match offset.checked_add(len as u64) {
        Some(end) if end <= file.len() => {}
        _ => return Err(eof()),
    }
    let bytes = file.read_at(offset, len)?;
    if bytes.len() < len {
        // The storage shrank or the backend served a short read inside the
        // file: never parse a partial structure.
        return Err(short_read());
    }
    Ok(bytes)
}

fn short_read() -> FormatError {
    FormatError::Storage(
        "short read inside the file (the storage shrank or the backend failed)".into(),
    )
}

/// A window of the file: up to `max` bytes read at `base`, fewer only at
/// the end of the file. Its [`Window::ensure`] reports a bounds failure
/// exactly as the whole-file check `ensure_len(file_data, base + rel, n)`
/// did — with the absolute position and the file's length — as long as
/// every position checked lies within the `max` bytes the window was asked
/// for: then a position past the window is past the end of the file.
pub(crate) struct Window<'a> {
    /// The bytes, from `base` on.
    pub bytes: Cow<'a, [u8]>,
    base: usize,
    file_len: usize,
}

impl<'a> Window<'a> {
    /// Read up to `max` bytes at `base`.
    pub fn read(file: &'a dyn Storage, base: u64, max: usize) -> Result<Self, FormatError> {
        Ok(Window {
            bytes: read_upto(file, base, max)?,
            base: usize::try_from(base).unwrap_or(usize::MAX),
            file_len: len_usize(file),
        })
    }

    /// Check that `[rel, rel + needed)` (relative to `base`) is in the file.
    pub fn ensure(&self, rel: usize, needed: usize) -> Result<(), FormatError> {
        match rel.checked_add(needed) {
            Some(end) if end <= self.bytes.len() => Ok(()),
            _ => Err(FormatError::UnexpectedEof {
                expected: self.base.saturating_add(rel).saturating_add(needed),
                available: self.file_len,
            }),
        }
    }
}

/// Up to `max` bytes from `offset` on: fewer only at the end of the
/// storage. For structures whose size is only known once their prefix has
/// been parsed and whose parsers bound-check what they are given.
#[inline]
pub fn read_upto(
    file: &dyn Storage,
    offset: u64,
    max: usize,
) -> Result<Cow<'_, [u8]>, FormatError> {
    let avail = file.len().saturating_sub(offset);
    let len = usize::try_from(avail).map_or(max, |a| a.min(max));
    let bytes = file.read_at(offset, len)?;
    if bytes.len() < len {
        return Err(short_read());
    }
    Ok(bytes)
}

/// Borrow the whole file for a code path that has not been converted to
/// [`Storage`] yet. On a backend without a contiguous view this is the
/// clean [`FormatError::ContiguousStorageRequired`] error, never a guess.
#[inline]
pub fn require_contiguous<'a>(
    file: &'a dyn Storage,
    what: &'static str,
) -> Result<&'a [u8], FormatError> {
    file.as_contiguous()
        .ok_or(FormatError::ContiguousStorageRequired(what))
}

/// A [`Storage`] over an in-memory buffer that serves every byte through
/// [`Storage::read_at`] (its [`Storage::as_contiguous`] is `None`, so no
/// parser can take the whole-slice shortcut), copies what it serves (as a
/// remote backend would), and counts the reads and bytes.
///
/// It is the equivalence harness of the range-read migration: parsing a
/// file through it must give exactly what parsing the `&[u8]` gives, and
/// the counters are the request counts a cacheless range reader would make.
#[derive(Debug)]
pub struct CountingStorage {
    data: Vec<u8>,
    reads: portable_atomic::AtomicU64,
    bytes: portable_atomic::AtomicU64,
}

impl CountingStorage {
    /// Serve `data` (the file from the superblock on).
    pub fn new(data: Vec<u8>) -> Self {
        CountingStorage {
            data,
            reads: portable_atomic::AtomicU64::new(0),
            bytes: portable_atomic::AtomicU64::new(0),
        }
    }

    /// Number of `read_at` calls served so far.
    pub fn reads(&self) -> u64 {
        self.reads.load(portable_atomic::Ordering::Relaxed)
    }

    /// Number of bytes served so far.
    pub fn bytes_read(&self) -> u64 {
        self.bytes.load(portable_atomic::Ordering::Relaxed)
    }

    /// Reset both counters.
    pub fn reset(&self) {
        self.reads.store(0, portable_atomic::Ordering::Relaxed);
        self.bytes.store(0, portable_atomic::Ordering::Relaxed);
    }
}

impl Storage for CountingStorage {
    fn read_at(&self, offset: u64, len: usize) -> Result<Cow<'_, [u8]>, FormatError> {
        let got = self.data.as_slice().read_at(offset, len)?;
        self.reads.fetch_add(1, portable_atomic::Ordering::Relaxed);
        self.bytes
            .fetch_add(got.len() as u64, portable_atomic::Ordering::Relaxed);
        Ok(Cow::Owned(got.into_owned()))
    }

    fn len(&self) -> u64 {
        self.data.len() as u64
    }
}

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

    #[test]
    fn slice_reads_are_borrowed_and_clamped() {
        let data: Vec<u8> = (0u8..10).collect();
        let s: &[u8] = &data;
        let dynamic: &dyn Storage = &s;
        assert_eq!(dynamic.len(), 10);
        let r = dynamic.read_at(2, 3).unwrap();
        assert!(matches!(r, Cow::Borrowed(_)));
        assert_eq!(&*r, &[2, 3, 4]);
        assert_eq!(&*dynamic.read_at(8, 5).unwrap(), &[8, 9]);
        assert!(dynamic.read_at(10, 5).unwrap().is_empty());
        assert!(dynamic.read_at(u64::MAX, 5).unwrap().is_empty());
        assert_eq!(dynamic.as_contiguous(), Some(&data[..]));
        let v: &dyn Storage = &data;
        assert_eq!(v.as_contiguous(), Some(&data[..]));
    }

    #[test]
    fn read_exact_matches_slice_bounds_errors() {
        let data = [0u8; 10];
        let s: &[u8] = &data;
        assert_eq!(&*read_exact_at(&s, 4, 6).unwrap(), &[0; 6]);
        assert_eq!(
            read_exact_at(&s, 4, 7).unwrap_err(),
            FormatError::UnexpectedEof {
                expected: 11,
                available: 10
            }
        );
        assert!(read_exact_at(&s, u64::MAX, 1).is_err());
        assert_eq!(read_upto(&s, 7, 100).unwrap().len(), 3);
        assert_eq!(read_upto(&s, 70, 100).unwrap().len(), 0);
    }

    #[test]
    fn counting_storage_counts_and_hides_the_slice() {
        let c = CountingStorage::new((0u8..10).collect());
        assert!(c.as_contiguous().is_none());
        let r = c.read_at(3, 4).unwrap();
        assert!(matches!(r, Cow::Owned(_)));
        assert_eq!(&*r, &[3, 4, 5, 6]);
        c.read_at(8, 4).unwrap();
        assert_eq!((c.reads(), c.bytes_read()), (2, 6));
        c.reset();
        assert_eq!((c.reads(), c.bytes_read()), (0, 0));
    }

    #[test]
    fn read_ranges_default_loops() {
        let data: Vec<u8> = (0u8..10).collect();
        let s: &[u8] = &data;
        let got = s.read_ranges(&[1..3, 5..9]).unwrap();
        assert_eq!(&*got[0], &[1, 2]);
        assert_eq!(&*got[1], &[5, 6, 7, 8]);
    }
}