netcdf_reader/

types.rs

use crate::error::{Error, Result};

/// A NetCDF dimension.
#[derive(Debug, Clone)]
pub struct NcDimension {
    pub name: String,
    pub size: u64,
    pub is_unlimited: bool,
}

/// A field within a compound (struct) type.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NcCompoundField {
    pub name: String,
    pub offset: u64,
    pub dtype: NcType,
}

/// A typed integer value used by NetCDF-4 enum definitions and values.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NcIntegerValue {
    I8(i8),
    U8(u8),
    I16(i16),
    U16(u16),
    I32(i32),
    U32(u32),
    I64(i64),
    U64(u64),
}

impl NcIntegerValue {
    /// Return the value as `i128` when it is lossless for the signed domain.
    pub fn as_i128(self) -> Option<i128> {
        match self {
            NcIntegerValue::I8(value) => Some(value as i128),
            NcIntegerValue::U8(value) => Some(value as i128),
            NcIntegerValue::I16(value) => Some(value as i128),
            NcIntegerValue::U16(value) => Some(value as i128),
            NcIntegerValue::I32(value) => Some(value as i128),
            NcIntegerValue::U32(value) => Some(value as i128),
            NcIntegerValue::I64(value) => Some(value as i128),
            NcIntegerValue::U64(value) => Some(i128::from(value)),
        }
    }

    /// Return the value as `u128` when it is non-negative.
    pub fn as_u128(self) -> Option<u128> {
        match self {
            NcIntegerValue::I8(value) => u128::try_from(value).ok(),
            NcIntegerValue::U8(value) => Some(value as u128),
            NcIntegerValue::I16(value) => u128::try_from(value).ok(),
            NcIntegerValue::U16(value) => Some(value as u128),
            NcIntegerValue::I32(value) => u128::try_from(value).ok(),
            NcIntegerValue::U32(value) => Some(value as u128),
            NcIntegerValue::I64(value) => u128::try_from(value).ok(),
            NcIntegerValue::U64(value) => Some(value as u128),
        }
    }
}

/// A named member of a NetCDF-4 enum type.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NcEnumMember {
    pub name: String,
    pub value: NcIntegerValue,
}

/// NetCDF data types.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum NcType {
    /// NC_BYTE (i8)
    Byte,
    /// NC_CHAR (u8/char)
    Char,
    /// NC_SHORT (i16)
    Short,
    /// NC_INT (i32)
    Int,
    /// NC_FLOAT (f32)
    Float,
    /// NC_DOUBLE (f64)
    Double,
    /// NC_UBYTE (u8, CDF-5)
    UByte,
    /// NC_USHORT (u16, CDF-5)
    UShort,
    /// NC_UINT (u32, CDF-5)
    UInt,
    /// NC_INT64 (i64, CDF-5)
    Int64,
    /// NC_UINT64 (u64, CDF-5)
    UInt64,
    /// NetCDF-4 only (variable-length string)
    String,
    /// NetCDF-4 enum type with an integer base type.
    Enum {
        base: Box<NcType>,
        members: Vec<NcEnumMember>,
    },
    /// NetCDF-4 compound type (struct with named fields).
    Compound {
        size: u32,
        fields: Vec<NcCompoundField>,
    },
    /// NetCDF-4 opaque type (uninterpreted byte blob).
    Opaque { size: u32, tag: String },
    /// NetCDF-4 array type (fixed-size array of a base type).
    Array { base: Box<NcType>, dims: Vec<u64> },
    /// NetCDF-4 variable-length type.
    VLen { base: Box<NcType> },
}

