gwseq-io 0.2.0

//! HiC header and footer.
//!
//! v8 and v9 differ in the width of nearly every
//! count and value here, and the version read out of the magic decides all of
//! them.

use indexmap::IndexMap;

use crate::error::{Error, Result};
use crate::genomic::ChrMap;
use crate::source::ByteSource;

use super::Unit;

/// How much to reserve for a count the file declares, before anything has been
/// read with it.
///
/// A reservation is not a read, so none of the checks the reads carry apply to
/// it, and a corrupt count asks the allocator for its full size before the first
/// of them runs. Capped rather than rejected: the lists this guards are small,
/// and a genuine one grows past the cap in a single reallocation.
const MAX_RESERVE: usize = 4096;

fn clamp_reserve(count: i64) -> usize {
    count.clamp(0, MAX_RESERVE as i64) as usize
}

#[derive(Debug, Clone)]
pub struct HiCHeader {
    pub version: i64,
    pub footer_position: u64,
    pub genome_id: String,
    /// Where the normalization vector index sits, and how long it is. Only v9
    /// carries them; -1 stands for a file that does not.
    pub nvi_position: i64,
    pub nvi_length: i64,
    pub attributes: IndexMap<String, String>,
    pub chr_map: ChrMap,
    pub bp_resolutions: Vec<i64>,
    pub frag_resolutions: Vec<i64>,
    /// Restriction-site positions of each chromosome, keyed by its id and in
    /// file order. Empty unless the file carries fragment resolutions.
    pub sites: IndexMap<String, Vec<i64>>,
}

impl HiCHeader {
    /// The resolutions the file was built at, for one unit.
    pub fn resolutions(&self, unit: Unit) -> &[i64] {
        match unit {
            Unit::Bp => &self.bp_resolutions,
            Unit::Frag => &self.frag_resolutions,
        }
    }
}

/// One entry of the master index: where a chromosome pair's matrix lives.
#[derive(Debug, Clone, Copy)]
pub struct HiCIndexItem {
    pub position: u64,
    pub size: i64,
}

#[derive(Debug, Clone)]
pub struct ExpectedValueVector {
    pub normalization: String,
    pub unit: String,
    pub bin_size: i64,
    pub values: Vec<f32>,
    pub chr_scale_factors: IndexMap<i64, f32>,
}

#[derive(Debug, Clone)]
pub struct NormalizationVector {
    pub normalization: String,
    pub chr_index: i64,
    pub unit: String,
    pub bin_size: i64,
    pub position: u64,
    pub byte_count: i64,
}

#[derive(Debug, Clone, Default)]
pub struct HiCFooter {
    /// The `nBytesV5` field the footer opens with: how much of the file the
    /// matrix section takes. Read but not acted on, and handed out for
    /// callers that want it.
    pub byte_count_v5: i64,
    pub master_index: IndexMap<String, HiCIndexItem>,
    pub expected_value_vectors: IndexMap<String, ExpectedValueVector>,
    pub normalization_vectors: IndexMap<String, NormalizationVector>,
    /// Sorted, so the reported list is stable.
    pub normalizations: Vec<String>,
    pub units: Vec<String>,
}

/// The key a vector is stored under. Built the same way on both sides, so a
/// lookup and an insertion cannot spell it differently.
pub fn vector_key(normalization: &str, bin_size: i64, unit: &str, chr: Option<i64>) -> String {
    let base = format!("normalization={normalization}|bin_size={bin_size}|unit={unit}");
    match chr {
        Some(index) => format!("chr_index={index}|{base}"),
        None => base,
    }
}

/// A sequential reader over a source, for the header and footer — both are
/// runs of variable-length fields with no offsets to seek by.
struct Cursor<'a> {
    source: &'a dyn ByteSource,
    offset: u64,
    buffer: bytes::Bytes,
    consumed: usize,
}

impl<'a> Cursor<'a> {
    fn new(source: &'a dyn ByteSource, offset: u64) -> Self {
        Self {
            source,
            offset,
            buffer: bytes::Bytes::new(),
            consumed: 0,
        }
    }

