gwseq-io 0.2.1

//! The compression header block: what the container preserved, and how every
//! data series in it is encoded.
//!
//! Format reference: `docs/cram_format_v3.1.md` §8.4.
//!
//! One of these governs every slice in its container, and it is the first thing
//! any read must fetch — §12 is explicit that a slice cannot be decoded without
//! it, which is why random access in CRAM costs two reads where BAM costs one.
//! It is also not cheap to parse: two maps of encodings, one of them keyed by
//! auxiliary tag. The reader memoises it per container for exactly that reason.
//!
//! Three parts, in order: a preservation map saying what the encoder kept, a
//! data series encoding map, and a tag encoding map. Each is a byte length and
//! a count, then that many entries — and the byte length is honoured rather
//! than assumed, so a map with a key this reader does not know does not
//! desynchronise the two that follow it.

use std::collections::HashMap;

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

use super::container::read_itf8;
use super::encoding::Encoding;

/// §8.4's substitution matrix: five bytes, one per reference base.
///
/// Each byte packs the four two-bit codes for substituting that base, high bits
/// first, in `ACGTN` order with the reference base itself left out. Stored
/// here already inverted — code to base, which is the direction a decoder
/// reads — because inverting it once per file beats inverting it once per
/// substituted base.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SubstitutionMatrix {
    /// `table[ref base][code]`, both indexed by the `ACGTN` order.
    table: [[u8; 4]; 5],
}

/// The five reference bases the `BS` series is defined over. Anything else —
/// an ambiguity code — must be written verbatim through `BA` instead.
const BASES: [u8; 5] = *b"ACGTN";

impl Default for SubstitutionMatrix {
    /// The identity-ish fallback used when a file omits `SM`, which the format
    /// forbids but which costs nothing to survive: every code maps to `N`, so a
    /// substitution decodes to an unknown base rather than to a wrong one.
    fn default() -> Self {
        Self {
            table: [[b'N'; 4]; 5],
        }
    }
}

impl SubstitutionMatrix {
    pub fn parse(bytes: &[u8]) -> Self {
        let mut table = [[b'N'; 4]; 5];
        for (r, &packed) in bytes.iter().enumerate().take(5) {
            // The substitutes for reference base r, in ACGTN order with r
            // itself skipped; the first gets the high two bits.
            let mut shift = 6;
            for &base in BASES.iter() {
                if base == BASES[r] {
                    continue;
                }
                let code = (packed >> shift) & 3;
                table[r][code as usize] = base;
                shift -= 2;
            }
        }
        Self { table }
    }

    /// The read base a substitution code names, given the reference base.
    ///
    /// Case-insensitive on the reference, as §10.6 requires ("all base
    /// comparisons should be done in a case-insensitive manner"). A reference
    /// base outside `ACGTN` is treated as `N`, which is what an out-of-range
    /// or ambiguous reference position is.
    pub fn substitute(&self, reference: u8, code: u8) -> u8 {
        let row = match reference.to_ascii_uppercase() {
            b'A' => 0,
            b'C' => 1,
            b'G' => 2,
            b'T' => 3,
            _ => 4,
        };
        self.table[row][(code & 3) as usize]
    }
}

/// §8.4's preservation map.
#[derive(Debug, Clone)]
pub struct PreservationMap {
    /// `RN`: read names are stored for every record.
    pub read_names_included: bool,
    /// `AP`: the alignment-start series is a delta rather than an absolute.
    pub ap_delta: bool,
    /// `RR`: the reference is needed to rebuild the bases.
    pub reference_required: bool,
    pub substitution_matrix: SubstitutionMatrix,
    /// `TD`: every distinct combination of auxiliary tag ids and types that
    /// occurs, each entry a run of three-byte `(id, id, type)` triples. The
    /// `TL` series indexes into this.
    pub tag_dictionary: Vec<Vec<[u8; 3]>>,
}