impl NcType {
    /// Size of a single element in bytes.
    pub fn size(&self) -> Result<usize> {
        match self {
            NcType::Byte | NcType::Char | NcType::UByte => Ok(1),
            NcType::Short | NcType::UShort => Ok(2),
            NcType::Int | NcType::UInt | NcType::Float => Ok(4),
            NcType::Int64 | NcType::UInt64 | NcType::Double => Ok(8),
            // Variable-length string; no fixed element size, but pointer-sized in memory.
            NcType::String => Ok(std::mem::size_of::<usize>()),
            NcType::Enum { base, .. } => base.size(),
            NcType::Compound { size, .. } => Ok(*size as usize),
            NcType::Opaque { size, .. } => Ok(*size as usize),
            NcType::Array { base, dims } => {
                let base_size = base.size()?;
                let count = dims.iter().try_fold(1usize, |acc, &dim| {
                    let dim = usize::try_from(dim).map_err(|_| {
                        Error::InvalidData(
                            "NetCDF array type dimension exceeds platform usize capacity"
                                .to_string(),
                        )
                    })?;
                    acc.checked_mul(dim).ok_or_else(|| {
                        Error::InvalidData(
                            "NetCDF array type element count exceeds platform usize capacity"
                                .to_string(),
                        )
                    })
                })?;
                base_size.checked_mul(count).ok_or_else(|| {
                    Error::InvalidData(
                        "NetCDF array type byte size exceeds platform usize capacity".to_string(),
                    )
                })
            }
            NcType::VLen { .. } => Ok(std::mem::size_of::<usize>()), // pointer-sized
        }
    }

    /// The numeric type code used in CDF-1/2/5 headers.
    pub fn classic_type_code(&self) -> Option<u32> {
        match self {
            NcType::Byte => Some(1),
            NcType::Char => Some(2),
            NcType::Short => Some(3),
            NcType::Int => Some(4),
            NcType::Float => Some(5),
            NcType::Double => Some(6),
            NcType::UByte => Some(7),
            NcType::UShort => Some(8),
            NcType::UInt => Some(9),
            NcType::Int64 => Some(10),
            NcType::UInt64 => Some(11),
            // Extended types are not valid in classic format.
            NcType::String
            | NcType::Enum { .. }
            | NcType::Compound { .. }
            | NcType::Opaque { .. }
            | NcType::Array { .. }
            | NcType::VLen { .. } => None,
        }
    }

    /// Returns true if this is a primitive numeric or string type.
    pub fn is_primitive(&self) -> bool {
        matches!(
            self,
            NcType::Byte
                | NcType::Char
                | NcType::Short
                | NcType::Int
                | NcType::Float
                | NcType::Double
                | NcType::UByte
                | NcType::UShort
                | NcType::UInt
                | NcType::Int64
                | NcType::UInt64
                | NcType::String
        )
    }
}

/// A NetCDF attribute value.
#[derive(Debug, Clone)]
pub enum NcAttrValue {
    Bytes(Vec<i8>),
    Chars(String),
    Shorts(Vec<i16>),
    Ints(Vec<i32>),
    Floats(Vec<f32>),
    Doubles(Vec<f64>),
    UBytes(Vec<u8>),
    UShorts(Vec<u16>),
    UInts(Vec<u32>),
    Int64s(Vec<i64>),
    UInt64s(Vec<u64>),
    Strings(Vec<String>),
}

impl NcAttrValue {
    /// Get the value as a string (for Chars or single-element Strings).
    pub fn as_string(&self) -> Option<String> {
        match self {
            NcAttrValue::Chars(s) => Some(s.clone()),
            NcAttrValue::Strings(v) if v.len() == 1 => Some(v[0].clone()),
            _ => None,
        }
    }

    /// Get the value as f64 (with numeric promotion from the first element).
    pub fn as_f64(&self) -> Option<f64> {
        match self {
            NcAttrValue::Bytes(v) => v.first().map(|&x| x as f64),
            NcAttrValue::Shorts(v) => v.first().map(|&x| x as f64),
            NcAttrValue::Ints(v) => v.first().map(|&x| x as f64),
            NcAttrValue::Floats(v) => v.first().map(|&x| x as f64),
            NcAttrValue::Doubles(v) => v.first().copied(),
            NcAttrValue::UBytes(v) => v.first().map(|&x| x as f64),
            NcAttrValue::UShorts(v) => v.first().map(|&x| x as f64),
            NcAttrValue::UInts(v) => v.first().map(|&x| x as f64),
            NcAttrValue::Int64s(v) => v.first().map(|&x| x as f64),
            NcAttrValue::UInt64s(v) => v.first().map(|&x| x as f64),
            NcAttrValue::Chars(_) | NcAttrValue::Strings(_) => None,
        }
    }

