anamnesis 0.6.7

// SPDX-License-Identifier: MIT OR Apache-2.0

//! Shared parsing utilities used by multiple format parsers (`NPZ`, `.pth`).

/// Soft cap on `Vec` / `HashMap` pre-allocation sized from a file-declared
/// count, shared by the `GGUF`, `NPZ`, and `.pth` parsers.
///
/// A header or archive can declare a large entry count: a `GGUF` header up to
/// `MAX_TENSOR_COUNT` / `MAX_KV_COUNT` (1 M each), or a zip central directory
/// with millions of (possibly empty) entries. Passing that count straight to
/// `with_capacity` commits the whole hint eagerly — ~175 MB for `GGUF` at 1 M
/// (empirically 114 MB for the metadata `HashMap` + 61 MB for the
/// `Vec<GgufTensorInfo>`), and ~1–2× the file size for a many-entries zip,
/// before a single entry is read. Clamping every trust-the-count pre-allocation
/// to this constant bounds the worst-case eager hint to a few dozen KB; the
/// container still grows as entries are inserted, costing at most a handful of
/// extra reallocs on legitimate files — imperceptible given parse is I/O-bound
/// and real files never approach the cap.
#[cfg(any(feature = "npz", feature = "pth", feature = "gguf"))]
pub(crate) const PREALLOC_SOFT_CAP: usize = 256;

/// Reverses the byte order of each element in `data` in-place.
///
/// Each contiguous `element_size`-byte chunk is reversed, converting
/// big-endian to little-endian (or vice versa). Elements that are not
/// an exact multiple of `element_size` are left untouched (handled by
/// `chunks_exact_mut`).
// VECTORIZED: scalar fallback — chunk.reverse() on a runtime-variable
// element_size prevents auto-vectorization. This is the big-endian path
// (<0.01% of ML files), so scalar performance is acceptable.
// EXPLICIT: in-place mutation on the read buffer avoids allocating a second
// buffer of equal size. CONVENTIONS Rule 6 (separate in/out) is waived here
// because the data is already in a dedicated Vec<u8> that serves as the output.
#[cfg(any(feature = "npz", feature = "pth"))]
pub(crate) fn byteswap_inplace(data: &mut [u8], element_size: usize) {
    for chunk in data.chunks_exact_mut(element_size) {
        chunk.reverse();
    }
}

/// Computes the product of a tensor's shape dimensions with overflow checking.
///
/// Returns `None` if the running product overflows `usize` — an adversarial
/// or malformed header declaring dimensions whose product cannot fit (e.g.
/// `[usize::MAX, 2]`). An empty shape yields `Some(1)` (the scalar
/// element-count convention).
///
/// Callers choose their own policy on `None`: a path that sizes an allocation
/// maps it to `AnamnesisError::Parse`, while a pure eligibility query treats
/// it as "not eligible". This is the
/// checked counterpart to a raw `shape.iter().product()`, which silently wraps
/// in release builds and panics in debug builds on overflow.
#[must_use]
pub(crate) fn checked_num_elements(shape: &[usize]) -> Option<usize> {
    shape.iter().try_fold(1usize, |acc, &d| acc.checked_mul(d))
}

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

    #[test]
    fn checked_num_elements_basic() {
        assert_eq!(checked_num_elements(&[]), Some(1));
        assert_eq!(checked_num_elements(&[16384, 2304]), Some(16384 * 2304));
        assert_eq!(checked_num_elements(&[0, 5]), Some(0));
    }

    #[test]
    fn checked_num_elements_overflow_is_none() {
        // Two near-`usize::MAX` dims whose product cannot fit.
        assert_eq!(checked_num_elements(&[usize::MAX, 2]), None);
    }

    #[cfg(any(feature = "npz", feature = "pth"))]
    #[test]
    fn byteswap_2byte() {
        let mut data = vec![0x01, 0x02, 0x03, 0x04];
        byteswap_inplace(&mut data, 2);
        assert_eq!(data, vec![0x02, 0x01, 0x04, 0x03]);
    }

    #[cfg(any(feature = "npz", feature = "pth"))]
    #[test]
    fn byteswap_4byte() {
        let mut data = vec![0x01, 0x02, 0x03, 0x04];
        byteswap_inplace(&mut data, 4);
        assert_eq!(data, vec![0x04, 0x03, 0x02, 0x01]);
    }

    #[cfg(any(feature = "npz", feature = "pth"))]
    #[test]
    fn byteswap_8byte() {
        let mut data = vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08];
        byteswap_inplace(&mut data, 8);
        assert_eq!(data, vec![0x08, 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01]);
    }

    #[cfg(any(feature = "npz", feature = "pth"))]
    #[test]
    fn byteswap_1byte_is_noop() {
        let mut data = vec![0xAA, 0xBB, 0xCC];
        byteswap_inplace(&mut data, 1);
        assert_eq!(data, vec![0xAA, 0xBB, 0xCC]);
    }

    #[cfg(any(feature = "npz", feature = "pth"))]
    #[test]
    fn byteswap_empty() {
        let mut data: Vec<u8> = vec![];
        byteswap_inplace(&mut data, 4);
        assert!(data.is_empty());
    }
}