impl Default for PreservationMap {
    /// §8.4: the three booleans default to true when absent. `SM` and `TD` are
    /// mandatory, and their defaults here are only what a malformed file falls
    /// back to rather than something to rely on.
    fn default() -> Self {
        Self {
            read_names_included: true,
            ap_delta: true,
            reference_required: true,
            substitution_matrix: SubstitutionMatrix::default(),
            tag_dictionary: Vec::new(),
        }
    }
}

/// Every data series' encoding, by its two-letter name.
///
/// A struct rather than a map: the record decoder reaches most of these once
/// per record, and a hash of a two-byte key per series per record is a cost
/// with nothing to show for it. The `Null` default is what an absent series
/// means, and reading one is an error rather than a zero — see
/// [`Encoding::Null`].
#[derive(Debug, Clone, Default)]
pub struct DataSeries {
    pub bf: Encoding,
    pub cf: Encoding,
    pub ri: Encoding,
    pub rl: Encoding,
    pub ap: Encoding,
    pub rg: Encoding,
    pub rn: Encoding,
    pub mf: Encoding,
    pub ns: Encoding,
    pub np: Encoding,
    pub ts: Encoding,
    pub nf: Encoding,
    pub tl: Encoding,
    pub fn_: Encoding,
    pub fc: Encoding,
    pub fp: Encoding,
    pub dl: Encoding,
    pub bb: Encoding,
    pub qq: Encoding,
    pub bs: Encoding,
    pub in_: Encoding,
    pub rs: Encoding,
    pub pd: Encoding,
    pub hc: Encoding,
    pub sc: Encoding,
    pub mq: Encoding,
    pub ba: Encoding,
    pub qs: Encoding,
}

impl DataSeries {
    fn set(&mut self, key: [u8; 2], encoding: Encoding) {
        let slot = match &key {
            b"BF" => &mut self.bf,
            b"CF" => &mut self.cf,
            b"RI" => &mut self.ri,
            b"RL" => &mut self.rl,
            b"AP" => &mut self.ap,
            b"RG" => &mut self.rg,
            b"RN" => &mut self.rn,
            b"MF" => &mut self.mf,
            b"NS" => &mut self.ns,
            b"NP" => &mut self.np,
            b"TS" => &mut self.ts,
            b"NF" => &mut self.nf,
            b"TL" => &mut self.tl,
            b"FN" => &mut self.fn_,
            b"FC" => &mut self.fc,
            b"FP" => &mut self.fp,
            b"DL" => &mut self.dl,
            b"BB" => &mut self.bb,
            b"QQ" => &mut self.qq,
            b"BS" => &mut self.bs,
            b"IN" => &mut self.in_,
            b"RS" => &mut self.rs,
            b"PD" => &mut self.pd,
            b"HC" => &mut self.hc,
            b"SC" => &mut self.sc,
            b"MQ" => &mut self.mq,
            b"BA" => &mut self.ba,
            b"QS" => &mut self.qs,
            // `TC` and `TN` are CRAM 1.0 leftovers that §8.4 says decoders
            // must silently skip. Some writers still emit them, empty.
            _ => return,
        };
        *slot = encoding;
    }

    fn each(&self) -> impl Iterator<Item = &Encoding> {
        [
            &self.bf, &self.cf, &self.ri, &self.rl, &self.ap, &self.rg, &self.rn, &self.mf,
            &self.ns, &self.np, &self.ts, &self.nf, &self.tl, &self.fn_, &self.fc, &self.fp,
            &self.dl, &self.bb, &self.qq, &self.bs, &self.in_, &self.rs, &self.pd, &self.hc,
            &self.sc, &self.mq, &self.ba, &self.qs,
        ]
        .into_iter()
    }
}

/// The parsed compression header block.
#[derive(Debug, Clone)]
pub struct CompressionHeader {
    pub preservation: PreservationMap,
    pub series: DataSeries,
    /// Auxiliary tag encodings, keyed as §8.4 keys them: `(id1 << 16) | (id2 <<
    /// 8) | type`, which is how the tag dictionary's three bytes read as an
    /// integer.
    pub tags: HashMap<i32, Encoding>,
}

