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

//! In-place editing (`clawhdf5.File(path, 'r+')`) through `FileEditor`:
//! turning what Python assigns into the bytes, selections and attribute
//! values the editor takes.
//!
//! Value conversion follows h5py (see `edit_helpers.py`, run inside the
//! extension module); what `FileEditor` cannot do is `NotImplementedError`
//! before anything is written.

use std::ffi::CString;

use clawhdf5_format::datatype::{
    CharacterSet, CompoundMember, Datatype, DatatypeByteOrder, EnumMember, StringPadding,
};
use clawhdf5_format::selection::Selection;
use clawhdf5_rs::AttrValue;
use pyo3::exceptions::{PyNotImplementedError, PyTypeError};
use pyo3::prelude::*;
use pyo3::sync::PyOnceLock;
use pyo3::types::{PyBytes, PyModule, PyTuple};

use crate::select::{Axis, Plan};

/// The helper module, compiled once.
pub(crate) fn helpers(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
    static HELPERS: PyOnceLock<Py<PyModule>> = PyOnceLock::new();
    let module = HELPERS.get_or_try_init(py, || -> PyResult<Py<PyModule>> {
        let code = CString::new(include_str!("edit_helpers.py"))
            .map_err(|e| PyTypeError::new_err(e.to_string()))?;
        Ok(PyModule::from_code(
            py,
            &code,
            c"clawhdf5/edit_helpers.py",
            c"clawhdf5._edit_helpers",
        )?
        .unbind())
    })?;
    Ok(module.bind(py))
}

fn not_implemented(what: impl std::fmt::Display) -> PyErr {
    PyNotImplementedError::new_err(format!(
        "{what} is not supported by clawhdf5's in-place editor"
    ))
}

/// How values for a dataset of type `dt` are converted (a category of
/// `edit_helpers._convert_array`), or why they cannot be written.
pub(crate) fn category(dt: &Datatype) -> PyResult<&'static str> {
    match dt {
        Datatype::Complex { .. } => Ok("complex"),
        Datatype::FixedPoint { .. } => Ok("int"),
        Datatype::FloatingPoint { .. } if crate::convert::is_ieee_float(dt) => Ok("float"),
        // Read as a wider IEEE float; writing would need the reverse
        // conversion (rounding into bfloat16, FP8, ...).
        Datatype::FloatingPoint { .. } => Err(not_implemented(
            "writing non-IEEE floats (bfloat16, FP8, FP6, FP4, ...)",
        )),
        Datatype::Enumeration {
            base_type, members, ..
        } => {
            let is_bool = base_type.type_size() == 1
                && members.len() == 2
                && members
                    .iter()
                    .any(|m| m.name == "FALSE" && m.value.first() == Some(&0))
                && members
                    .iter()
                    .any(|m| m.name == "TRUE" && m.value.first() == Some(&1));
            Ok(if is_bool { "bool" } else { "enum" })
        }
        Datatype::String {
            padding: StringPadding::NullPad,
            ..
        } => Ok("string"),
        Datatype::String { padding, .. } => Err(not_implemented(format!(
            "writing fixed-length strings padded {padding:?} (libhdf5 converts them \
             differently from numpy)"
        ))),
        Datatype::Compound { size, members } => {
            if is_complex(*size, members) {
                return Ok("complex");
            }
            check_exact(dt)?;
            Ok("exact")
        }
        Datatype::Opaque { .. } => Ok("exact"),
        Datatype::Array { .. } => Err(not_implemented("writing HDF5 array-type elements")),
        Datatype::VariableLength { .. } => Err(not_implemented("writing variable-length data")),
        Datatype::Reference { .. } => Err(not_implemented("writing references")),
        Datatype::BitField { .. } => Err(not_implemented("writing bitfields")),
        Datatype::Time { .. } => Err(not_implemented("writing time values")),
    }
}

/// h5py's complex numbers: a compound of two identical floats `r`, `i`.
fn is_complex(size: u32, members: &[CompoundMember]) -> bool {
    matches!(members, [r, i] if r.name == "r" && i.name == "i"
        && r.datatype == i.datatype
        && matches!(r.datatype, Datatype::FloatingPoint { size: fs, .. }
            if r.byte_offset == 0 && i.byte_offset == u64::from(fs) && size == 2 * fs))
}

