clawhdf5/crates/clawhdf5-py/src/convert.rs

//! HDF5 datatypes as numpy dtypes, and element bytes as numpy arrays.
//!
//! The dtype a file's datatype maps to is the one h5py reports for it
//! (byte order kept, compound offsets and padding kept, `r`/`i` compounds as
//! complex, the `FALSE`/`TRUE` enum as `bool`, fixed strings as `S<n>`,
//! variable-length data as `object`). For every fixed-size type that dtype
//! describes the file's element bytes exactly, so the bytes the library
//! returns become the array's buffer as they are: the `Vec<u8>` is handed to
//! numpy without a copy and viewed as the dtype.
//!
//! Anything this mapping cannot describe exactly — non-IEEE floats, integers
//! with padding bits, VAX byte order, references, bitfields, time, and
//! variable-length members inside compounds or arrays — is a `TypeError`,
//! never a best-effort guess.

use std::collections::HashMap;

use clawhdf5_format::datatype::{CharacterSet, Datatype, DatatypeByteOrder};
use clawhdf5_format::global_heap::GlobalHeapCollection;
use numpy::PyArray1;
use pyo3::exceptions::{PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict, PyList, PyString, PyTuple};

/// How the elements of a datatype become Python values.
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum Layout {
    /// Fixed-size elements numpy reads as they are.
    Fixed,
    /// A top-level HDF5 array type: elements are viewed as the base dtype and
    /// the array's dimensions are appended to the shape (as h5py does).
    Subarray(Vec<usize>),
    /// Variable-length string: a global heap reference per element.
    VlString { utf8: bool },
    /// Variable-length sequence of a fixed-size base type.
    VlSequence,
}

/// Everything needed to turn a dataset's or attribute's bytes into numpy.
pub(crate) struct Converter {
    /// The dtype reported to Python (`Dataset.dtype`).
    pub dtype: Py<PyAny>,
    /// The dtype the element bytes are viewed as: `dtype` itself, the base
    /// of a subarray, or the base of a variable-length sequence.
    pub view: Py<PyAny>,
    pub layout: Layout,
    /// Bytes per element in the raw buffer the library returns.
    pub elem_size: usize,
    /// For variable-length data, bytes per unit of an element's stored
    /// length: 1 for strings, the base type's size for sequences.
    pub vl_unit: usize,
}

fn unsupported(what: impl std::fmt::Display) -> PyErr {
    PyTypeError::new_err(format!(
        "clawhdf5 cannot read this datatype into numpy: {what}"
    ))
}

fn byte_order_char(order: &DatatypeByteOrder, size: u32) -> PyResult<&'static str> {
    if size == 1 {
        return Ok("|");
    }
    match order {
        DatatypeByteOrder::LittleEndian => Ok("<"),
        DatatypeByteOrder::BigEndian => Ok(">"),
        DatatypeByteOrder::Vax => Err(unsupported("VAX byte order")),
    }
}

/// The numpy format string of an integer type, if it is a plain one.
fn int_format(dt: &Datatype) -> PyResult<String> {
    match dt {
        Datatype::FixedPoint {
            size,
            byte_order,
            signed,
            bit_offset,
            bit_precision,
        } => {
            if !matches!(size, 1 | 2 | 4 | 8) {
                return Err(unsupported(format!("{size}-byte integer")));
            }
            if *bit_offset != 0 || u32::from(*bit_precision) != size * 8 {
                return Err(unsupported(format!(
                    "integer with {bit_precision} significant bits at offset {bit_offset} in {size} bytes"
                )));
            }
            let kind = if *signed { 'i' } else { 'u' };
            Ok(format!(
                "{}{kind}{size}",
                byte_order_char(byte_order, *size)?
            ))
        }
        other => Err(unsupported(format!("{other:?} is not an integer"))),
    }
}