    /// Get the value as a vector of f64 (with numeric promotion).
    pub fn as_f64_vec(&self) -> Option<Vec<f64>> {
        match self {
            NcAttrValue::Bytes(v) => Some(v.iter().map(|&x| x as f64).collect()),
            NcAttrValue::Shorts(v) => Some(v.iter().map(|&x| x as f64).collect()),
            NcAttrValue::Ints(v) => Some(v.iter().map(|&x| x as f64).collect()),
            NcAttrValue::Floats(v) => Some(v.iter().map(|&x| x as f64).collect()),
            NcAttrValue::Doubles(v) => Some(v.clone()),
            NcAttrValue::UBytes(v) => Some(v.iter().map(|&x| x as f64).collect()),
            NcAttrValue::UShorts(v) => Some(v.iter().map(|&x| x as f64).collect()),
            NcAttrValue::UInts(v) => Some(v.iter().map(|&x| x as f64).collect()),
            NcAttrValue::Int64s(v) => Some(v.iter().map(|&x| x as f64).collect()),
            NcAttrValue::UInt64s(v) => Some(v.iter().map(|&x| x as f64).collect()),
            NcAttrValue::Chars(_) | NcAttrValue::Strings(_) => None,
        }
    }
}

/// A NetCDF attribute.
#[derive(Debug, Clone)]
pub struct NcAttribute {
    pub name: String,
    pub value: NcAttrValue,
}

/// A NetCDF variable (metadata only -- data is read on demand).
#[derive(Debug, Clone)]
pub struct NcVariable {
    pub name: String,
    pub dimensions: Vec<NcDimension>,
    pub dtype: NcType,
    pub attributes: Vec<NcAttribute>,
    /// For classic: file byte offset to the start of this variable's data.
    /// For nc4: HDF5 dataset object header address.
    pub(crate) data_offset: u64,
    /// Total data size in bytes (for non-record variables).
    pub(crate) _data_size: u64,
    /// Whether this variable uses the unlimited (record) dimension.
    pub(crate) is_record_var: bool,
    /// Size of one record slice in bytes (only meaningful for record variables).
    pub(crate) record_size: u64,
}

impl NcVariable {
    /// Variable name.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Variable dimensions.
    pub fn dimensions(&self) -> &[NcDimension] {
        &self.dimensions
    }

    /// Returns the dimension for a CF/NetCDF coordinate variable.
    ///
    /// A coordinate variable is one-dimensional and has the same name as its
    /// dimension. NetCDF-4 stores these as HDF5 dimension scales, but they are
    /// exposed here with the same shape as classic NetCDF coordinate variables.
    pub fn coordinate_dimension(&self) -> Option<&NcDimension> {
        match self.dimensions.as_slice() {
            [dim] if dim.name == self.name => Some(dim),
            _ => None,
        }
    }

    /// Returns true when this variable is a CF/NetCDF coordinate variable.
    pub fn is_coordinate_variable(&self) -> bool {
        self.coordinate_dimension().is_some()
    }

    /// Returns true when this variable is the coordinate variable for a named
    /// dimension.
    pub fn is_coordinate_variable_for(&self, dimension_name: &str) -> bool {
        self.coordinate_dimension()
            .is_some_and(|dim| dim.name == dimension_name)
    }

    /// Variable data type.
    pub fn dtype(&self) -> &NcType {
        &self.dtype
    }

    /// Shape of the variable as a vector of dimension sizes.
    pub fn shape(&self) -> Vec<u64> {
        self.dimensions.iter().map(|d| d.size).collect()
    }