    /// Make at least `wanted` unread bytes available.
    ///
    /// Re-anchoring and reading are both done to `self`, and an early return
    /// between them would leave the cursor describing a buffer it no longer
    /// holds — which is exactly the bug this had: an error returned after
    /// `offset` had advanced but before `buffer` was replaced sent the cursor
    /// back to the start of the file, and the header parsed itself again from
    /// its own magic. So the buffer is trimmed *first*, and every path after
    /// that leaves the cursor consistent whether the read succeeds or not.
    fn fill(&mut self, wanted: usize) -> Result<()> {
        if self.buffer.len() - self.consumed >= wanted {
            return Ok(());
        }
        const CHUNK: usize = 1 << 16;
        self.buffer = self.buffer.slice(self.consumed..);
        self.offset += self.consumed as u64;
        self.consumed = 0;

        while self.buffer.len() < wanted {
            let at = self.offset + self.buffer.len() as u64;
            let more = self
                .source
                .read_at(at, CHUNK.max(wanted - self.buffer.len()))?;
            if more.is_empty() {
                return Err(Error::corrupt(
                    self.source.path(),
                    at,
                    "hic header ended early",
                ));
            }
            let mut joined = bytes::BytesMut::with_capacity(self.buffer.len() + more.len());
            joined.extend_from_slice(&self.buffer);
            joined.extend_from_slice(&more);
            self.buffer = joined.freeze();
        }
        Ok(())
    }

    fn take(&mut self, n: usize) -> Result<bytes::Bytes> {
        self.fill(n)?;
        let out = self.buffer.slice(self.consumed..self.consumed + n);
        self.consumed += n;
        Ok(out)
    }

    fn i32(&mut self) -> Result<i32> {
        let b = self.take(4)?;
        Ok(i32::from_le_bytes([b[0], b[1], b[2], b[3]]))
    }
    fn i64(&mut self) -> Result<i64> {
        let b = self.take(8)?;
        Ok(i64::from_le_bytes([
            b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7],
        ]))
    }
    fn f32(&mut self) -> Result<f32> {
        let b = self.take(4)?;
        Ok(f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
    }
    fn f64(&mut self) -> Result<f64> {
        let b = self.take(8)?;
        Ok(f64::from_le_bytes([
            b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7],
        ]))
    }

    /// A NUL-terminated string, which is how the format writes every name.
    fn cstr(&mut self) -> Result<String> {
        let mut wanted = 64usize;
        loop {
            // A short read is not a failure here — the terminator may already be
            // in what is buffered — and `fill` leaves the cursor consistent
            // either way, so its error is deliberately dropped.
            let _ = self.fill(wanted);
            let rest = &self.buffer[self.consumed..];
            if let Some(at) = memchr::memchr(0, rest) {
                let out = String::from_utf8_lossy(&rest[..at]).into_owned();
                self.consumed += at + 1;
                return Ok(out);
            }
            if rest.len() < wanted {
                return Err(Error::corrupt(
                    self.source.path(),
                    self.offset,
                    "hic header string is not NUL-terminated",
                ));
            }
            wanted *= 2;
        }
    }
}