/// The numpy format string of an IEEE 754 binary16/32/64 type.
fn float_format(dt: &Datatype) -> PyResult<String> {
    let Datatype::FloatingPoint {
        size,
        byte_order,
        bit_offset,
        bit_precision,
        exponent_location,
        exponent_size,
        mantissa_location,
        mantissa_size,
        exponent_bias,
    } = dt
    else {
        return Err(unsupported(format!("{dt:?} is not a float")));
    };
    // (exponent location, exponent size, mantissa size, bias) of IEEE 754.
    let ieee = match size {
        2 => (10, 5, 10, 15),
        4 => (23, 8, 23, 127),
        8 => (52, 11, 52, 1023),
        _ => return Err(unsupported(format!("{size}-byte float"))),
    };
    let layout = (
        *exponent_location,
        *exponent_size,
        *mantissa_size,
        *exponent_bias,
    );
    if *bit_offset != 0
        || u32::from(*bit_precision) != size * 8
        || *mantissa_location != 0
        || layout != ieee
    {
        return Err(unsupported(format!(
            "non-IEEE {size}-byte float (exponent {exponent_size} bits at {exponent_location}, \
             mantissa {mantissa_size} bits at {mantissa_location}, bias {exponent_bias})"
        )));
    }
    Ok(format!("{}f{size}", byte_order_char(byte_order, *size)?))
}

/// `r`/`i` compounds of two identical IEEE floats are complex numbers in h5py.
fn complex_format(
    size: u32,
    members: &[clawhdf5_format::datatype::CompoundMember],
) -> Option<String> {
    let [re, im] = members else { return None };
    if re.name != "r" || im.name != "i" || re.datatype != im.datatype {
        return None;
    }
    let Datatype::FloatingPoint {
        size: fsize,
        byte_order,
        ..
    } = &re.datatype
    else {
        return None;
    };
    if !matches!(fsize, 4 | 8)
        || re.byte_offset != 0
        || im.byte_offset != u64::from(*fsize)
        || size != 2 * fsize
    {
        return None;
    }
    float_format(&re.datatype).ok()?;
    let order = byte_order_char(byte_order, *fsize).ok()?;
    Some(format!("{order}c{}", 2 * fsize))
}

/// The members of an enum as `{name: value}`.
fn enum_members<'py>(
    py: Python<'py>,
    base: &Datatype,
    members: &[clawhdf5_format::datatype::EnumMember],
) -> PyResult<Bound<'py, PyDict>> {
    let signed = matches!(base, Datatype::FixedPoint { signed: true, .. });
    let dict = PyDict::new(py);
    for m in members {
        let value: Py<PyAny> = if signed {
            let v = clawhdf5_format::data_read::read_as_i64(&m.value, base)
                .map_err(|e| PyValueError::new_err(format!("enum member {}: {e}", m.name)))?;
            let v = *v.first().ok_or_else(|| {
                PyValueError::new_err(format!("enum member {} has no value", m.name))
            })?;
            v.into_pyobject(py)?.into_any().unbind()
        } else {
            let v = clawhdf5_format::data_read::read_as_u64(&m.value, base)
                .map_err(|e| PyValueError::new_err(format!("enum member {}: {e}", m.name)))?;
            let v = *v.first().ok_or_else(|| {
                PyValueError::new_err(format!("enum member {} has no value", m.name))
            })?;
            v.into_pyobject(py)?.into_any().unbind()
        };
        dict.set_item(&m.name, value)?;
    }
    Ok(dict)
}

/// Whether an enum is h5py's boolean: a one-byte integer with exactly the
/// members `FALSE` = 0 and `TRUE` = 1.
fn is_h5py_bool(base: &Datatype, members: &[clawhdf5_format::datatype::EnumMember]) -> bool {
    if base.type_size() != 1 || members.len() != 2 {
        return false;
    }
    let value = |name: &str| {
        members
            .iter()
            .find(|m| m.name == name)
            .and_then(|m| m.value.first().copied())
    };
    value("FALSE") == Some(0) && value("TRUE") == Some(1)
}

fn np_dtype<'py>(py: Python<'py>, spec: impl IntoPyObject<'py>) -> PyResult<Bound<'py, PyAny>> {
    py.import("numpy")?.getattr("dtype")?.call1((spec,))
}

fn np_dtype_with_metadata<'py>(
    py: Python<'py>,
    spec: impl IntoPyObject<'py>,
    metadata: Bound<'py, PyDict>,
) -> PyResult<Bound<'py, PyAny>> {
    let kwargs = PyDict::new(py);
    kwargs.set_item("metadata", metadata)?;
    py.import("numpy")?
        .getattr("dtype")?
        .call((spec,), Some(&kwargs))
}