    /// Variable attributes.
    pub fn attributes(&self) -> &[NcAttribute] {
        &self.attributes
    }

    /// Find an attribute by name.
    pub fn attribute(&self, name: &str) -> Option<&NcAttribute> {
        self.attributes.iter().find(|a| a.name == name)
    }

    /// Number of dimensions.
    pub fn ndim(&self) -> usize {
        self.dimensions.len()
    }

    /// Total number of elements.
    pub fn num_elements(&self) -> Result<u64> {
        match self.dimensions.as_slice() {
            [] => Ok(1), // scalar
            [dim] => Ok(dim.size),
            dimensions => {
                let mut total = 1u64;
                for dim in dimensions {
                    total = total.checked_mul(dim.size).ok_or_else(|| {
                        Error::InvalidData(
                            "NetCDF variable element count overflows u64".to_string(),
                        )
                    })?;
                }
                Ok(total)
            }
        }
    }
}

/// A NetCDF group (NetCDF-4 only; classic files have one implicit root group).
#[derive(Debug, Clone)]
pub struct NcGroup {
    pub name: String,
    pub dimensions: Vec<NcDimension>,
    pub variables: Vec<NcVariable>,
    pub attributes: Vec<NcAttribute>,
    pub groups: Vec<NcGroup>,
}

impl NcGroup {
    /// Find a variable by name in this group.
    pub fn variable(&self, name: &str) -> Option<&NcVariable> {
        let (group_path, variable_name) = split_parent_path(name)?;
        let group = self.group(group_path)?;
        group.variables.iter().find(|v| v.name == variable_name)
    }

    /// Find a dimension by name in this group.
    pub fn dimension(&self, name: &str) -> Option<&NcDimension> {
        let (group_path, dimension_name) = split_parent_path(name)?;
        let group = self.group(group_path)?;
        group.dimensions.iter().find(|d| d.name == dimension_name)
    }

    /// Find the coordinate variable for a dimension in this group.
    ///
    /// `name` may be a local dimension name or a path relative to this group,
    /// for example `time` or `forecast/time`.
    pub fn coordinate_variable(&self, name: &str) -> Option<&NcVariable> {
        let (group_path, dimension_name) = split_parent_path(name)?;
        let group = self.group(group_path)?;
        group
            .variables
            .iter()
            .find(|var| var.is_coordinate_variable_for(dimension_name))
    }

    /// Iterate over coordinate variables declared in this group.
    pub fn coordinate_variables(&self) -> impl Iterator<Item = &NcVariable> {
        self.variables
            .iter()
            .filter(|var| var.is_coordinate_variable())
    }

    /// Find an attribute by name in this group.
    pub fn attribute(&self, name: &str) -> Option<&NcAttribute> {
        let (group_path, attribute_name) = split_parent_path(name)?;
        let group = self.group(group_path)?;
        group.attributes.iter().find(|a| a.name == attribute_name)
    }

    /// Find a child group by relative path.
    pub fn group(&self, name: &str) -> Option<&NcGroup> {
        let trimmed = name.trim_matches('/');
        if trimmed.is_empty() {
            return Some(self);
        }

        let mut group = self;
        for component in trimmed.split('/').filter(|part| !part.is_empty()) {
            group = group.groups.iter().find(|child| child.name == component)?;
        }

        Some(group)
    }
}

fn split_parent_path(path: &str) -> Option<(&str, &str)> {
    let trimmed = path.trim_matches('/');
    if trimmed.is_empty() {
        return None;
    }

    match trimmed.rsplit_once('/') {
        Some((group_path, leaf_name)) if !leaf_name.is_empty() => Some((group_path, leaf_name)),
        Some(_) => None,
        None => Some(("", trimmed)),
    }
}

pub(crate) fn checked_usize_from_u64(value: u64, context: &str) -> crate::Result<usize> {
    usize::try_from(value)
        .map_err(|_| crate::Error::InvalidData(format!("{context} exceeds platform usize")))
}

