gwseq-io 0.2.1

//! Block compression — the *external* methods of §14, one per `method` byte.
//!
//! These are a layer below the encodings in [`super::encoding`], and the two
//! are easy to confuse. An **encoding** says how one data series is laid out
//! within a block (a Huffman code, a fixed-width integer, a byte array with a
//! terminator); a **compression method** says how the block's bytes are
//! squeezed once that layout is decided. A slice's `QS` series is
//! `EXTERNAL(id=12)` — an encoding — and block 12 is `rans4x16` — a method.
//!
//! Five of the nine are written here from the *CRAMcodecs* document rather than
//! taken from a crate, because no crate offers them: they exist for this format
//! and nothing else. That is the risk in this module and it is worth naming —
//! an entropy decoder that is subtly wrong does not fail, it returns plausible
//! bytes. Every one of them is tested against an encoder written alongside it,
//! so the round trip is real rather than a re-run of the same misreading, and
//! against the worked examples the specification prints.
//!
//! The other four are not exercises. `raw` is a copy, gzip and the `.crai` are
//! `flate2`, and [`bzip2`] and [`lzma`] are dependencies — a general-purpose
//! compressor is somebody else's job. All nine of §14's methods are here.
//!
//! # What each one costs to leave out
//!
//! Measured by re-encoding one file every way `samtools` offers and counting
//! the block methods that came back:
//!
//! | Method | Where it turns up |
//! |---|---|
//! | `rans4x8` | 72–74% of the bytes of *every* CRAM 3.0 file, at every profile |
//! | `rans4x16` | 74–80% of a 3.1 file |
//! | `nametok` | 12–14% of a 3.1 file, always the read names |
//! | `fqzcomp` | 71% of a 3.1 `small` or `archive` file, always the qualities |
//! | `bzip2` | 8% of a `small` or `archive` file, either version — but on `BF`, `AP`, `MQ` and `NF` |
//! | `arith` | 6% of a 3.1 `archive` file — but on eleven of its eighteen series |
//! | `lzma` | only with an explicit `use_lzma=1`; no profile turns it on — but then it is 70% of the file |
//!
//! The last two rows are why byte share is the wrong measure. A codec holding
//! alignment positions is not eight percent of a file, it is all of whether
//! the file opens.
//!
//! Format reference: `docs/cram_codecs_v3.1.md`.

use std::io::Read as _;

use bytes::Bytes;

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

use super::container::{CompressionMethod, MAX_BLOCK_RAW_SIZE};

pub(crate) mod arith;
pub(crate) mod bzip2;
pub(crate) mod fqzcomp;
pub(crate) mod lzma;
pub(crate) mod rans4x16;
pub(crate) mod rans4x8;
pub(crate) mod tokenise;

/// Undo one block's compression.
///
/// `raw_size` is what the block header declared, and every method is handed it
/// rather than discovering the length itself: the formats here mostly carry
/// their own output length too, and a disagreement between the two is a
/// corrupt file worth catching at the boundary instead of halfway through a
/// decode.
pub fn decode(
    method: CompressionMethod,
    data: &[u8],
    raw_size: usize,
    path: &str,
    offset: u64,
) -> Result<Bytes> {
    let out = match method {
        CompressionMethod::Raw => {
            if data.len() != raw_size {
                return Err(Error::corrupt(
                    path,
                    offset,
                    format!(
                        "a raw block declares {raw_size} bytes and carries {}",
                        data.len()
                    ),
                ));
            }
            return Ok(Bytes::copy_from_slice(data));
        }
        CompressionMethod::Gzip => gzip(data, raw_size, path, offset)?,
        CompressionMethod::Rans4x8 => rans4x8::decode(data, path, offset)?,
        CompressionMethod::Rans4x16 => rans4x16::decode(data, path, offset)?,
        CompressionMethod::Arith => arith::decode(data, path, offset)?,
        CompressionMethod::NameTok => tokenise::decode(data, path, offset)?,
        CompressionMethod::Bzip2 => bzip2::decode(data, raw_size, path, offset)?,
        CompressionMethod::Fqzcomp => fqzcomp::decode(data, path, offset)?,
        CompressionMethod::Lzma => lzma::decode(data, raw_size, path, offset)?,
    };
    if out.len() != raw_size {
        return Err(Error::corrupt(
            path,
            offset,
            format!(
                "a {} block declares {raw_size} bytes and decoded to {}",
                method.name(),
                out.len()
            ),
        ));
    }
    Ok(Bytes::from(out))
}

