gwseq-io 0.2.1

//! Slices — the unit of random access, and the unit of decoding.
//!
//! Format reference: `docs/cram_format_v3.1.md` §8.5.
//!
//! A slice is a header block plus a core block plus one external block per data
//! series it uses. Nothing smaller can be read: the entropy coders run the
//! length of a block, and `NF` reaches from one record to another within the
//! slice, so ten thousand records come back or none do. That is the cost model
//! for every query — a two-hundred-base locus decodes whichever slices overlap
//! it, whole.

use bytes::Bytes;

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

use super::container::{
    read_itf8, read_itf8_array, read_ltf8, Block, BlockContentType, ContainerHeader,
};

/// Reference id meaning "this slice holds unmapped, unplaced reads".
pub const UNMAPPED: i32 = -1;
/// Reference id meaning "the `RI` series says, per record".
pub const MULTI_REF: i32 = -2;

/// §8.5's slice header block.
#[derive(Debug, Clone)]
pub struct SliceHeader {
    pub ref_id: i32,
    pub start: i32,
    pub span: i32,
    pub n_records: i32,
    pub record_counter: i64,
    pub n_blocks: i32,
    pub block_content_ids: Vec<i32>,
    /// The block holding the reference bases for this slice, or -1 when the
    /// reference is external.
    pub embedded_reference_id: i32,
    /// All zeros when the slice does not use an external reference, and then
    /// not to be checked — §11's fourth rule.
    pub reference_md5: [u8; 16],
}

impl SliceHeader {
    pub fn parse(data: &[u8], path: &str) -> Result<Self> {
        let mut cursor = LeCursor::new(data, 0, path);
        let ref_id = read_itf8(&mut cursor)?;
        let start = read_itf8(&mut cursor)?;
        let span = read_itf8(&mut cursor)?;
        let n_records = read_itf8(&mut cursor)?;
        let record_counter = read_ltf8(&mut cursor)?;
        let n_blocks = read_itf8(&mut cursor)?;
        let block_content_ids = read_itf8_array(&mut cursor)?;
        let embedded_reference_id = read_itf8(&mut cursor)?;
        let mut reference_md5 = [0u8; 16];
        // §8.5 puts the MD5 before an optional-tag field that no tag is yet
        // defined for. A slice header cut short of it is not fatal — nothing
        // here depends on the MD5 — so a short read leaves it zero, which
        // already means "do not check me".
        if cursor.remaining() >= 16 {
            reference_md5.copy_from_slice(cursor.take(16)?);
        }
        if n_records < 0 {
            return Err(Error::corrupt(
                path,
                0,
                format!("a slice of {n_records} records"),
            ));
        }
        Ok(Self {
            ref_id,
            start,
            span,
            n_records,
            record_counter,
            n_blocks,
            block_content_ids,
            embedded_reference_id,
            reference_md5,
        })
    }

    /// Whether the `RI` series decides each record's reference.
    pub fn is_multi_ref(&self) -> bool {
        self.ref_id == MULTI_REF
    }

    /// The half-open reference span this slice covers, 0-based.
    ///
    /// `None` when the slice is unmapped or multi-reference, where §8.5 says
    /// the start and span fields are to be ignored.
    pub fn range(&self) -> Option<(i64, i64)> {
        if self.ref_id == UNMAPPED || self.is_multi_ref() || self.ref_id < 0 {
            return None;
        }
        let start = i64::from(self.start) - 1;
        Some((start.max(0), start.max(0) + i64::from(self.span).max(0)))
    }

    /// Whether the MD5 is one to check, or the sixteen zeros that mean it is
    /// not.
    pub fn has_reference_md5(&self) -> bool {
        self.reference_md5 != [0u8; 16]
    }
}