impl CompressionHeader {
    pub fn parse(data: &[u8], path: &str) -> Result<Self> {
        let mut cursor = LeCursor::new(data, 0, path);
        let preservation = Self::parse_preservation(&mut cursor)?;
        let series = Self::parse_series(&mut cursor)?;
        let tags = Self::parse_tags(&mut cursor)?;
        Ok(Self {
            preservation,
            series,
            tags,
        })
    }

    /// The byte span of one of the three maps, so an unknown entry can be
    /// stepped over rather than guessed at.
    fn map_extent(cursor: &mut LeCursor<'_>) -> Result<(usize, usize)> {
        let n_bytes = read_itf8(cursor)?;
        if n_bytes < 0 {
            return Err(Error::corrupt(
                cursor.path(),
                cursor.file_offset(),
                format!("a compression header map of {n_bytes} bytes"),
            ));
        }
        // The declared size covers the count and the entries, from here.
        let end = cursor.position() + n_bytes as usize;
        let count = read_itf8(cursor)?;
        if count < 0 {
            return Err(Error::corrupt(
                cursor.path(),
                cursor.file_offset(),
                format!("a compression header map of {count} entries"),
            ));
        }
        Ok((end, count as usize))
    }

    fn parse_preservation(cursor: &mut LeCursor<'_>) -> Result<PreservationMap> {
        let (end, count) = Self::map_extent(cursor)?;
        let mut map = PreservationMap::default();
        for _ in 0..count {
            let key: [u8; 2] = cursor.take(2)?.try_into().expect("two bytes");
            match &key {
                b"RN" => map.read_names_included = cursor.take(1)?[0] != 0,
                b"AP" => map.ap_delta = cursor.take(1)?[0] != 0,
                b"RR" => map.reference_required = cursor.take(1)?[0] != 0,
                b"SM" => map.substitution_matrix = SubstitutionMatrix::parse(cursor.take(5)?),
                b"TD" => {
                    let n = read_itf8(cursor)?;
                    if n < 0 {
                        return Err(Error::corrupt(
                            cursor.path(),
                            cursor.file_offset(),
                            format!("a tag dictionary of {n} bytes"),
                        ));
                    }
                    map.tag_dictionary = parse_tag_dictionary(cursor.take(n as usize)?);
                }
                // An unknown key has no length this reader can know, so the
                // only safe thing is to stop reading entries and let the
                // declared extent carry the cursor past them.
                _ => break,
            }
        }
        cursor.seek(end)?;
        Ok(map)
    }

    fn parse_series(cursor: &mut LeCursor<'_>) -> Result<DataSeries> {
        let (end, count) = Self::map_extent(cursor)?;
        let mut series = DataSeries::default();
        for _ in 0..count {
            let key: [u8; 2] = cursor.take(2)?.try_into().expect("two bytes");
            series.set(key, Encoding::read(cursor)?);
        }
        cursor.seek(end)?;
        Ok(series)
    }

    fn parse_tags(cursor: &mut LeCursor<'_>) -> Result<HashMap<i32, Encoding>> {
        let (end, count) = Self::map_extent(cursor)?;
        let mut tags = HashMap::with_capacity(count.min(1024));
        for _ in 0..count {
            let key = read_itf8(cursor)?;
            tags.insert(key, Encoding::read(cursor)?);
        }
        cursor.seek(end)?;
        Ok(tags)
    }

    /// Every external block id any series or tag reads.
    pub fn block_ids(&self) -> Vec<i32> {
        let mut out = Vec::new();
        for encoding in self.series.each() {
            encoding.block_ids(&mut out);
        }
        for encoding in self.tags.values() {
            encoding.block_ids(&mut out);
        }
        out.sort_unstable();
        out.dedup();
        out
    }