/// The numpy dtype of a fixed-size datatype, whose element bytes numpy can
/// read as they are.
pub(crate) fn fixed_dtype<'py>(py: Python<'py>, dt: &Datatype) -> PyResult<Bound<'py, PyAny>> {
    match dt {
        Datatype::FixedPoint { .. } => np_dtype(py, int_format(dt)?),
        Datatype::FloatingPoint { .. } => np_dtype(py, float_format(dt)?),
        Datatype::String { size, charset, .. } => {
            if *size == 0 {
                return Err(unsupported("zero-length fixed string"));
            }
            let meta = PyDict::new(py);
            let enc = match charset {
                CharacterSet::Ascii => "ascii",
                CharacterSet::Utf8 => "utf-8",
            };
            meta.set_item("h5py_encoding", enc)?;
            np_dtype_with_metadata(py, format!("S{size}"), meta)
        }
        Datatype::Opaque { size, .. } => {
            if *size == 0 {
                return Err(unsupported("zero-length opaque type"));
            }
            np_dtype(py, format!("V{size}"))
        }
        Datatype::Enumeration {
            base_type, members, ..
        } => {
            let base = int_format(base_type)?;
            if is_h5py_bool(base_type, members) {
                return np_dtype(py, "?");
            }
            let meta = PyDict::new(py);
            meta.set_item("enum", enum_members(py, base_type, members)?)?;
            np_dtype_with_metadata(py, base, meta)
        }
        Datatype::Compound { size, members } => {
            if let Some(c) = complex_format(*size, members) {
                return np_dtype(py, c);
            }
            let names = PyList::empty(py);
            let formats = PyList::empty(py);
            let offsets = PyList::empty(py);
            for m in members {
                let end = m.byte_offset.checked_add(u64::from(m.datatype.type_size()));
                if end.is_none_or(|end| end > u64::from(*size)) {
                    return Err(PyValueError::new_err(format!(
                        "compound member {} lies outside the {size}-byte compound",
                        m.name
                    )));
                }
                names.append(&m.name)?;
                formats.append(fixed_dtype(py, &m.datatype).map_err(|e| {
                    unsupported(format!("compound member {}: {}", m.name, e.value(py)))
                })?)?;
                offsets.append(m.byte_offset)?;
            }
            let spec = PyDict::new(py);
            spec.set_item("names", names)?;
            spec.set_item("formats", formats)?;
            spec.set_item("offsets", offsets)?;
            spec.set_item("itemsize", size)?;
            np_dtype(py, spec)
        }
        Datatype::Array {
            base_type,
            dimensions,
        } => {
            let base = fixed_dtype(py, base_type)?;
            let dims = PyTuple::new(py, dimensions)?;
            np_dtype(py, (base, dims))
        }
        Datatype::VariableLength { is_string, .. } => Err(unsupported(if *is_string {
            "variable-length string inside a compound or array type"
        } else {
            "variable-length sequence inside a compound or array type"
        })),
        Datatype::Reference { .. } => Err(unsupported("object/region references")),
        Datatype::BitField { .. } => Err(unsupported("bitfield")),
        Datatype::Time { .. } => Err(unsupported("time")),
    }
}

impl Converter {
    /// The converter for a dataset's or attribute's datatype.
    pub(crate) fn new(py: Python<'_>, dt: &Datatype, offset_size: u8) -> PyResult<Self> {
        match dt {
            Datatype::VariableLength {
                is_string: true,
                charset,
                ..
            } => {
                let utf8 = matches!(charset, Some(CharacterSet::Utf8));
                let meta = PyDict::new(py);
                if utf8 {
                    meta.set_item("vlen", py.get_type::<PyString>())?;
                } else {
                    meta.set_item("vlen", py.get_type::<PyBytes>())?;
                }
                let dtype = np_dtype_with_metadata(py, "O", meta)?;
                Ok(Self {
                    view: dtype.clone().unbind(),
                    dtype: dtype.unbind(),
                    layout: Layout::VlString { utf8 },
                    elem_size: vl_ref_size(offset_size)?,
                    vl_unit: 1,
                })
            }
            Datatype::VariableLength {
                is_string: false,
                base_type,
                ..
            } => {
                let base = fixed_dtype(py, base_type)?;
                let meta = PyDict::new(py);
                meta.set_item("vlen", &base)?;
                let dtype = np_dtype_with_metadata(py, "O", meta)?;
                Ok(Self {
                    dtype: dtype.unbind(),
                    view: base.unbind(),
                    layout: Layout::VlSequence,
                    elem_size: vl_ref_size(offset_size)?,
                    vl_unit: base_type.type_size() as usize,
                })
            }
            Datatype::Array {
                base_type,
                dimensions,
            } => {
                let dtype = fixed_dtype(py, dt)?;
                let base = fixed_dtype(py, base_type)?;
                Ok(Self {
                    dtype: dtype.unbind(),
                    view: base.unbind(),
                    layout: Layout::Subarray(dimensions.iter().map(|&d| d as usize).collect()),
                    elem_size: dt.type_size() as usize,
                    vl_unit: 0,
                })
            }
            _ => {
                let dtype = fixed_dtype(py, dt)?;
                Ok(Self {
                    view: dtype.clone().unbind(),
                    dtype: dtype.unbind(),
                    layout: Layout::Fixed,
                    elem_size: dt.type_size() as usize,
                    vl_unit: 0,
                })
            }
        }
    }

