clawhdf5/crates/clawhdf5-netcdf4/src/dimension.rs

//! NetCDF-4 dimension representation.
//!
//! Dimensions in NetCDF-4 are stored as HDF5 datasets with the CLASS=DIMENSION_SCALE
//! attribute and a `_Netcdf4Dimid` attribute. Unlimited dimensions are detected via
//! the HDF5 dataspace max_dimensions (u64::MAX indicates unlimited); their length
//! is the largest extent of the variables attached to them.

use std::collections::HashMap;

use clawhdf5::AttrValue;

use crate::error::Error;

/// A NetCDF-4 dimension.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Dimension {
    /// Name of this dimension.
    pub name: String,
    /// Current size of this dimension.
    pub size: u64,
    /// Whether this dimension is unlimited (extensible).
    pub is_unlimited: bool,
}

/// A dimension scale of one group: the dataset that defines a dimension.
#[derive(Debug, Clone)]
pub(crate) struct Scale {
    /// Object header address of the scale's dataset (what a variable's
    /// `DIMENSION_LIST` references).
    pub address: u64,
    /// Its `_Netcdf4Dimid` (what a variable's `_Netcdf4Coordinates` lists).
    pub dimid: Option<i64>,
    /// Index of its dimension in [`GroupDims::dims`].
    pub dim: usize,
}

/// The dimensions a group defines, with the scales that define them.
#[derive(Debug, Clone, Default)]
pub(crate) struct GroupDims {
    /// The group's dimensions, in `_Netcdf4Dimid` order (then discovery order).
    pub dims: Vec<Dimension>,
    /// The dimension scales behind `dims`; empty when the group has no
    /// dimension scale and `dims` were inferred from 1-D datasets.
    pub scales: Vec<Scale>,
}

impl GroupDims {
    /// The dimension defined by the scale at `address`.
    pub fn by_address(&self, address: u64) -> Option<&Dimension> {
        self.scales
            .iter()
            .find(|s| s.address == address)
            .map(|s| &self.dims[s.dim])
    }

    /// The dimension whose scale has `_Netcdf4Dimid` `id`.
    pub fn by_dimid(&self, id: i64) -> Option<&Dimension> {
        self.scales
            .iter()
            .find(|s| s.dimid == Some(id))
            .map(|s| &self.dims[s.dim])
    }
}

/// The dimensions of an HDF5 group (root or subgroup).
///
/// NetCDF-4 stores dimensions as datasets with `CLASS=DIMENSION_SCALE`. A fixed
/// dimension's size is the dataset's first (and typically only) shape extent.
/// Unlimited dimensions have `max_dimensions[0] == u64::MAX` in the HDF5 dataspace;
/// their size is computed by `unlimited_len`. A group with no dimension
/// scale at all (not written by a netCDF library) gets one dimension per
/// 1-D dataset instead.
pub(crate) fn group_dims(
    file: &clawhdf5::File,
    group: &clawhdf5::Group<'_>,
) -> Result<GroupDims, Error> {
    let addresses: HashMap<String, u64> = group.entries()?.into_iter().collect();
    let dataset_names = group.datasets()?;
    // (dimid, dimension, scale address), in discovery order.
    let mut found: Vec<(Option<i64>, Dimension, u64)> = Vec::new();

    for ds_name in &dataset_names {
        let ds = group.dataset(ds_name)?;
        let attrs = ds.attrs()?;
        if !is_dimension_scale(&attrs) {
            continue;
        }
        let Some(&address) = addresses.get(ds_name) else {
            continue;
        };
        let shape = ds.shape()?;
        let is_unlimited = is_unlimited(&ds);
        let size = if is_unlimited {
            unlimited_len(file, &attrs, &shape)
        } else {
            shape.first().copied().unwrap_or(0)
        };
        let dim = Dimension {
            name: ds_name.clone(),
            size,
            is_unlimited,
        };
        found.push((get_dimid(&attrs), dim, address));
    }

    if found.is_empty() {
        // Fallback: infer dimensions from dataset shapes and names.
        // In NetCDF-4, coordinate variables are datasets whose name matches
        // a dimension name. If there are no explicit DIMENSION_SCALE attributes,
        // we look for 1-D datasets that might be coordinate variables.
        let mut dims = Vec::new();
        for ds_name in &dataset_names {
            let ds = group.dataset(ds_name)?;
            let shape = ds.shape()?;
            if shape.len() == 1 {
                dims.push(Dimension {
                    name: ds_name.clone(),
                    size: shape[0],
                    is_unlimited: is_unlimited(&ds),
                });
            }
        }
        return Ok(GroupDims {
            dims,
            scales: Vec::new(),
        });
    }

    // By dimid; scales without one keep their discovery order after those
    // with one (the sort is stable).
    found.sort_by_key(|(id, ..)| (id.is_none(), id.unwrap_or(0)));
    let mut out = GroupDims::default();
    for (i, (dimid, dim, address)) in found.into_iter().enumerate() {
        out.dims.push(dim);
        out.scales.push(Scale {
            address,
            dimid,
            dim: i,
        });
    }
    Ok(out)
}