/// Inflate a gzip member.
///
/// `raw_size` is a ceiling as well as an expectation — see the crate's rule
/// about never reserving what a file names. The header's number is trusted for
/// the reserve only after `Block::parse` has checked it against
/// `MAX_BLOCK_RAW_SIZE`.
fn gzip(data: &[u8], raw_size: usize, path: &str, offset: u64) -> Result<Vec<u8>> {
    let mut out = Vec::with_capacity(raw_size.min(1 << 20));
    // Bounded, not merely checked afterwards. `raw_size` is capped at
    // `MAX_BLOCK_RAW_SIZE` by `Block::parse`, but that caps the *reserve*: a
    // block declaring a thousand raw bytes whose stream inflates to four
    // gigabytes would allocate all four before the mismatch below noticed. The
    // extra byte is what turns "filled the bound exactly" into "had more to
    // give", which the caller's length check then reports.
    inflate_bounded(
        flate2::read::MultiGzDecoder::new(data),
        &mut out,
        raw_size,
        "gzip",
        path,
        offset,
    )?;
    Ok(out)
}

/// Read a decoder out to at most `limit` bytes.
///
/// The same shape the bbi and HiC inflates use. A decoder that still has bytes
/// at the limit is a decompression bomb and is refused by name rather than by
/// the length mismatch it would eventually cause.
pub(crate) fn inflate_bounded<R: std::io::Read>(
    decoder: R,
    out: &mut Vec<u8>,
    limit: usize,
    what: &str,
    path: &str,
    offset: u64,
) -> Result<()> {
    let capped = limit.min(MAX_BLOCK_RAW_SIZE) as u64;
    let read = decoder.take(capped + 1).read_to_end(out).map_err(|e| {
        Error::corrupt(
            path,
            offset,
            format!("could not inflate a {what} block: {e}"),
        )
    })?;
    if read as u64 > capped {
        return Err(Error::corrupt(
            path,
            offset,
            format!("a {what} block inflating past the {limit} bytes it declared"),
        ));
    }
    Ok(())
}

/// A short read inside a codec, which is always a corrupt file rather than a
/// caller's mistake.
pub(crate) fn short(path: &str, offset: u64, what: &str) -> Error {
    Error::corrupt(
        path,
        offset,
        format!("{what} ran past the end of the block"),
    )
}

/// The ceiling on any length a codec's own stream names.
///
/// These formats nest — a stripe holds four sub-streams, each with its own
/// declared length — so the block header's `raw_size` does not bound what the
/// bytes inside can ask for. Same rule as everywhere else in the crate: clamp
/// where the number is read, because a failed allocation aborts the process
/// and no `Result` can carry that.
pub(crate) const MAX_CODEC_LEN: usize = 1 << 30;

/// §3.5's bit packing: an alphabet of at most sixteen, several to a byte.
///
/// Shared with the arithmetic coder, which §4 defines by pointing back here
/// — "the same `DecodePackMeta` and `DecodePack` functions are used" — so the
/// messages below name the transform rather than either codec.
pub(crate) struct PackMeta {
    /// Padded to the full width of the code, so an unused code has an entry.
    map: Vec<u8>,
    pub(crate) packed_len: usize,
    n_symbols: usize,
}

impl PackMeta {
    pub(crate) fn read(reader: &mut ByteReader<'_>) -> Result<Self> {
        let n_symbols = reader.u8()? as usize;
        if n_symbols == 0 || n_symbols > 16 {
            return Err(Error::corrupt(
                reader.path(),
                reader.offset(),
                format!("a pack map of {n_symbols} symbols, which cannot be packed"),
            ));
        }
        let mut map = reader.take(n_symbols)?.to_vec();
        // A code wider than the alphabet is legal input and means nothing —
        // three symbols still travel two bits at a time, so the fourth code can
        // occur. Padded rather than checked so that an unused code decodes to a
        // nul instead of indexing past the map.
        let width = match n_symbols {
            1 => 1,
            2 => 2,
            3..=4 => 4,
            _ => 16,
        };
        map.resize(width, 0);
        Ok(Self {
            map,
            packed_len: reader.length()?,
            n_symbols,
        })
    }

    pub(crate) fn unpack(
        &self,
        data: &[u8],
        len: usize,
        path: &str,
        offset: u64,
    ) -> Result<Vec<u8>> {
        // One symbol: nothing was packed, because nothing needed to be.
        let per_byte = match self.n_symbols {
            1 => return Ok(vec![self.map[0]; len]),
            2 => 8,
            3..=4 => 4,
            _ => 2,
        };
        let bits = 8 / per_byte;
        let mask = (1u8 << bits) - 1;
        let wanted = len.div_ceil(per_byte);
        if data.len() < wanted {
            return Err(Error::corrupt(
                path,
                offset,
                format!(
                    "a packed stream of {} bytes where {wanted} were needed for {len} values",
                    data.len()
                ),
            ));
        }
        let mut out = Vec::with_capacity(len.min(1 << 20));
        'outer: for &byte in data {
            for k in 0..per_byte {
                if out.len() >= len {
                    break 'outer;
                }
                out.push(self.map[((byte >> (k * bits)) & mask) as usize]);
            }
        }
        Ok(out)
    }
}