    /// The tag list a `TL` value names.
    pub fn tag_list(&self, index: i32, path: &str) -> Result<&[[u8; 3]]> {
        self.preservation
            .tag_dictionary
            .get(usize::try_from(index).unwrap_or(usize::MAX))
            .map(Vec::as_slice)
            .ok_or_else(|| {
                Error::corrupt(
                    path,
                    0,
                    format!(
                        "a record names tag list {index}, and the dictionary holds {}",
                        self.preservation.tag_dictionary.len()
                    ),
                )
            })
    }
}

/// `TD`: NUL-separated lists, each a run of three-byte tag descriptors.
///
/// A trailing NUL leaves an empty final field, which is not an entry — but an
/// empty list in the middle *is* one, and means a record with no tags at all.
fn parse_tag_dictionary(data: &[u8]) -> Vec<Vec<[u8; 3]>> {
    let mut out = Vec::new();
    let mut rest = data;
    while !rest.is_empty() {
        let end = memchr::memchr(0, rest).unwrap_or(rest.len());
        let (entry, tail) = rest.split_at(end);
        out.push(entry.chunks_exact(3).map(|c| [c[0], c[1], c[2]]).collect());
        rest = tail.get(1..).unwrap_or(&[]);
    }
    out
}

/// The integer key §8.4 gives a tag: its two id bytes and its type byte, read
/// as a big-endian three-byte number.
pub fn tag_key(tag: [u8; 3]) -> i32 {
    (i32::from(tag[0]) << 16) | (i32::from(tag[1]) << 8) | i32::from(tag[2])
}

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

    /// §10.6's worked matrix, both directions. The specification prints the
    /// packed bytes and the decode table they invert to, so this is its example
    /// checked rather than a re-statement of the code.
    #[test]
    fn the_spec_substitution_matrix_inverts_to_the_table_it_prints() {
        let matrix = SubstitutionMatrix::parse(&[0x63, 0x4b, 0x87, 0x27, 0x1b]);
        let expected: [(u8, [u8; 4]); 5] = [
            (b'A', *b"TCGN"),
            (b'C', *b"GATN"),
            (b'G', *b"CTAN"),
            (b'T', *b"AGCN"),
            (b'N', *b"ACGT"),
        ];
        for (reference, row) in expected {
            for (code, base) in row.iter().enumerate() {
                assert_eq!(
                    matrix.substitute(reference, code as u8),
                    *base,
                    "ref {} code {code}",
                    reference as char
                );
            }
        }
    }

    /// Every row of a well-formed matrix is a permutation of the four bases the
    /// reference is not, which is what makes a substitution reversible at all.
    ///
    /// "Well-formed" is the load-bearing word: a packed byte is free to give two
    /// substitutes the same code, and then one base is simply unreachable. That
    /// is legal input and decodes to `N`, so the property is asserted over
    /// matrices built from real permutations rather than from arbitrary bytes.
    #[test]
    fn every_row_of_a_well_formed_substitution_matrix_is_a_permutation() {
        let permutations: Vec<[u8; 4]> = {
            let mut out = Vec::new();
            for a in 0..4u8 {
                for b in 0..4u8 {
                    for c in 0..4u8 {
                        for d in 0..4u8 {
                            let code = [a, b, c, d];
                            let mut sorted = code;
                            sorted.sort_unstable();
                            if sorted == [0, 1, 2, 3] {
                                out.push(code);
                            }
                        }
                    }
                }
            }
            out
        };
        assert_eq!(permutations.len(), 24);

        for (i, _) in permutations.iter().enumerate() {
            // A different permutation for each of the five reference bases.
            let packed: Vec<u8> = (0..5)
                .map(|r| {
                    let p = permutations[(i + r) % permutations.len()];
                    (p[0] << 6) | (p[1] << 4) | (p[2] << 2) | p[3]
                })
                .collect();
            let matrix = SubstitutionMatrix::parse(&packed);
            for &reference in BASES.iter() {
                let mut seen: Vec<u8> = (0..4).map(|c| matrix.substitute(reference, c)).collect();
                seen.sort_unstable();
                let mut want: Vec<u8> = BASES.iter().copied().filter(|b| *b != reference).collect();
                want.sort_unstable();
                assert_eq!(seen, want, "ref {}", reference as char);
            }
        }
    }

    /// A code two substitutes share leaves one unreachable, and that code reads
    /// as `N` rather than as some base the file never named.
    #[test]
    fn a_code_no_substitute_claims_reads_as_an_unknown_base() {
        // Reference A, every substitute given code 0.
        let matrix = SubstitutionMatrix::parse(&[0x00, 0, 0, 0, 0]);
        assert_eq!(matrix.substitute(b'A', 1), b'N');
        assert_eq!(matrix.substitute(b'A', 2), b'N');
    }

    #[test]
    fn the_reference_base_is_matched_case_insensitively() {
        let matrix = SubstitutionMatrix::parse(&[0x63, 0x4b, 0x87, 0x27, 0x1b]);
        assert_eq!(matrix.substitute(b'a', 0), matrix.substitute(b'A', 0));
        assert_eq!(matrix.substitute(b'g', 2), matrix.substitute(b'G', 2));
        // Anything outside ACGTN reads as N.
        assert_eq!(matrix.substitute(b'M', 1), matrix.substitute(b'N', 1));
    }

    #[test]
    fn the_tag_dictionary_splits_into_lists_of_three_byte_descriptors() {
        let entries = parse_tag_dictionary(b"X1CBCZSAZ\0X1CBCZ\0");
        assert_eq!(entries.len(), 2);
        assert_eq!(entries[0], vec![*b"X1C", *b"BCZ", *b"SAZ"]);
        assert_eq!(entries[1], vec![*b"X1C", *b"BCZ"]);
    }

    /// §8.4's own example of a tag key: `OQ:Z` is `{0x4F, 0x51, 0x5A}`, which
    /// reads as the integer `0x004F515A`.
    #[test]
    fn a_tag_key_is_its_three_bytes_read_as_an_integer() {
        assert_eq!(tag_key(*b"OQZ"), 0x004F_515A);
    }

    /// The empty compression header the EOF container carries: three maps, each
    /// one byte long and holding nothing.
    #[test]
    fn the_eof_containers_empty_compression_header_parses() {
        let header = CompressionHeader::parse(&[0x01, 0x00, 0x01, 0x00, 0x01, 0x00], "test")
            .expect("an empty header");
        assert!(header.tags.is_empty());
        assert!(header.series.bf.is_null());
        // The booleans default to true when the map does not mention them.
        assert!(header.preservation.read_names_included);
        assert!(header.preservation.ap_delta);
        assert!(header.preservation.reference_required);
    }

    #[test]
    fn a_preservation_map_reads_its_flags_and_matrix() {
        let mut body = vec![0u8];
        body.extend_from_slice(b"RN");
        body.push(0);
        body.extend_from_slice(b"AP");
        body.push(0);
        body.extend_from_slice(b"SM");
        body.extend_from_slice(&[0x63, 0x4b, 0x87, 0x27, 0x1b]);
        body.extend_from_slice(b"TD");
        body.push(4);
        body.extend_from_slice(b"X1C\0");
        body[0] = 4; // entry count

        let mut data = vec![body.len() as u8];
        data.extend_from_slice(&body);
        // Two empty maps behind it.
        data.extend_from_slice(&[0x01, 0x00, 0x01, 0x00]);

        let header = CompressionHeader::parse(&data, "test").expect("a header");
        assert!(!header.preservation.read_names_included);
        assert!(!header.preservation.ap_delta);
        assert!(header.preservation.reference_required); // absent, so true
        assert_eq!(header.preservation.tag_dictionary, vec![vec![*b"X1C"]]);
        assert_eq!(
            header.preservation.substitution_matrix.substitute(b'A', 0),
            b'T'
        );
    }
}