/// Compound members written byte for byte from the same numpy dtype: fine
/// unless libhdf5 would convert them on the way (strings padded other than
/// with NULs), or the editor cannot write them at all.
fn check_exact(dt: &Datatype) -> PyResult<()> {
    match dt {
        Datatype::Compound { members, .. } => {
            members.iter().try_for_each(|m| check_exact(&m.datatype))
        }
        Datatype::Array { base_type, .. } | Datatype::Complex { base_type, .. } => {
            check_exact(base_type)
        }
        Datatype::String {
            padding: StringPadding::NullPad,
            ..
        }
        | Datatype::FixedPoint { .. }
        | Datatype::FloatingPoint { .. }
        | Datatype::Enumeration { .. }
        | Datatype::Opaque { .. }
        | Datatype::BitField { .. } => Ok(()),
        Datatype::String { .. } => Err(not_implemented(
            "writing compounds with strings not padded with NULs",
        )),
        Datatype::VariableLength { .. } => Err(not_implemented(
            "writing compounds with variable-length members",
        )),
        Datatype::Reference { .. } => Err(not_implemented("writing references")),
        Datatype::Time { .. } => Err(not_implemented("writing time values")),
    }
}

/// Largest point selection an index-list write builds (one coordinate
/// vector per element).
const MAX_POINTS: usize = 1 << 22;

/// The selection `plan` writes, whose elements are numbered as the value's
/// (row-major over the selection's shape).
pub(crate) fn selection(plan: &Plan, dims: &[u64]) -> PyResult<Selection> {
    if plan.axes.is_empty() {
        return Ok(Selection::All);
    }
    if plan.list_axis().is_none() {
        let (reads, _) = plan.reads(dims, None, 1);
        return match <[_; 1]>::try_from(reads) {
            Ok([read]) => Ok(read.sel),
            Err(_) => Err(PyTypeError::new_err("internal error: several hyperslabs")),
        };
    }
    // An index list: the points, in the value's order.
    let per_axis: Vec<Vec<u64>> = plan
        .axes
        .iter()
        .map(|a| match a {
            Axis::Index(i) => vec![*i],
            Axis::Slice { start, step, count } => (0..*count).map(|k| start + k * step).collect(),
            Axis::List(v) => v.clone(),
        })
        .collect();
    let n = per_axis
        .iter()
        .try_fold(1usize, |acc, v| acc.checked_mul(v.len()))
        .filter(|&n| n <= MAX_POINTS)
        .ok_or_else(|| {
            not_implemented(format!(
                "an index-list write of more than {MAX_POINTS} elements (write it in slices)"
            ))
        })?;
    let mut points = Vec::with_capacity(n);
    let mut at = vec![0usize; per_axis.len()];
    for _ in 0..n {
        points.push(at.iter().zip(&per_axis).map(|(&i, v)| v[i]).collect());
        for d in (0..at.len()).rev() {
            at[d] += 1;
            if at[d] < per_axis[d].len() {
                break;
            }
            at[d] = 0;
        }
    }
    Ok(Selection::Points(points))
}

/// The bytes to write for `value` under `plan`, in the dataset's dtype.
pub(crate) fn dataset_bytes(
    py: Python<'_>,
    value: &Bound<'_, PyAny>,
    dtype: &Bound<'_, PyAny>,
    category: &str,
    plan: &Plan,
    chunks: Option<&[u64]>,
) -> PyResult<Vec<u8>> {
    let shape = PyTuple::new(py, plan.out_shape())?;
    let fancy = plan.list_axis().is_some();
    let chunk_elems = chunks.map_or(0, |c| c.iter().fold(1u64, |a, &d| a.saturating_mul(d)));
    let bytes = helpers(py)?.call_method1(
        "dataset_values",
        (value, dtype, category, shape, fancy, chunk_elems),
    )?;
    Ok(bytes.cast::<PyBytes>()?.as_bytes().to_vec())
}

fn ieee_float(size: u32, byte_order: DatatypeByteOrder) -> Option<Datatype> {
    let (exponent_location, exponent_size, mantissa_size, exponent_bias) = match size {
        2 => (10, 5, 10, 15),
        4 => (23, 8, 23, 127),
        8 => (52, 11, 52, 1023),
        _ => return None,
    };
    Some(Datatype::FloatingPoint {
        size,
        byte_order,
        bit_offset: 0,
        bit_precision: (size * 8) as u16,
        exponent_location,
        exponent_size,
        mantissa_location: 0,
        mantissa_size,
        exponent_bias,
    })
}