pub fn read_header(source: &dyn ByteSource) -> Result<HiCHeader> {
    let path = source.path();
    let mut c = Cursor::new(source, 0);

    let magic = c.take(4)?;
    if &magic[..3] != b"HIC" {
        return Err(Error::format(
            path,
            format!(
                "not a hic file (magic: '{}')",
                String::from_utf8_lossy(&magic[..3])
            ),
        ));
    }
    let version = c.i32()? as i64;
    // Bounded above as well as below: a layout past 9 does not exist yet, and
    // reading one as a 9 fails somewhere deep in the footer instead of here.
    if !(6..=9).contains(&version) {
        return Err(Error::format(
            path,
            format!("hic version {version} unsupported (6 to 9)"),
        ));
    }
    let footer_position = c.i64()? as u64;
    let genome_id = c.cstr()?;
    // The normalization vector index. Only v9 carries one, and -1 is what a
    // file without it reports rather than a position of zero.
    let (nvi_position, nvi_length) = if version > 8 {
        (c.i64()?, c.i64()?)
    } else {
        (-1, -1)
    };

    let attribute_count = c.i32()?;
    let mut attributes = IndexMap::with_capacity(clamp_reserve(attribute_count as i64));
    for _ in 0..attribute_count.max(0) {
        let key = c.cstr()?;
        attributes.insert(key, c.cstr()?);
    }

    let chr_count = c.i32()?;
    let mut chrs = Vec::with_capacity(clamp_reserve(chr_count as i64));
    for index in 0..chr_count.max(0) {
        let id = c.cstr()?;
        let size = if version > 8 {
            c.i64()?
        } else {
            c.i32()? as i64
        };
        chrs.push((id, size, index as usize));
    }

    let bp_count = c.i32()?;
    let mut bp_resolutions = Vec::with_capacity(clamp_reserve(bp_count as i64));
    for _ in 0..bp_count.max(0) {
        bp_resolutions.push(c.i32()? as i64);
    }
    let frag_count = c.i32()?;
    let mut frag_resolutions = Vec::with_capacity(clamp_reserve(frag_count as i64));
    for _ in 0..frag_count.max(0) {
        frag_resolutions.push(c.i32()? as i64);
    }
    // One list of site positions per chromosome, in the order the chromosomes
    // were listed above, and the whole section **absent** from a file carrying
    // no fragment resolutions.
    //
    // Absent means absent: reading a count regardless takes the body's first
    // four bytes for one, which passes unnoticed only while the body opens on
    // the matrix of chromosome 0 and those bytes are its index, zero. Laid out
    // any other way, the header fails to read or comes back with a site that is
    // not one.
    //
    // Keyed off the chromosome list rather than the map, since two chromosomes
    // sharing a name would collapse into one entry of a map keyed by it and
    // leave the remaining lists read against the wrong chromosomes.
    let mut sites = IndexMap::new();
    if !frag_resolutions.is_empty() {
        for (id, _, _) in &chrs {
            let count = c.i32()?;
            let mut chr_sites = Vec::with_capacity(clamp_reserve(count as i64));
            for _ in 0..count.max(0) {
                chr_sites.push(c.i32()? as i64);
            }
            sites.insert(id.clone(), chr_sites);
        }
    }

    Ok(HiCHeader {
        version,
        footer_position,
        genome_id,
        nvi_position,
        nvi_length,
        attributes,
        chr_map: ChrMap::from_indexed_entries(chrs),
        bp_resolutions,
        frag_resolutions,
        sites,
    })
}