/// The start of the `NAME` attribute netCDF-C gives a dimension scale that
/// is only a dimension, not also a (coordinate) variable.
const PURE_DIMENSION_NAME: &str = "This is a netCDF dimension but not a netCDF variable";

/// The current length of an unlimited dimension, as netCDF-C reports it
/// (`NC4_inq_dim` → `nc4_find_dim_len`): the largest current extent, along
/// the dimension, of the variables that use it, in any group; 0 when none
/// has been written. netCDF-C does not extend a dimension scale that is not
/// also a variable, so such a scale's own extent (0) is not counted; a
/// coordinate variable's is. The variables are the scale's attachments,
/// listed with the axis they use in its `REFERENCE_LIST` attribute (the
/// mirror of each variable's `DIMENSION_LIST`). Attachments that cannot be
/// read are skipped; without a readable `REFERENCE_LIST` the length is the
/// scale's own extent, as before.
fn unlimited_len(file: &clawhdf5::File, attrs: &HashMap<String, AttrValue>, shape: &[u64]) -> u64 {
    let own = shape.first().copied().unwrap_or(0);
    let is_variable = !is_pure_dimension(attrs);
    let Some(refs) = reference_list(file, attrs) else {
        return own;
    };
    refs.into_iter()
        .filter_map(|(address, axis)| {
            let shape = file.dataset_at(address).ok()?.shape().ok()?;
            shape.get(usize::try_from(axis).ok()?).copied()
        })
        .chain(is_variable.then_some(own))
        .max()
        .unwrap_or(0)
}

/// The `(dataset address, axis)` pairs of a dimension scale's
/// `REFERENCE_LIST` attribute (HDF5 dimension scales: a compound of an
/// object reference `dataset` and an integer `dimension`), or `None` when it
/// is missing or not in that form.
fn reference_list(
    file: &clawhdf5::File,
    attrs: &HashMap<String, AttrValue>,
) -> Option<Vec<(u64, u64)>> {
    use clawhdf5_format::data_read::{read_compound_field, read_object_references};
    use clawhdf5_format::datatype::{Datatype, DatatypeByteOrder};
    let Some(AttrValue::Raw { datatype, data, .. }) = attrs.get("REFERENCE_LIST") else {
        return None;
    };
    let dataset = read_compound_field(data, datatype, "dataset").ok()?;
    let addresses = read_object_references(
        &dataset.raw_data,
        &dataset.datatype,
        file.superblock().offset_size,
    )
    .ok()?;
    let dimension = read_compound_field(data, datatype, "dimension").ok()?;
    let Datatype::FixedPoint {
        size, byte_order, ..
    } = dimension.datatype
    else {
        return None;
    };
    let size = usize::try_from(size).ok().filter(|s| (1..=8).contains(s))?;
    let axes = dimension.raw_data.chunks_exact(size).map(|b| {
        let mut v = [0u8; 8];
        match byte_order {
            DatatypeByteOrder::BigEndian => {
                v[8 - size..].copy_from_slice(b);
                u64::from_be_bytes(v)
            }
            _ => {
                v[..size].copy_from_slice(b);
                u64::from_le_bytes(v)
            }
        }
    });
    if axes.len() != addresses.len() {
        return None;
    }
    Some(addresses.into_iter().map(|r| r.address).zip(axes).collect())
}

