conspire 0.7.7

The Rust interface to conspire.
Documentation
mod checksum;
mod chunk;
mod message;
#[cfg(test)]
mod test;

use crate::io::netcdf::format::{
    AttValue, Attribute, DimSpec, NC_DOUBLE, NC_FLOAT, NC_INT, VarSpec,
};
use checksum::jenkins;
use message::*;
use std::io::{Result as IoResult, Write};

const DEFLATE_LEVEL: u32 = 4;

fn elem_size(xtype: i32) -> usize {
    match xtype {
        NC_INT | NC_FLOAT => 4,
        NC_DOUBLE => 8,
        other => panic!("cannot write netCDF external type {other} to HDF5"),
    }
}

enum Storage {
    Contiguous(Vec<u8>),
    Chunked(chunk::Plan),
}

struct VarPlan {
    dims: Vec<u64>,
    storage: Storage,
}

fn plan_var(var: &VarSpec, dim_lens: &[u64], data: Vec<u8>, threads: usize) -> VarPlan {
    let dims: Vec<u64> = var.dimids.iter().map(|&d| dim_lens[d]).collect();
    let storage = if dims.is_empty() {
        Storage::Contiguous(data)
    } else {
        Storage::Chunked(chunk::plan(&dims, &data, elem_size(var.xtype), threads))
    };
    VarPlan { dims, storage }
}

fn attribute_message(att: &Attribute) -> Vec<u8> {
    match &att.value {
        AttValue::Text(s) => attr_text(&att.name, s),
        AttValue::Int(v) => attr_i32(&att.name, v),
        AttValue::Float(v) => attr_f32(&att.name, v),
    }
}

fn dim_scale_messages(index: usize, len: u64) -> Vec<(u8, Vec<u8>)> {
    vec![
        (MSG_DATASPACE, dataspace(&[len])),
        (MSG_DATATYPE, datatype(NC_INT)),
        (MSG_FILL, fill_default()),
        (MSG_LAYOUT, layout_unallocated(len * 4)),
        (MSG_ATTRIBUTE, attr_dimension_scale()),
        (MSG_ATTRIBUTE, attr_pure_dimension_name(len)),
        (MSG_ATTRIBUTE, attr_i32("_Netcdf4Dimid", &[index as i32])),
    ]
}

fn var_messages(
    var: &VarSpec,
    plan: &VarPlan,
    btree_addr: u64,
    data_addr: u64,
) -> Vec<(u8, Vec<u8>)> {
    let mut msgs = vec![
        (MSG_DATASPACE, dataspace(&plan.dims)),
        (MSG_DATATYPE, datatype(var.xtype)),
        (MSG_FILL, fill_value(var.xtype)),
    ];
    match &plan.storage {
        Storage::Contiguous(bytes) => {
            msgs.push((MSG_LAYOUT, layout_contiguous(data_addr, bytes.len() as u64)))
        }
        Storage::Chunked(p) => {
            msgs.push((
                MSG_LAYOUT,
                layout_chunked(btree_addr, &p.chunk_dims, p.elem),
            ));
            msgs.push((
                MSG_FILTER,
                filter_pipeline(elem_size(var.xtype) as u32, DEFLATE_LEVEL),
            ));
        }
    }
    if !var.dimids.is_empty() {
        let dimids: Vec<i32> = var.dimids.iter().map(|&d| d as i32).collect();
        msgs.push((MSG_ATTRIBUTE, attr_i32("_Netcdf4Coordinates", &dimids)));
    }
    for att in &var.atts {
        msgs.push((MSG_ATTRIBUTE, attribute_message(att)));
    }
    msgs
}

fn object_header(messages: &[(u8, Vec<u8>)]) -> Vec<u8> {
    let chunk0: usize = messages.iter().map(|(_, b)| 4 + b.len()).sum();
    let mut v = b"OHDR".to_vec();
    v.push(2);
    v.push(0x02);
    v.extend_from_slice(&(chunk0 as u32).to_le_bytes());
    for (typ, body) in messages {
        let size = u16::try_from(body.len()).expect("object header message exceeds 64 KiB");
        v.push(*typ);
        v.extend_from_slice(&size.to_le_bytes());
        v.push(0);
        v.extend_from_slice(body);
    }
    let sum = jenkins(&v);
    v.extend_from_slice(&sum.to_le_bytes());
    v
}

fn object_header_len(messages: &[(u8, Vec<u8>)]) -> usize {
    10 + messages.iter().map(|(_, b)| 4 + b.len()).sum::<usize>() + 4
}

fn superblock(root_oh: u64, eof: u64) -> Vec<u8> {
    let mut v = vec![
        0x89, b'H', b'D', b'F', b'\r', b'\n', 0x1a, b'\n', 2, 8, 8, 0,
    ];
    v.extend_from_slice(&0u64.to_le_bytes());
    v.extend_from_slice(&u64::MAX.to_le_bytes());
    v.extend_from_slice(&eof.to_le_bytes());
    v.extend_from_slice(&root_oh.to_le_bytes());
    let sum = jenkins(&v);
    v.extend_from_slice(&sum.to_le_bytes());
    v
}