pub fn read_footer(source: &dyn ByteSource, header: &HiCHeader) -> Result<HiCFooter> {
    let version = header.version;
    let mut c = Cursor::new(source, header.footer_position);

    // Read before the footer exists rather than assigned into a default one:
    // it is the first field on the wire, and a `Default` that is immediately
    // overwritten reads as if it had a meaningful zero.
    let byte_count_v5 = if version > 8 {
        c.i64()?
    } else {
        c.i32()? as i64
    };
    let mut footer = HiCFooter {
        byte_count_v5,
        ..Default::default()
    };

    let master_count = c.i32()?;
    for _ in 0..master_count.max(0) {
        let key = c.cstr()?;
        let position = c.i64()? as u64;
        let size = c.i32()? as i64;
        footer
            .master_index
            .insert(key, HiCIndexItem { position, size });
    }

    let mut normalizations: Vec<String> = Vec::new();
    let mut units: Vec<String> = Vec::new();
    let note = |set: &mut Vec<String>, value: &str| {
        if !set.iter().any(|v| v == value) {
            set.push(value.to_string());
        }
    };

    // Two runs: the unnormalized expected values, then the normalized ones.
    for is_normalized in [false, true] {
        let count = c.i32()?;
        for _ in 0..count.max(0) {
            let normalization = if is_normalized {
                c.cstr()?.to_ascii_lowercase()
            } else {
                "none".to_string()
            };
            note(&mut normalizations, &normalization);
            let unit = c.cstr()?.to_ascii_lowercase();
            note(&mut units, &unit);
            let bin_size = c.i32()? as i64;

            let values = if version > 8 {
                let n = c.i64()?;
                let mut values = Vec::with_capacity(clamp_reserve(n));
                for _ in 0..n.max(0) {
                    values.push(c.f32()?);
                }
                values
            } else {
                let n = c.i32()? as i64;
                let mut values = Vec::with_capacity(clamp_reserve(n));
                for _ in 0..n.max(0) {
                    values.push(c.f64()? as f32);
                }
                values
            };

            let factor_count = c.i32()?;
            let mut chr_scale_factors = IndexMap::with_capacity(clamp_reserve(factor_count as i64));
            for _ in 0..factor_count.max(0) {
                let chr_index = c.i32()? as i64;
                let factor = if version > 8 {
                    c.f32()?
                } else {
                    c.f64()? as f32
                };
                chr_scale_factors.insert(chr_index, factor);
            }

            let key = vector_key(&normalization, bin_size, &unit, None);
            footer.expected_value_vectors.insert(
                key,
                ExpectedValueVector {
                    normalization,
                    unit,
                    bin_size,
                    values,
                    chr_scale_factors,
                },
            );
        }
    }

    let vector_count = c.i32()?;
    for _ in 0..vector_count.max(0) {
        let normalization = c.cstr()?.to_ascii_lowercase();
        note(&mut normalizations, &normalization);
        let chr_index = c.i32()? as i64;
        let unit = c.cstr()?.to_ascii_lowercase();
        note(&mut units, &unit);
        let bin_size = c.i32()? as i64;
        let position = c.i64()? as u64;
        let byte_count = if version > 8 {
            c.i64()?
        } else {
            c.i32()? as i64
        };
        let key = vector_key(&normalization, bin_size, &unit, Some(chr_index));
        footer.normalization_vectors.insert(
            key,
            NormalizationVector {
                normalization,
                chr_index,
                unit,
                bin_size,
                position,
                byte_count,
            },
        );
    }

    normalizations.sort();
    units.sort();
    footer.normalizations = normalizations;
    footer.units = units;
    Ok(footer)
}

/// Expected values for one chromosome at one resolution, scaled.
///
/// An expected value of zero cannot divide, and the contact it would have valued
/// comes back NaN — see `process_record` in `hic/block.rs`, which is private
/// and so not linked from here.
pub fn compute_expected_values(
    footer: &HiCFooter,
    chr_index: i64,
    unit: &str,
    bin_size: i64,
    normalization: &str,
) -> Result<Vec<f32>> {
    let key = vector_key(normalization, bin_size, unit, None);
    let vector = footer
        .expected_value_vectors
        .get(&key)
        .ok_or_else(|| Error::invalid(format!("expected value vector {key} not found")))?;
    let factor = vector.chr_scale_factors.get(&chr_index).ok_or_else(|| {
        Error::invalid(format!(
            "expected value vector {} not found",
            vector_key(normalization, bin_size, unit, Some(chr_index))
        ))
    })?;
    let scale = 1.0f32 / factor;
    Ok(vector.values.iter().map(|v| v * scale).collect())
}