/// The HDF5 datatype h5py writes for a numpy dtype string (`'<i4'`,
/// `'|b1'`, `'>f8'`, `'<c16'`, `'|S5'`).
fn datatype_of(dtype: &str) -> Option<Datatype> {
    let order = match dtype.as_bytes().first()? {
        b'<' | b'|' | b'=' => DatatypeByteOrder::LittleEndian,
        b'>' => DatatypeByteOrder::BigEndian,
        _ => return None,
    };
    let kind = dtype.as_bytes().get(1)?;
    let size: u32 = dtype.get(2..)?.parse().ok()?;
    match kind {
        b'b' if size == 1 => Some(Datatype::Enumeration {
            size: 1,
            base_type: Box::new(Datatype::FixedPoint {
                size: 1,
                byte_order: DatatypeByteOrder::LittleEndian,
                signed: true,
                bit_offset: 0,
                bit_precision: 8,
            }),
            members: vec![
                EnumMember {
                    name: "FALSE".into(),
                    value: vec![0],
                },
                EnumMember {
                    name: "TRUE".into(),
                    value: vec![1],
                },
            ],
        }),
        b'i' | b'u' if matches!(size, 1 | 2 | 4 | 8) => Some(Datatype::FixedPoint {
            size,
            byte_order: order,
            signed: *kind == b'i',
            bit_offset: 0,
            bit_precision: (size * 8) as u16,
        }),
        b'f' => ieee_float(size, order),
        b'c' => {
            let part = ieee_float(size / 2, order)?;
            Some(Datatype::Compound {
                size,
                members: vec![
                    CompoundMember {
                        name: "r".into(),
                        byte_offset: 0,
                        datatype: part.clone(),
                    },
                    CompoundMember {
                        name: "i".into(),
                        byte_offset: u64::from(size / 2),
                        datatype: part,
                    },
                ],
            })
        }
        b'S' if size > 0 => Some(Datatype::String {
            size,
            padding: StringPadding::NullPad,
            charset: CharacterSet::Ascii,
        }),
        _ => None,
    }
}

/// An attribute value as h5py would store it (`attrs[name] = value`, or
/// `attrs.create(name, data, shape, dtype)`), except that `str` data is
/// stored as fixed-length UTF-8 strings (h5py stores variable-length ones,
/// which the editor cannot write).
pub(crate) fn attr_value(
    py: Python<'_>,
    value: &Bound<'_, PyAny>,
    dtype: Option<&Bound<'_, PyAny>>,
    shape: Option<&Bound<'_, PyAny>>,
) -> PyResult<AttrValue> {
    if value.is_instance_of::<crate::PyEmpty>() {
        return Err(not_implemented(
            "writing an empty (null dataspace) attribute",
        ));
    }
    let (kind, dt, dims, data): (String, String, Vec<u64>, Vec<u8>) = helpers(py)?
        .call_method1("attr_value", (value, dtype, shape))?
        .extract()?;
    let datatype = if kind == "str" {
        let size: u32 = dt
            .parse()
            .map_err(|_| PyTypeError::new_err("bad string size"))?;
        Datatype::String {
            size,
            padding: StringPadding::NullPad,
            charset: CharacterSet::Utf8,
        }
    } else {
        datatype_of(&dt).ok_or_else(|| not_implemented(format!("an attribute of dtype {dt}")))?
    };
    Ok(AttrValue::Raw {
        datatype,
        shape: dims,
        data,
    })
}

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

    #[test]
    fn numpy_dtypes_map_to_h5py_types() {
        assert!(matches!(
            datatype_of("<i4"),
            Some(Datatype::FixedPoint {
                size: 4,
                signed: true,
                byte_order: DatatypeByteOrder::LittleEndian,
                ..
            })
        ));
        assert!(matches!(
            datatype_of(">u2"),
            Some(Datatype::FixedPoint {
                size: 2,
                signed: false,
                byte_order: DatatypeByteOrder::BigEndian,
                ..
            })
        ));
        assert!(matches!(
            datatype_of("<f2"),
            Some(Datatype::FloatingPoint { size: 2, .. })
        ));
        let c = datatype_of("<c16").unwrap();
        assert!(matches!(&c, Datatype::Compound { size: 16, members } if is_complex(16, members)));
        assert_eq!(category(&c).unwrap(), "complex");
        let b = datatype_of("|b1").unwrap();
        assert_eq!(category(&b).unwrap(), "bool");
        assert!(matches!(
            datatype_of("|S5"),
            Some(Datatype::String { size: 5, .. })
        ));
        assert!(datatype_of("<f16").is_none());
        assert!(datatype_of("<M8").is_none());
        assert!(datatype_of("|S0").is_none());
    }

    #[test]
    fn list_writes_become_points_in_value_order() {
        let plan = Plan {
            axes: vec![
                Axis::List(vec![1, 4]),
                Axis::Slice {
                    start: 0,
                    step: 2,
                    count: 2,
                },
                Axis::Index(3),
            ],
            fields: vec![],
            scalar: false,
        };
        let sel = selection(&plan, &[5, 4, 4]).unwrap();
        assert_eq!(
            sel,
            Selection::Points(vec![
                vec![1, 0, 3],
                vec![1, 2, 3],
                vec![4, 0, 3],
                vec![4, 2, 3]
            ])
        );
    }
}