/// A forward byte reader with the widths these codecs are written in.
///
/// Separate from [`crate::bytes::LeCursor`] because the two disagree on what a
/// variable-width integer is: this one speaks `uint7`, seven bits per byte with
/// the top bit as the continuation flag, big-endian — which is neither ITF8 nor
/// anything `LeCursor` offers.
pub(crate) struct ByteReader<'a> {
    data: &'a [u8],
    pos: usize,
    path: &'a str,
    /// Where `data[0]` sits in the file, so an error names a file offset.
    base: u64,
}

impl<'a> ByteReader<'a> {
    pub fn new(data: &'a [u8], path: &'a str, base: u64) -> Self {
        Self {
            data,
            pos: 0,
            path,
            base,
        }
    }

    pub fn remaining(&self) -> usize {
        self.data.len() - self.pos
    }

    pub fn is_empty(&self) -> bool {
        self.remaining() == 0
    }

    fn fail(&self, what: &str) -> Error {
        short(self.path, self.base + self.pos as u64, what)
    }

    pub fn u8(&mut self) -> Result<u8> {
        let byte = *self.data.get(self.pos).ok_or_else(|| self.fail("a byte"))?;
        self.pos += 1;
        Ok(byte)
    }

    pub fn u16(&mut self) -> Result<u16> {
        Ok(u16::from_le_bytes(
            self.take(2)?.try_into().expect("two bytes"),
        ))
    }

    pub fn u32(&mut self) -> Result<u32> {
        Ok(u32::from_le_bytes(
            self.take(4)?.try_into().expect("four bytes"),
        ))
    }

    /// A `uint7`: seven bits a byte, most significant group first, the top bit
    /// set on every byte but the last.
    ///
    /// Capped at five bytes. Without that a run of `0x80`s is an infinite
    /// shift that silently keeps the low bits — and the fuzzer finds it.
    pub fn uint7(&mut self) -> Result<u32> {
        let mut value: u32 = 0;
        for _ in 0..5 {
            let byte = self.u8()?;
            value = (value << 7) | u32::from(byte & 0x7f);
            if byte < 128 {
                return Ok(value);
            }
        }
        Err(Error::corrupt(
            self.path,
            self.base + self.pos as u64,
            "a uint7 longer than the 32 bits it can hold",
        ))
    }

    /// An ITF8, §1.4 — the *container* format's variable-width integer, which
    /// is big-endian with the byte count in the leading one-bits.
    ///
    /// Only rANS 4x8 wants it: its frequency tables predate the `uint7` the
    /// 3.1 codecs read. The logic is
    /// [`super::container::read_itf8`]'s, over this reader instead of a
    /// [`crate::bytes::LeCursor`] — including the
    /// five-byte form, where the last byte gives up its *high* nibble. If one
    /// of the two is ever corrected, the other is wrong.
    pub fn itf8(&mut self) -> Result<u32> {
        let first = self.u8()?;
        Ok(if first & 0x80 == 0 {
            u32::from(first)
        } else if first & 0x40 == 0 {
            (u32::from(first & 0x7f) << 8) | u32::from(self.u8()?)
        } else if first & 0x20 == 0 {
            let rest = self.take(2)?;
            (u32::from(first & 0x3f) << 16) | (u32::from(rest[0]) << 8) | u32::from(rest[1])
        } else if first & 0x10 == 0 {
            let rest = self.take(3)?;
            (u32::from(first & 0x1f) << 24)
                | (u32::from(rest[0]) << 16)
                | (u32::from(rest[1]) << 8)
                | u32::from(rest[2])
        } else {
            let rest = self.take(4)?;
            (u32::from(first & 0x0f) << 28)
                | (u32::from(rest[0]) << 20)
                | (u32::from(rest[1]) << 12)
                | (u32::from(rest[2]) << 4)
                | u32::from(rest[3] & 0x0f)
        })
    }

    /// A `uint7` used as a length, checked against this reader's ceiling.
    pub fn length(&mut self) -> Result<usize> {
        let value = self.uint7()? as usize;
        if value > MAX_CODEC_LEN {
            return Err(Error::corrupt(
                self.path,
                self.base + self.pos as u64,
                format!("a declared length of {value} bytes, past this reader's ceiling"),
            ));
        }
        Ok(value)
    }

    pub fn take(&mut self, n: usize) -> Result<&'a [u8]> {
        let end = self
            .pos
            .checked_add(n)
            .filter(|end| *end <= self.data.len())
            .ok_or_else(|| self.fail(&format!("{n} bytes")))?;
        let out = &self.data[self.pos..end];
        self.pos = end;
        Ok(out)
    }

    pub fn path(&self) -> &'a str {
        self.path
    }

    pub fn offset(&self) -> u64 {
        self.base + self.pos as u64
    }
}