/// A slice with its blocks read and decompressed.
#[derive(Debug)]
pub struct Slice {
    pub header: SliceHeader,
    /// Where this slice starts in the file.
    ///
    /// Carried so that an error from the record decoder can name a place. The
    /// decoder works in a bit stream and several external blocks at once, so
    /// there is no one cursor to report — but "this slice, this record" is
    /// enough to find the problem, and it is what the BAM decoder gives.
    pub offset: u64,
    /// The bit stream. Empty when every series is external, which is the common
    /// shape — a real 3.1 file puts nothing here at all.
    pub core: Bytes,
    /// External blocks by content id.
    pub external: Vec<(i32, Bytes)>,
    /// The embedded reference bases, when the slice carries them.
    pub embedded_reference: Option<Bytes>,
}

impl Slice {
    /// Read and decompress the slice at `offset`.
    ///
    /// The whole slice is fetched in one read rather than block by block: its
    /// size is known from the index or from the container's next landmark, and
    /// over HTTP the difference between one range request and thirty is the
    /// difference between a query and a wait.
    pub fn read(source: &dyn ByteSource, offset: u64, size: usize) -> Result<Self> {
        let path = source.path();
        let data = source.read_exact_at(offset, size)?;
        Self::parse(&data, offset, path)
    }

    /// Just the header block of the slice at `offset`.
    ///
    /// `read` decompresses *every* block to reach the header, which is the
    /// right trade when the records are wanted and badly wrong when they are
    /// not: building an index from a file with several slices per container
    /// turns into entropy-decoding the whole file. Only the first block is
    /// touched here.
    pub fn read_header(source: &dyn ByteSource, offset: u64, size: usize) -> Result<SliceHeader> {
        let path = source.path();
        let data = source.read_exact_at(offset, size)?;
        let head = Block::parse(&data, offset, path)?;
        if head.content_type != BlockContentType::SliceHeader {
            return Err(Error::corrupt(
                path,
                offset,
                format!(
                    "expected a slice header block and found {:?}",
                    head.content_type
                ),
            ));
        }
        SliceHeader::parse(&head.data, path)
    }

    pub fn parse(data: &[u8], offset: u64, path: &str) -> Result<Self> {
        let head = Block::parse(data, offset, path)?;
        if head.content_type != BlockContentType::SliceHeader {
            return Err(Error::corrupt(
                path,
                offset,
                format!(
                    "expected a slice header block and found {:?}",
                    head.content_type
                ),
            ));
        }
        let header = SliceHeader::parse(&head.data, path)?;

        let mut core = Bytes::new();
        // Clamped by what is left to read, not by the number the header names.
        // A block is at least five bytes of header plus a CRC32, so the count
        // cannot exceed that; `i32::MAX` here is an 85 GB reservation, granted
        // lazily on this machine and an abort on a Linux box with strict
        // overcommit. Same rule as everywhere else in the crate.
        const SMALLEST_BLOCK: usize = 9;
        let declared = header.n_blocks.max(0) as usize;
        let possible = data.len().saturating_sub(head.total_size) / SMALLEST_BLOCK + 1;
        let mut external = Vec::with_capacity(declared.min(possible));
        let mut embedded_reference = None;
        let mut at = head.total_size;
        for _ in 0..header.n_blocks.max(0) {
            let block = Block::parse(
                data.get(at..).ok_or_else(|| {
                    Error::corrupt(path, offset + at as u64, "a slice cut short of its blocks")
                })?,
                offset + at as u64,
                path,
            )?;
            at += block.total_size;
            match block.content_type {
                BlockContentType::Core => core = block.data,
                BlockContentType::External => {
                    if block.content_id == header.embedded_reference_id
                        && header.embedded_reference_id >= 0
                    {
                        embedded_reference = Some(block.data.clone());
                    }
                    external.push((block.content_id, block.data));
                }
                other => {
                    return Err(Error::corrupt(
                        path,
                        offset + at as u64,
                        format!("a {other:?} block inside a slice"),
                    ))
                }
            }
        }
        Ok(Self {
            header,
            offset,
            core,
            external,
            embedded_reference,
        })
    }

    /// The external blocks as the encodings want them.
    pub fn streams(&self) -> impl Iterator<Item = (i32, &[u8])> {
        self.external.iter().map(|(id, data)| (*id, data.as_ref()))
    }
}