pub fn read_normalization_vector(
    source: &dyn ByteSource,
    footer: &HiCFooter,
    version: i64,
    chr_index: i64,
    unit: &str,
    bin_size: i64,
    normalization: &str,
) -> Result<Vec<f32>> {
    let key = vector_key(normalization, bin_size, unit, Some(chr_index));
    let vector = footer
        .normalization_vectors
        .get(&key)
        .ok_or_else(|| Error::invalid(format!("normalization vector {key} not found")))?;
    let buffer = source.read_at(vector.position, vector.byte_count.max(0) as usize)?;

    // The count is read out of the file and then drives a read of that many
    // values, so it is checked against what the footer said the vector occupies.
    let (header_size, value_size) = if version > 8 {
        (8usize, 4usize)
    } else {
        (4, 8)
    };
    if buffer.len() < header_size {
        return Err(Error::corrupt(
            source.path(),
            vector.position,
            format!("normalization vector {key} is truncated"),
        ));
    }
    let count = if version > 8 {
        i64::from_le_bytes(buffer[..8].try_into().expect("checked length"))
    } else {
        i32::from_le_bytes(buffer[..4].try_into().expect("checked length")) as i64
    };
    // Widened and checked rather than multiplied in `usize`: `count` comes
    // straight out of the file, and `count * value_size` overflows for a value
    // a file can perfectly well name — which is a panic in a debug build and, in
    // a release one, a wrapped number small enough to pass the very check it is
    // part of.
    let declared = (count.max(0) as u64)
        .checked_mul(value_size as u64)
        .and_then(|bytes| bytes.checked_add(header_size as u64));
    if count < 0 || declared.is_none_or(|needed| needed > buffer.len() as u64) {
        return Err(Error::corrupt(
            source.path(),
            vector.position,
            format!(
                "normalization vector {key} declares {count} values, which do not fit its {} bytes",
                buffer.len()
            ),
        ));
    }
    let body = &buffer[header_size..];
    Ok(if version > 8 {
        body.chunks_exact(4)
            .take(count as usize)
            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
            .collect()
    } else {
        body.chunks_exact(8)
            .take(count as usize)
            .map(|c| f64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) as f32)
            .collect()
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::source::testing::MemorySource;

    /// A v8 header with two chromosomes and two resolutions.
    fn header_bytes(version: i32, frag: bool) -> Vec<u8> {
        let mut b = b"HIC\0".to_vec();
        b.extend_from_slice(&version.to_le_bytes());
        b.extend_from_slice(&4096i64.to_le_bytes()); // footer position
        b.extend_from_slice(b"mm10\0");
        if version > 8 {
            b.extend_from_slice(&0i64.to_le_bytes());
            b.extend_from_slice(&0i64.to_le_bytes());
        }
        b.extend_from_slice(&1i32.to_le_bytes()); // one attribute
        b.extend_from_slice(b"software\0made-up\0");
        b.extend_from_slice(&2i32.to_le_bytes()); // two chromosomes
        for (name, size) in [("chr1", 1000i64), ("chr2", 2000)] {
            b.extend_from_slice(name.as_bytes());
            b.push(0);
            if version > 8 {
                b.extend_from_slice(&size.to_le_bytes());
            } else {
                b.extend_from_slice(&(size as i32).to_le_bytes());
            }
        }
        b.extend_from_slice(&2i32.to_le_bytes()); // bp resolutions
        b.extend_from_slice(&5000i32.to_le_bytes());
        b.extend_from_slice(&10000i32.to_le_bytes());
        if frag {
            b.extend_from_slice(&1i32.to_le_bytes());
            b.extend_from_slice(&1i32.to_le_bytes());
            // And then the site lists the section only carries when it does:
            // two positions for chr1, none for chr2.
            b.extend_from_slice(&2i32.to_le_bytes());
            b.extend_from_slice(&100i32.to_le_bytes());
            b.extend_from_slice(&250i32.to_le_bytes());
            b.extend_from_slice(&0i32.to_le_bytes());
        } else {
            b.extend_from_slice(&0i32.to_le_bytes());
        }
        b
    }

    #[test]
    fn reads_a_v8_header() {
        let source = MemorySource::new(header_bytes(8, false));
        let h = read_header(&source).unwrap();
        assert_eq!(h.version, 8);
        assert_eq!(h.genome_id, "mm10");
        assert_eq!(h.footer_position, 4096);
        assert_eq!(
            h.attributes.get("software").map(String::as_str),
            Some("made-up")
        );
        assert_eq!(h.chr_map.names(), ["chr1", "chr2"]);
        assert_eq!(h.chr_map.resolve("chr2").unwrap().size, 2000);
        assert_eq!(h.bp_resolutions, [5000, 10000]);
        assert!(h.frag_resolutions.is_empty());
    }

    #[test]
    fn a_v9_header_widens_its_chromosome_sizes() {
        let source = MemorySource::new(header_bytes(9, false));
        let h = read_header(&source).unwrap();
        assert_eq!(h.version, 9);
        assert_eq!(h.chr_map.resolve("chr2").unwrap().size, 2000);
    }

    #[test]
    fn resolutions_are_reported_per_unit() {
        let source = MemorySource::new(header_bytes(8, true));
        let h = read_header(&source).unwrap();
        assert_eq!(h.resolutions(Unit::Bp), [5000, 10000]);
        assert_eq!(h.resolutions(Unit::Frag), [1]);
        // The site lists come with them, keyed in chromosome order.
        assert_eq!(h.sites.keys().collect::<Vec<_>>(), ["chr1", "chr2"]);
        assert_eq!(h.sites["chr1"], [100, 250]);
        assert!(h.sites["chr2"].is_empty());
    }

    #[test]
    fn a_file_without_fragment_resolutions_reads_no_site_section_at_all() {
        // Reading a count regardless would take the body's first four bytes for
        // one; here there is no body, so it would fail outright.
        let source = MemorySource::new(header_bytes(8, false));
        assert!(read_header(&source).unwrap().sites.is_empty());
    }

    #[test]
    fn only_a_v9_header_carries_a_normalization_vector_index() {
        assert_eq!(
            read_header(&MemorySource::new(header_bytes(8, false)))
                .unwrap()
                .nvi_position,
            -1
        );
        assert_eq!(
            read_header(&MemorySource::new(header_bytes(9, false)))
                .unwrap()
                .nvi_position,
            0
        );
    }

    #[test]
    fn a_bad_magic_and_an_unsupported_version_are_refused() {
        let source = MemorySource::new(b"NOPE\x08\0\0\0".to_vec());
        let err = read_header(&source).unwrap_err().to_string();
        assert!(err.contains("not a hic file"), "{err}");

        for version in [5i32, 10, 99] {
            let source = MemorySource::new(header_bytes(version, false));
            let err = read_header(&source).unwrap_err().to_string();
            assert!(
                err.contains("unsupported (6 to 9)"),
                "version {version}: {err}"
            );
        }
    }

    #[test]
    fn a_truncated_header_is_corrupt_not_a_panic() {
        let mut bytes = header_bytes(8, false);
        bytes.truncate(20);
        assert!(matches!(
            read_header(&MemorySource::new(bytes)),
            Err(Error::Corrupt { .. })
        ));
    }

    #[test]
    fn the_vector_key_is_built_the_same_way_both_ways() {
        assert_eq!(
            vector_key("kr", 5000, "bp", None),
            "normalization=kr|bin_size=5000|unit=bp"
        );
        assert_eq!(
            vector_key("kr", 5000, "bp", Some(3)),
            "chr_index=3|normalization=kr|bin_size=5000|unit=bp"
        );
    }

    #[test]
    fn expected_values_are_scaled_by_the_chromosomes_own_factor() {
        let mut footer = HiCFooter::default();
        let mut factors = IndexMap::new();
        factors.insert(0i64, 2.0f32);
        footer.expected_value_vectors.insert(
            vector_key("none", 5000, "bp", None),
            ExpectedValueVector {
                normalization: "none".into(),
                unit: "bp".into(),
                bin_size: 5000,
                values: vec![10.0, 20.0],
                chr_scale_factors: factors,
            },
        );
        // 1/2 is the scale, so the values halve.
        let got = compute_expected_values(&footer, 0, "bp", 5000, "none").unwrap();
        assert_eq!(got, [5.0, 10.0]);
        // A chromosome the vector has no factor for names the key it looked for.
        let err = compute_expected_values(&footer, 7, "bp", 5000, "none")
            .unwrap_err()
            .to_string();
        assert!(err.contains("chr_index=7"), "{err}");
    }
}