fn emit<W: Write>(out: &mut W, written: &mut u64, bytes: &[u8]) -> IoResult<()> {
    *written += bytes.len() as u64;
    out.write_all(bytes)
}

pub(in crate::io::netcdf) fn write<W: Write>(
    dims: &[DimSpec],
    global: &[Attribute],
    vars: &[VarSpec],
    data: Vec<Vec<u8>>,
    threads: usize,
    out: &mut W,
) -> IoResult<()> {
    let dim_lens: Vec<u64> = dims.iter().map(|d| d.len).collect();
    let plans: Vec<VarPlan> = vars
        .iter()
        .zip(data)
        .map(|(v, d)| plan_var(v, &dim_lens, d, threads))
        .collect();

    let dim_scale_msgs: Vec<Vec<(u8, Vec<u8>)>> = (0..dims.len())
        .map(|i| dim_scale_messages(i, dim_lens[i]))
        .collect();

    let placeholder_var_msgs: Vec<Vec<(u8, Vec<u8>)>> = vars
        .iter()
        .zip(&plans)
        .map(|(v, p)| var_messages(v, p, 0, 0))
        .collect();

    const SUPERBLOCK: u64 = 48;
    let mut cursor = SUPERBLOCK;
    let root_oh = cursor;

    let root_placeholder = root_messages(
        dims,
        vars,
        global,
        &vec![0u64; dims.len()],
        &vec![0u64; vars.len()],
    );
    cursor += object_header_len(&root_placeholder) as u64;

    let mut dim_oh = Vec::with_capacity(dims.len());
    for msgs in &dim_scale_msgs {
        dim_oh.push(cursor);
        cursor += object_header_len(msgs) as u64;
    }
    let mut var_oh = Vec::with_capacity(vars.len());
    for msgs in &placeholder_var_msgs {
        var_oh.push(cursor);
        cursor += object_header_len(msgs) as u64;
    }

    let mut var_btree = vec![0u64; vars.len()];
    let mut var_data = vec![0u64; vars.len()];
    let mut var_chunk_addrs: Vec<Vec<u64>> = vec![Vec::new(); vars.len()];
    for (i, plan) in plans.iter().enumerate() {
        match &plan.storage {
            Storage::Contiguous(bytes) => {
                var_data[i] = cursor;
                cursor += bytes.len() as u64;
            }
            Storage::Chunked(p) => {
                var_btree[i] = cursor;
                cursor += chunk::btree_node(p, &vec![0u64; p.blobs.len()]).len() as u64;
                for blob in &p.blobs {
                    var_chunk_addrs[i].push(cursor);
                    cursor += blob.bytes.len() as u64;
                }
            }
        }
    }
    let eof = cursor;

    let mut written = 0u64;
    emit(out, &mut written, &superblock(root_oh, eof))?;
    emit(
        out,
        &mut written,
        &object_header(&root_messages(dims, vars, global, &dim_oh, &var_oh)),
    )?;
    for msgs in &dim_scale_msgs {
        emit(out, &mut written, &object_header(msgs))?;
    }
    for (i, (var, plan)) in vars.iter().zip(&plans).enumerate() {
        emit(
            out,
            &mut written,
            &object_header(&var_messages(var, plan, var_btree[i], var_data[i])),
        )?;
    }
    for (i, plan) in plans.into_iter().enumerate() {
        match plan.storage {
            Storage::Contiguous(bytes) => emit(out, &mut written, &bytes)?,
            Storage::Chunked(p) => {
                emit(
                    out,
                    &mut written,
                    &chunk::btree_node(&p, &var_chunk_addrs[i]),
                )?;
                for blob in &p.blobs {
                    emit(out, &mut written, &blob.bytes)?;
                }
            }
        }
    }
    debug_assert_eq!(written, eof);
    Ok(())
}

fn root_messages(
    dims: &[DimSpec],
    vars: &[VarSpec],
    global: &[Attribute],
    dim_oh: &[u64],
    var_oh: &[u64],
) -> Vec<(u8, Vec<u8>)> {
    let mut msgs = vec![(MSG_LINK_INFO, link_info()), (MSG_GROUP_INFO, group_info())];
    for (d, &addr) in dims.iter().zip(dim_oh) {
        msgs.push((MSG_LINK, link(&d.name, addr)));
    }
    for (v, &addr) in vars.iter().zip(var_oh) {
        msgs.push((MSG_LINK, link(&v.name, addr)));
    }
    for att in global {
        msgs.push((MSG_ATTRIBUTE, attribute_message(att)));
    }
    msgs
}