pub(crate) fn checked_mul_u64(lhs: u64, rhs: u64, context: &str) -> crate::Result<u64> {
    lhs.checked_mul(rhs)
        .ok_or_else(|| crate::Error::InvalidData(format!("{context} exceeds u64 capacity")))
}

pub(crate) fn checked_shape_elements(shape: &[u64], context: &str) -> crate::Result<u64> {
    shape
        .iter()
        .try_fold(1u64, |acc, &dim| checked_mul_u64(acc, dim, context))
}

/// Hyperslab selection for reading slices of NetCDF variables.
///
/// Each element corresponds to one dimension of the variable.
#[derive(Debug, Clone)]
pub struct NcSliceInfo {
    pub selections: Vec<NcSliceInfoElem>,
}

/// A single dimension's selection within a hyperslab.
#[derive(Debug, Clone)]
pub enum NcSliceInfoElem {
    /// Select a single index (reduces dimensionality).
    Index(u64),
    /// Select a range with stride.
    Slice { start: u64, end: u64, step: u64 },
}

impl NcSliceInfo {
    /// Create a selection that reads everything for an `ndim`-dimensional variable.
    pub fn all(ndim: usize) -> Self {
        NcSliceInfo {
            selections: vec![
                NcSliceInfoElem::Slice {
                    start: 0,
                    end: u64::MAX,
                    step: 1,
                };
                ndim
            ],
        }
    }
}

