clawhdf5/crates/clawhdf5-tools/src/stat.rs

//! `h5rs stat`: object, layout, filter and storage statistics, laid out like
//! h5stat's report.

use std::collections::BTreeMap;

use clawhdf5_format::data_layout::DataLayout;
use clawhdf5_format::dataspace::DataspaceType;

use crate::cli::{Args, Out};
use crate::h5::{H5, Kind, LinkKind};
use crate::info::{self, DsInfo};

pub const USAGE: &str = "\
usage: h5rs stat FILE

Print statistics for FILE, in the layout of h5stat's report: object counts,
links, dataset ranks, layouts, filters, attribute counts, raw-data size and
a file-space summary. Metadata space is not broken down by structure (h5stat
does); it is reported as one figure together with free space.

Exit status: 0 printed, 1 some object could not be read (counted as
\"unreadable\"), 2 error.";

#[derive(Default)]
struct Stats {
    groups: u64,
    datasets: u64,
    datatypes: u64,
    other: u64,
    unreadable: u64,
    links: u64,
    max_links_to_object: u64,
    max_objects_in_group: u64,
    group_sizes: BTreeMap<usize, u64>,
    ranks: BTreeMap<usize, u64>,
    max_1d: u64,
    layout: [u64; 4],
    external: u64,
    no_filter: u64,
    filter_counts: BTreeMap<&'static str, u64>,
    raw_data: u64,
    raw_errors: u64,
    attr_objects: u64,
    max_attrs: u64,
    attr_counts: BTreeMap<usize, u64>,
}

const FILTER_KINDS: [&str; 8] = [
    "GZIP",
    "SHUFFLE",
    "FLETCHER32",
    "SZIP",
    "NBIT",
    "SCALEOFFSET",
    "USER-DEFINED",
    "",
];

fn filter_kind(id: u16) -> &'static str {
    match id {
        1 => "GZIP",
        2 => "SHUFFLE",
        3 => "FLETCHER32",
        4 => "SZIP",
        5 => "NBIT",
        6 => "SCALEOFFSET",
        _ => "USER-DEFINED",
    }
}

pub fn run(args: &mut Args, out: &mut Out) -> std::io::Result<i32> {
    let mut file = None;
    while let Some(a) = args.next() {
        match a.as_str() {
            "-h" | "--help" => {
                writeln!(out.o, "{USAGE}")?;
                return Ok(0);
            }
            s if s.starts_with('-') && s.len() > 1 => {
                return args.usage_error(out, &format!("unknown option {a}"), USAGE);
            }
            _ if file.is_none() => file = Some(a),
            _ => return args.usage_error(out, &format!("unexpected argument {a}"), USAGE),
        }
    }
    let Some(file) = file else {
        return args.usage_error(out, "missing FILE", USAGE);
    };
    let h5 = match H5::open(std::path::Path::new(&file)) {
        Ok(h) => h,
        Err(e) => {
            writeln!(out.e, "h5rs stat: {e}")?;
            return Ok(2);
        }
    };
    let mut s = Stats::default();
    let mut errors: Vec<String> = Vec::new();
    let walk = h5.walk(|it| {
        if let Some(l) = it.link
            && !matches!(l.kind, LinkKind::Hard(_))
        {
            s.links += 1;
            return;
        }
        if it.first_path.is_some() {
            return;
        }
        let Some(Ok(h)) = it.header else {
            if let Some(Err(e)) = it.header {
                s.unreadable += 1;
                errors.push(format!("{}: {e}", it.path));
            }
            return;
        };
        s.max_links_to_object = s.max_links_to_object.max(info::link_count(h));
        match h5.attributes(h) {
            Ok((a, errs)) => {
                let n = a.len() + errs.len();
                if n > 0 {
                    s.attr_objects += 1;
                    s.max_attrs = s.max_attrs.max(n as u64);
                    *s.attr_counts.entry(n).or_default() += 1;
                }
                for e in errs {
                    errors.push(format!("{}: attribute: {e}", it.path));
                }
            }
            Err(e) => errors.push(format!("{}: {e}", it.path)),
        }
        match Kind::of(h) {
            Kind::Group => {
                s.groups += 1;
                match h5.links(h) {
                    Ok(l) => {
                        s.max_objects_in_group = s.max_objects_in_group.max(l.len() as u64);
                        *s.group_sizes.entry(l.len()).or_default() += 1;
                    }
                    Err(e) => {
                        s.unreadable += 1;
                        errors.push(format!("{}: {e}", it.path));
                    }
                }
            }
            Kind::Datatype => s.datatypes += 1,
            Kind::Unknown => {
                // The root group of a file with no links at all has no
                // group messages; count it as a group.
                if it.link.is_none() {
                    s.groups += 1;
                    *s.group_sizes.entry(0).or_default() += 1;
                } else {
                    s.other += 1;
                }
            }
            Kind::Dataset => {
                s.datasets += 1;
                let info = DsInfo::read(&h5, it.path, h);
                if let Ok(ds) = &info.ds {
                    let rank = match ds.space_type {
                        DataspaceType::Simple => ds.dimensions.len(),
                        _ => 0,
                    };
                    *s.ranks.entry(rank).or_default() += 1;
                    if rank == 1 {
                        s.max_1d = s.max_1d.max(ds.dimensions[0]);
                    }
                }
                match &info.layout {
                    Ok(l) => {
                        let i = match l {
                            DataLayout::Compact { .. } => 0,
                            DataLayout::Contiguous { .. } => 1,
                            DataLayout::Chunked { .. } => 2,
                            DataLayout::Virtual { .. } => 3,
                        };
                        s.layout[i] += 1;
                    }
                    Err(e) => errors.push(format!("{}: {e}", it.path)),
                }
                if info.external {
                    s.external += 1;
                }
                match &info.filters {
                    Ok(Some(p)) if !p.filters.is_empty() => {
                        let mut kinds: Vec<&'static str> =
                            p.filters.iter().map(|f| filter_kind(f.filter_id)).collect();
                        kinds.sort_unstable();
                        kinds.dedup();
                        for k in kinds {
                            *s.filter_counts.entry(k).or_default() += 1;
                        }
                    }
                    Ok(_) => s.no_filter += 1,
                    Err(e) => errors.push(format!("{}: {e}", it.path)),
                }
                match info::allocated_bytes(&h5, &info) {
                    Ok(b) => s.raw_data = s.raw_data.saturating_add(b),
                    Err(e) => {
                        s.raw_errors += 1;
                        errors.push(format!("{}: storage size: {e}", it.path));
                    }
                }
            }
        }
    });
    if let Err(e) = walk {
        errors.push(e.to_string());
    }
    report(&h5, &file, &s, out)?;
    for e in &errors {
        writeln!(out.e, "h5rs stat: {e}")?;
    }
    Ok(if errors.is_empty() { 0 } else { 1 })
}