/// Where every slice of a container sits, and how long each is.
///
/// The landmarks are offsets from the end of the container header, and the last
/// slice runs to the end of the container — which is the one length the format
/// does not write down anywhere.
pub fn slice_extents(container: &ContainerHeader) -> Vec<(u64, usize)> {
    let mut out = Vec::with_capacity(container.landmarks.len());
    for (i, &landmark) in container.landmarks.iter().enumerate() {
        let start = container.landmark_offset(landmark);
        let end = match container.landmarks.get(i + 1) {
            Some(&next) => container.landmark_offset(next),
            None => container.end_offset(),
        };
        out.push((start, end.saturating_sub(start) as usize));
    }
    out
}

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

    fn header_bytes(ref_id: &[u8], start: u8, span: u8, ids: &[u8]) -> Vec<u8> {
        let mut out = Vec::new();
        out.extend_from_slice(ref_id);
        out.push(start);
        out.push(span);
        out.push(2); // two records
        out.push(0); // record counter
        out.push(ids.len() as u8 + 1); // block count: core plus the externals
        out.push(ids.len() as u8);
        out.extend_from_slice(ids);
        out.extend_from_slice(&[0xff, 0xff, 0xff, 0xff, 0x0f]); // embedded ref: -1
        out.extend_from_slice(&[0u8; 16]);
        out
    }

    #[test]
    fn a_slice_header_reads_its_fields() {
        let bytes = header_bytes(&[0x00], 101, 50, &[11, 12]);
        let header = SliceHeader::parse(&bytes, "test").expect("a slice header");
        assert_eq!(header.ref_id, 0);
        assert_eq!(header.start, 101);
        assert_eq!(header.span, 50);
        assert_eq!(header.n_records, 2);
        assert_eq!(header.block_content_ids, vec![11, 12]);
        assert_eq!(header.embedded_reference_id, -1);
        assert!(!header.has_reference_md5());
        // 1-based in the file, 0-based out.
        assert_eq!(header.range(), Some((100, 150)));
    }

    #[test]
    fn an_unmapped_or_multi_ref_slice_has_no_range_to_report() {
        let unmapped = SliceHeader::parse(
            &header_bytes(&[0xff, 0xff, 0xff, 0xff, 0x0f], 0, 0, &[]),
            "test",
        )
        .expect("unmapped");
        assert_eq!(unmapped.ref_id, UNMAPPED);
        assert_eq!(unmapped.range(), None);
        assert!(!unmapped.is_multi_ref());

        let multi = SliceHeader::parse(
            &header_bytes(&[0xff, 0xff, 0xff, 0xff, 0x0e], 0, 0, &[]),
            "test",
        )
        .expect("multi-ref");
        assert_eq!(multi.ref_id, MULTI_REF);
        assert!(multi.is_multi_ref());
        assert_eq!(multi.range(), None);
    }

    /// A slice header cut short of its MD5 still parses, because nothing that
    /// matters comes after it and sixteen zeros already mean "unchecked".
    #[test]
    fn a_slice_header_without_its_md5_still_parses() {
        let mut bytes = header_bytes(&[0x00], 1, 1, &[11]);
        bytes.truncate(bytes.len() - 16);
        let header = SliceHeader::parse(&bytes, "test").expect("a slice header");
        assert!(!header.has_reference_md5());
    }

    #[test]
    fn slice_extents_run_the_last_slice_to_the_end_of_its_container() {
        let container = ContainerHeader {
            length: 1000,
            ref_id: 0,
            start: 1,
            span: 1,
            n_records: 3,
            record_counter: 0,
            bases: 0,
            n_blocks: 4,
            landmarks: vec![100, 400, 700],
            offset: 26,
            header_len: 30,
        };
        // Blocks begin at 56; the container ends at 1056.
        assert_eq!(
            slice_extents(&container),
            vec![(156, 300), (456, 300), (756, 300)]
        );
    }
}