    pub(crate) fn is_vl(&self) -> bool {
        matches!(self.layout, Layout::VlString { .. } | Layout::VlSequence)
    }

    /// An empty array of `shape` (some dimension is zero).
    pub(crate) fn empty<'py>(
        &self,
        py: Python<'py>,
        shape: &[usize],
    ) -> PyResult<Bound<'py, PyAny>> {
        let np = py.import("numpy")?;
        match &self.layout {
            Layout::Subarray(dims) => {
                let mut full = shape.to_vec();
                full.extend_from_slice(dims);
                np.call_method1("empty", (PyTuple::new(py, full)?, self.view.bind(py)))
            }
            _ => np.call_method1("empty", (PyTuple::new(py, shape)?, self.dtype.bind(py))),
        }
    }

    /// Turn decoded element data into a numpy array of `shape`.
    ///
    /// `str_values` decodes variable-length strings to `str` (what h5py
    /// does for attributes) instead of `bytes` (what it does for datasets).
    pub(crate) fn to_array<'py>(
        &self,
        py: Python<'py>,
        data: Elements,
        shape: &[usize],
        str_values: bool,
    ) -> PyResult<Bound<'py, PyAny>> {
        let n: usize = shape.iter().product();
        match (data, &self.layout) {
            (Elements::Bytes(bytes), Layout::Fixed) => {
                bytes_as_array(py, bytes, self.view.bind(py), shape)
            }
            (Elements::Bytes(bytes), Layout::Subarray(dims)) => {
                let mut full = shape.to_vec();
                full.extend_from_slice(dims);
                bytes_as_array(py, bytes, self.view.bind(py), &full)
            }
            (Elements::Vl(items), Layout::VlString { .. }) => {
                check_count(items.len(), n)?;
                let mut objs: Vec<Py<PyAny>> = Vec::with_capacity(items.len());
                for item in items {
                    let obj = if str_values {
                        PyBytes::new(py, &item)
                            .call_method1("decode", ("utf-8", "surrogateescape"))?
                            .unbind()
                    } else {
                        PyBytes::new(py, &item).into_any().unbind()
                    };
                    objs.push(obj);
                }
                object_array(py, objs, shape)
            }
            (Elements::Vl(items), Layout::VlSequence) => {
                check_count(items.len(), n)?;
                let base = self.view.bind(py);
                let itemsize: usize = base.getattr("itemsize")?.extract()?;
                let mut objs: Vec<Py<PyAny>> = Vec::with_capacity(items.len());
                for item in items {
                    if item.len() % itemsize != 0 {
                        return Err(PyValueError::new_err(format!(
                            "variable-length element of {} bytes is not a whole number of {itemsize}-byte values",
                            item.len()
                        )));
                    }
                    let len = item.len() / itemsize;
                    objs.push(bytes_as_array(py, item, base, &[len])?.unbind());
                }
                object_array(py, objs, shape)
            }
            _ => Err(PyValueError::new_err(
                "internal error: element data does not match the datatype",
            )),
        }
    }
}

/// Element data as read, before it becomes numpy.
pub(crate) enum Elements {
    /// The elements' bytes, back to back.
    Bytes(Vec<u8>),
    /// Each variable-length element's bytes, resolved from the global heap.
    Vl(Vec<Vec<u8>>),
}

fn vl_ref_size(offset_size: u8) -> PyResult<usize> {
    // The library sizes a variable-length element as 16 bytes (a length, an
    // 8-byte heap address and an index) whatever the file's offset size.
    // Refuse the other sizes rather than read misaligned references.
    if offset_size != 8 {
        return Err(unsupported(format!(
            "variable-length data in a file with {offset_size}-byte offsets"
        )));
    }
    Ok(4 + usize::from(offset_size) + 4)
}