fn report(h5: &H5, file: &str, s: &Stats, out: &mut Out) -> std::io::Result<()> {
    let o = &mut out.o;
    writeln!(o, "Filename: {file}")?;
    writeln!(o, "File information")?;
    writeln!(o, "\t# of unique groups: {}", s.groups)?;
    writeln!(o, "\t# of unique datasets: {}", s.datasets)?;
    writeln!(o, "\t# of unique named datatypes: {}", s.datatypes)?;
    writeln!(o, "\t# of unique links: {}", s.links)?;
    writeln!(o, "\t# of unique other: {}", s.other)?;
    if s.unreadable > 0 {
        writeln!(o, "\t# of unreadable objects: {}", s.unreadable)?;
    }
    writeln!(o, "\tMax. # of links to object: {}", s.max_links_to_object)?;
    writeln!(
        o,
        "\tMax. # of objects in group: {}",
        s.max_objects_in_group
    )?;
    let sb = h5.sb();
    writeln!(o, "Superblock:")?;
    writeln!(o, "\tVersion: {}", sb.version)?;
    writeln!(o, "\tSize of offsets: {} bytes", sb.offset_size)?;
    writeln!(o, "\tSize of lengths: {} bytes", sb.length_size)?;
    writeln!(o, "\tUser block: {} bytes", h5.file.user_block_size())?;
    writeln!(o, "Group bins:")?;
    for (n, c) in &s.group_sizes {
        writeln!(o, "\t# of groups with {n} link(s): {c}")?;
    }
    writeln!(o, "\tTotal # of groups: {}", s.groups)?;
    writeln!(o, "Dataset dimension information:")?;
    writeln!(
        o,
        "\tMax. rank of datasets: {}",
        s.ranks.keys().next_back().copied().unwrap_or(0)
    )?;
    writeln!(o, "\tDataset ranks:")?;
    for (r, c) in &s.ranks {
        writeln!(o, "\t\t# of dataset with rank {r}: {c}")?;
    }
    writeln!(o, "1-D Dataset information:")?;
    writeln!(o, "\tMax. dimension size of 1-D datasets: {}", s.max_1d)?;
    writeln!(o, "Dataset storage information:")?;
    writeln!(o, "\tTotal raw data size: {}", s.raw_data)?;
    if s.raw_errors > 0 {
        writeln!(
            o,
            "\tDatasets whose storage size could not be read: {}",
            s.raw_errors
        )?;
    }
    writeln!(o, "Dataset layout information:")?;
    for (i, n) in ["COMPACT", "CONTIG", "CHUNKED", "VIRTUAL"]
        .iter()
        .enumerate()
    {
        writeln!(o, "\tDataset layout counts[{n}]: {}", s.layout[i])?;
    }
    writeln!(o, "\tDatasets with external raw data: {}", s.external)?;
    writeln!(o, "Dataset filters information:")?;
    writeln!(o, "\tNumber of datasets with:")?;
    writeln!(o, "\t\tNO filter: {}", s.no_filter)?;
    for k in FILTER_KINDS.iter().filter(|k| !k.is_empty()) {
        writeln!(
            o,
            "\t\t{k} filter: {}",
            s.filter_counts.get(k).copied().unwrap_or(0)
        )?;
    }
    writeln!(o, "Attribute information:")?;
    for (n, c) in &s.attr_counts {
        writeln!(o, "\t# of objects with {n} attribute(s): {c}")?;
    }
    writeln!(
        o,
        "\tTotal # of objects with attributes: {}",
        s.attr_objects
    )?;
    writeln!(o, "\tMax. # of attributes to objects: {}", s.max_attrs)?;
    let total = std::fs::metadata(&h5.path).map(|m| m.len()).unwrap_or(0);
    let ub = h5.file.user_block_size();
    writeln!(o, "Summary of file space information:")?;
    writeln!(o, "  User block: {ub} bytes")?;
    writeln!(o, "  Raw data: {} bytes", s.raw_data)?;
    writeln!(
        o,
        "  Metadata and free space: {} bytes",
        total.saturating_sub(s.raw_data).saturating_sub(ub)
    )?;
    writeln!(o, "Total space: {total} bytes")
}