/// The dimension scale attached to each axis of a variable, from its
/// `DIMENSION_LIST` attribute (HDF5 dimension scales: one variable-length
/// sequence of object references per axis) — the address of the scale
/// netCDF-C takes for the axis, or `None` for an axis with none. netCDF-C's
/// `dimscale_visitor` lets `H5DSiterate_scales` visit every scale attached
/// to the axis and keeps the last, so with several (h5py's `attach_scale`
/// twice) the last one is the axis's dimension. `None` overall when the
/// attribute is missing or not in that form.
pub(crate) fn dimension_list(
    file: &clawhdf5::File,
    attrs: &HashMap<String, AttrValue>,
) -> Option<Vec<Option<u64>>> {
    use clawhdf5_format::data_read::read_object_references;
    use clawhdf5_format::datatype::Datatype;
    use clawhdf5_format::vl_data::VlResolver;
    let Some(AttrValue::Raw { datatype, data, .. }) = attrs.get("DIMENSION_LIST") else {
        return None;
    };
    let Datatype::VariableLength {
        is_string: false,
        base_type,
        ..
    } = datatype
    else {
        return None;
    };
    let sb = file.superblock();
    let base_size = usize::try_from(base_type.type_size()).ok()?;
    let sequences = VlResolver::new_in(file.storage(), sb.offset_size, sb.length_size)
        .sequences(data, base_size)
        .ok()?;
    sequences
        .iter()
        .map(|refs| {
            let refs = read_object_references(refs, base_type, sb.offset_size).ok()?;
            Some(refs.iter().rev().find(|r| !r.is_null()).map(|r| r.address))
        })
        .collect()
}

/// Whether a dataset is a dimension scale that is only a dimension, not a
/// netCDF variable: netCDF-C and h5netcdf give it this `NAME`, and netCDF-C
/// does not list it among the variables.
pub(crate) fn is_pure_dimension(attrs: &HashMap<String, AttrValue>) -> bool {
    is_dimension_scale(attrs)
        && matches!(
            attrs.get("NAME"),
            Some(AttrValue::String(n)) if n.starts_with(PURE_DIMENSION_NAME)
        )
}

/// Check if a dataset's attributes mark it as a dimension scale.
pub(crate) fn is_dimension_scale(attrs: &HashMap<String, AttrValue>) -> bool {
    if let Some(AttrValue::String(class)) = attrs.get("CLASS") {
        return class == "DIMENSION_SCALE";
    }
    false
}

/// Get the _Netcdf4Dimid attribute value if present.
pub(crate) fn get_dimid(attrs: &HashMap<String, AttrValue>) -> Option<i64> {
    match attrs.get("_Netcdf4Dimid") {
        Some(AttrValue::I64(id)) => Some(*id),
        Some(AttrValue::U64(id)) => Some(*id as i64),
        _ => None,
    }
}

/// Whether a dataset's first axis is unlimited (`max_dimensions[0] ==
/// u64::MAX` in its dataspace).
fn is_unlimited(ds: &clawhdf5::Dataset<'_>) -> bool {
    matches!(ds.max_dimensions(), Ok(Some(max_dims)) if max_dims.first() == Some(&u64::MAX))
}