#[cfg(feature = "netcdf4")]
impl NcSliceInfo {
    /// Convert to hdf5_reader::SliceInfo for NC4 delegation.
    pub(crate) fn to_hdf5_slice_info(&self) -> hdf5_reader::SliceInfo {
        hdf5_reader::SliceInfo {
            selections: self
                .selections
                .iter()
                .map(|s| match s {
                    NcSliceInfoElem::Index(idx) => hdf5_reader::SliceInfoElem::Index(*idx),
                    NcSliceInfoElem::Slice { start, end, step } => {
                        hdf5_reader::SliceInfoElem::Slice {
                            start: *start,
                            end: *end,
                            step: *step,
                        }
                    }
                })
                .collect(),
        }
    }
}

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

    fn sample_group_tree() -> NcGroup {
        NcGroup {
            name: "/".to_string(),
            dimensions: vec![NcDimension {
                name: "root_dim".to_string(),
                size: 2,
                is_unlimited: false,
            }],
            variables: vec![NcVariable {
                name: "root_var".to_string(),
                dimensions: vec![],
                dtype: NcType::Int,
                attributes: vec![],
                data_offset: 0,
                _data_size: 0,
                is_record_var: false,
                record_size: 4,
            }],
            attributes: vec![NcAttribute {
                name: "title".to_string(),
                value: NcAttrValue::Chars("root".to_string()),
            }],
            groups: vec![NcGroup {
                name: "obs".to_string(),
                dimensions: vec![NcDimension {
                    name: "time".to_string(),
                    size: 3,
                    is_unlimited: false,
                }],
                variables: vec![NcVariable {
                    name: "temperature".to_string(),
                    dimensions: vec![],
                    dtype: NcType::Float,
                    attributes: vec![],
                    data_offset: 0,
                    _data_size: 0,
                    is_record_var: false,
                    record_size: 4,
                }],
                attributes: vec![],
                groups: vec![NcGroup {
                    name: "surface".to_string(),
                    dimensions: vec![],
                    variables: vec![NcVariable {
                        name: "pressure".to_string(),
                        dimensions: vec![],
                        dtype: NcType::Double,
                        attributes: vec![],
                        data_offset: 0,
                        _data_size: 0,
                        is_record_var: false,
                        record_size: 8,
                    }],
                    attributes: vec![NcAttribute {
                        name: "units".to_string(),
                        value: NcAttrValue::Chars("hPa".to_string()),
                    }],
                    groups: vec![],
                }],
            }],
        }
    }

    #[test]
    fn group_path_lookup() {
        let root = sample_group_tree();

        let surface = root.group("obs/surface").unwrap();
        assert_eq!(surface.name, "surface");
        assert!(root.group("/obs/surface").is_some());
        assert!(root.group("missing").is_none());
    }

    #[test]
    fn variable_path_lookup() {
        let root = sample_group_tree();

        assert_eq!(root.variable("root_var").unwrap().name(), "root_var");
        assert_eq!(
            root.variable("obs/temperature").unwrap().dtype(),
            &NcType::Float
        );
        assert_eq!(
            root.variable("/obs/surface/pressure").unwrap().dtype(),
            &NcType::Double
        );
        assert!(root.variable("pressure").is_none());
    }

    #[test]
    fn dimension_and_attribute_path_lookup() {
        let root = sample_group_tree();

        assert_eq!(root.dimension("root_dim").unwrap().size, 2);
        assert_eq!(root.dimension("obs/time").unwrap().size, 3);
        assert_eq!(
            root.attribute("title").unwrap().value.as_string().unwrap(),
            "root"
        );
        assert_eq!(
            root.attribute("obs/surface/units")
                .unwrap()
                .value
                .as_string()
                .unwrap(),
            "hPa"
        );
    }

    #[test]
    fn coordinate_variable_detection_and_lookup() {
        let time_dim = NcDimension {
            name: "time".to_string(),
            size: 3,
            is_unlimited: false,
        };
        let lat_dim = NcDimension {
            name: "lat".to_string(),
            size: 2,
            is_unlimited: false,
        };
        let time = NcVariable {
            name: "time".to_string(),
            dimensions: vec![time_dim.clone()],
            dtype: NcType::Double,
            attributes: vec![],
            data_offset: 0,
            _data_size: 0,
            is_record_var: false,
            record_size: 8,
        };
        let temperature = NcVariable {
            name: "temperature".to_string(),
            dimensions: vec![time_dim.clone(), lat_dim.clone()],
            dtype: NcType::Float,
            attributes: vec![],
            data_offset: 0,
            _data_size: 0,
            is_record_var: false,
            record_size: 4,
        };
        let group = NcGroup {
            name: "/".to_string(),
            dimensions: vec![time_dim, lat_dim],
            variables: vec![time.clone(), temperature],
            attributes: vec![],
            groups: vec![],
        };

        assert!(time.is_coordinate_variable());
        assert_eq!(time.coordinate_dimension().unwrap().name, "time");
        assert_eq!(group.coordinate_variable("time").unwrap().name(), "time");
        assert!(group.coordinate_variable("lat").is_none());

        let names: Vec<&str> = group.coordinate_variables().map(NcVariable::name).collect();
        assert_eq!(names, vec!["time"]);
    }

    #[test]
    fn checked_shape_elements_overflow() {
        let err = checked_shape_elements(&[u64::MAX, 2], "test overflow").unwrap_err();
        assert!(matches!(err, crate::Error::InvalidData(_)));
    }

    #[test]
    fn array_type_size_rejects_overflow() {
        let ty = NcType::Array {
            base: Box::new(NcType::UInt64),
            dims: vec![u64::MAX, 2],
        };

        let err = ty.size().unwrap_err();
        assert!(err.to_string().contains("array type"));
    }

    #[test]
    fn variable_num_elements_rejects_overflow() {
        let var = NcVariable {
            name: "huge".to_string(),
            dimensions: vec![
                NcDimension {
                    name: "x".to_string(),
                    size: u64::MAX,
                    is_unlimited: false,
                },
                NcDimension {
                    name: "y".to_string(),
                    size: 2,
                    is_unlimited: false,
                },
            ],
            dtype: NcType::Float,
            attributes: vec![],
            data_offset: 0,
            _data_size: 0,
            is_record_var: false,
            record_size: 4,
        };

        let err = var.num_elements().unwrap_err();
        assert!(err.to_string().contains("element count"));
    }
}