fn check_count(got: usize, want: usize) -> PyResult<()> {
    if got != want {
        return Err(PyValueError::new_err(format!(
            "read {got} elements, expected {want}"
        )));
    }
    Ok(())
}

/// A numpy array over `bytes` without copying them: the `Vec` becomes the
/// array's buffer and is viewed as `dtype` with `shape`.
pub(crate) fn bytes_as_array<'py>(
    py: Python<'py>,
    bytes: Vec<u8>,
    dtype: &Bound<'py, PyAny>,
    shape: &[usize],
) -> PyResult<Bound<'py, PyAny>> {
    let itemsize: usize = dtype.getattr("itemsize")?.extract()?;
    let n: usize = shape.iter().product();
    if n.checked_mul(itemsize) != Some(bytes.len()) {
        return Err(PyValueError::new_err(format!(
            "read {} bytes, expected {n} elements of {itemsize} bytes",
            bytes.len()
        )));
    }
    let shape = PyTuple::new(py, shape)?;
    if n == 0 {
        return py.import("numpy")?.call_method1("empty", (shape, dtype));
    }
    let raw = PyArray1::from_vec(py, bytes);
    let arr = raw
        .call_method1("view", (dtype,))?
        .call_method1("reshape", (shape,))?;
    // A `Vec<u8>` carries no alignment promise. numpy copes with unaligned
    // arrays, but slowly and not in every routine, so hand out an aligned
    // copy in the (allocator-dependent, rare) case the buffer is not.
    if !arr
        .getattr("flags")?
        .getattr("aligned")?
        .extract::<bool>()?
    {
        return arr.call_method0("copy");
    }
    Ok(arr)
}

fn object_array<'py>(
    py: Python<'py>,
    objs: Vec<Py<PyAny>>,
    shape: &[usize],
) -> PyResult<Bound<'py, PyAny>> {
    let arr = PyArray1::from_vec(py, objs);
    arr.call_method1("reshape", (PyTuple::new(py, shape)?,))
}

/// Resolve variable-length elements (global heap references in `raw`) to
/// their bytes: each element's stored length times `unit` (1 for strings,
/// the base type's size for sequences). Pure Rust, so it runs without the
/// GIL.
pub(crate) fn resolve_vl(
    file_data: &[u8],
    raw: &[u8],
    count: usize,
    offset_size: u8,
    length_size: u8,
    unit: usize,
) -> Result<Vec<Vec<u8>>, String> {
    let refs = clawhdf5_format::vl_data::parse_vl_references(raw, count as u64, offset_size)
        .map_err(|e| e.to_string())?;
    let undefined = match offset_size {
        2 => 0xFFFF,
        4 => 0xFFFF_FFFF,
        _ => u64::MAX,
    };
    let mut collections: HashMap<u64, GlobalHeapCollection> = HashMap::new();
    let mut out = Vec::with_capacity(refs.len());
    for vl in &refs {
        if vl.collection_address == 0 || vl.collection_address == undefined {
            if vl.length != 0 {
                return Err(format!(
                    "variable-length element of length {} has no heap address",
                    vl.length
                ));
            }
            out.push(Vec::new());
            continue;
        }
        let coll = match collections.entry(vl.collection_address) {
            std::collections::hash_map::Entry::Occupied(e) => e.into_mut(),
            std::collections::hash_map::Entry::Vacant(e) => {
                let addr = usize::try_from(vl.collection_address)
                    .map_err(|_| "global heap address out of range".to_string())?;
                e.insert(
                    GlobalHeapCollection::parse(file_data, addr, length_size)
                        .map_err(|e| e.to_string())?,
                )
            }
        };
        let index = u16::try_from(vl.object_index)
            .map_err(|_| format!("global heap object index {} out of range", vl.object_index))?;
        let obj = coll.get_object(index).ok_or_else(|| {
            format!(
                "global heap object {index} not found in the collection at {}",
                vl.collection_address
            )
        })?;
        let need = (vl.length as usize)
            .checked_mul(unit)
            .ok_or("variable-length element too long")?;
        if need > obj.data.len() {
            return Err(format!(
                "variable-length element of {need} bytes in a {}-byte heap object",
                obj.data.len()
            ));
        }
        out.push(obj.data[..need].to_vec());
    }
    Ok(out)
}