gwseq-io 0.2.1

//! The CRAM index.
//!
//! Format reference: `docs/cram_format_v3.1.md` §12.
//!
//! A `.crai` is a gzipped tab-separated file with one line per slice — not a
//! binning index like a `.bai`, just a list. §12 is explicit that this sits
//! outside the format proper and may be replaced, which is part of why it is so
//! plain.
//!
//! Two things follow from its shape. A line gives the slice's offset *relative
//! to the end of its container header*, and the header's length is not in the
//! index, so reaching a slice means reading its container header first — the
//! two reads §12 warns about, and the same header the compression header is
//! read from anyway. And because it is a list rather than a tree, a query is a
//! binary search over one reference's slices rather than a bin walk.
//!
//! # Building one
//!
//! §12: "Indexing a CRAM file is deemed to be a lightweight operation because
//! it usually does not require any CRAM records to be read." That is true
//! enough to be worth acting on — [`CramIndex::build`] walks the container
//! headers and makes the index, so a file with no `.crai` beside it is still
//! randomly accessible. It costs one small read per container, which is a
//! thousand reads for a ten-million-read file and imperceptible locally.

use std::collections::HashMap;
use std::io::Read as _;

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

use super::container::{ContainerHeader, FILE_DEFINITION_SIZE};
use super::slice::{Slice, MULTI_REF};

/// A container's compression header, which is always its first block.
fn container_compression(
    source: &dyn ByteSource,
    container: &ContainerHeader,
    path: &str,
) -> Result<super::compression::CompressionHeader> {
    let first = container
        .landmarks
        .first()
        .copied()
        .map(|landmark| landmark.max(0) as usize)
        .unwrap_or(container.length.max(0) as usize);
    let data = source.read_exact_at(container.blocks_offset(), first.max(1))?;
    let block = super::container::Block::parse(&data, container.blocks_offset(), path)?;
    super::compression::CompressionHeader::parse(&block.data, path)
}

/// One slice, as the index names it.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct IndexEntry {
    pub ref_id: i32,
    /// 0-based half-open span on the reference. Meaningless, and both zero, for
    /// an unmapped or multi-reference slice.
    pub start: i64,
    pub end: i64,
    /// Absolute offset of the container header.
    pub container_offset: u64,
    /// The slice's offset from the end of that header.
    pub landmark: u64,
    /// Bytes of the slice, its header block included.
    pub size: usize,
    /// The largest `end` of this entry and every entry before it in start
    /// order — htslib's "extended end".
    ///
    /// Set by [`CramIndex::sort`], not by whoever built the entry, because it
    /// is a property of the list and not of the slice. It exists to make the
    /// binary search in [`CramIndex::slices`] legal: `end` alone is not
    /// monotone over a list sorted by `start`, and a running maximum is, by
    /// construction.
    max_end: i64,
}

/// Every slice of a file, by reference.
#[derive(Debug, Clone, Default)]
pub struct CramIndex {
    /// Sorted by start within each reference.
    by_ref: HashMap<i32, Vec<IndexEntry>>,
    /// Slices holding more than one reference, which every query must consider
    /// because only the records inside say what they hold.
    multi: Vec<IndexEntry>,
}

impl CramIndex {
    /// Parse a `.crai`, which is gzip whatever it is called.
    pub fn parse(data: &[u8], path: &str) -> Result<Self> {
        let mut text = Vec::new();
        if crate::source::is_gzipped(data) {
            // Bounded: a `.crai` is a line per slice, so a thousand slices is
            // forty kilobytes and this ceiling is four orders of magnitude
            // past any real one. Without it a small file claiming to be an
            // index inflates without limit.
            const MAX_INDEX_SIZE: u64 = 256 << 20;
            let read = flate2::read::MultiGzDecoder::new(data)
                .take(MAX_INDEX_SIZE + 1)
                .read_to_end(&mut text)
                .map_err(|e| Error::format(path, format!("could not inflate this index: {e}")))?;
            if read as u64 > MAX_INDEX_SIZE {
                return Err(Error::format(
                    path,
                    format!("this index inflates past the {MAX_INDEX_SIZE} bytes any index needs"),
                ));
            }
        } else {
            text.extend_from_slice(data);
        }

        let mut index = Self::default();
        for (number, line) in text.split(|b| *b == b'\n').enumerate() {
            let line = line.strip_suffix(b"\r").unwrap_or(line);
            if line.is_empty() {
                continue;
            }
            let mut fields = line.split(|b| *b == b'\t');
            let mut next = |what: &str| -> Result<i64> {
                fields
                    .next()
                    .and_then(|f| std::str::from_utf8(f).ok())
                    .and_then(|s| s.trim().parse::<i64>().ok())
                    .ok_or_else(|| {
                        Error::format(
                            path,
                            format!("line {} of this index has no readable {what}", number + 1),
                        )
                    })
            };
            let ref_id = next("reference id")? as i32;
            let start = next("alignment start")?;
            let span = next("alignment span")?;
            let container_offset = next("container offset")?.max(0) as u64;
            let landmark = next("slice offset")?.max(0) as u64;
            let size = next("slice size")?.max(0) as usize;
            index.push(IndexEntry {
                ref_id,
                // The index is 1-based like the rest of the format.
                start: (start - 1).max(0),
                end: (start - 1).max(0) + span.max(0),
                container_offset,
                landmark,
                size,
                // Replaced by `sort`, which is the only thing that can know it.
                max_end: i64::MIN,
            });
        }
        if index.is_empty() {
            return Err(Error::format(path, "this index lists no slices"));
        }
        index.sort();
        Ok(index)
    }

    /// Build an index by walking the file's container headers.
    ///
    /// No records are read. A container holding one slice — which is what
    /// `samtools` writes — is described by its header alone; a container with
    /// several needs each slice header block, which are small and uncompressed.
    ///
    /// The walk starts *after* the header container, not at the file
    /// definition. §8 makes the first container the SAM header's, and it holds
    /// no slices — but `samtools` still gives it landmarks (two, for a header
    /// block and the padding that lets the header be rewritten in place), so
    /// there is nothing in its header to distinguish it from a container of
    /// two slices. Walking from the file definition indexes the header block
    /// as if it were a slice, and the failure surfaces much later as a slice
    /// header that is a `FileHeader`.
    pub fn build(source: &dyn ByteSource) -> Result<Self> {
        let path = source.path();
        let length = source.len()?;
        let mut index = Self::default();
        let header = ContainerHeader::read(source, FILE_DEFINITION_SIZE as u64)?;
        let mut offset = header.end_offset();

        while offset < length {
            let container = ContainerHeader::read(source, offset)?;
            if container.is_eof() {
                break;
            }
            let extents = super::slice::slice_extents(&container);
            // A container with one slice is its own description, so its header
            // spares a read of the slice's.
            if extents.len() == 1 && container.ref_id != MULTI_REF {
                let (start, size) = extents[0];
                index.push(IndexEntry {
                    ref_id: container.ref_id,
                    start: (i64::from(container.start) - 1).max(0),
                    end: (i64::from(container.start) - 1).max(0) + i64::from(container.span).max(0),
                    container_offset: container.offset,
                    landmark: start - container.blocks_offset(),
                    size,
                    // Replaced by `sort`, which is the only thing that can know it.
                    max_end: i64::MIN,
                });
            } else {
                // The compression header, needed only to read a
                // multi-reference slice's `RI`. Read once per container and
                // only when there is one to read.
                let mut compression = None;
                for (start, size) in extents {
                    // The header alone, which is all a single-reference slice
                    // needs; a multi-reference one is re-read in full below.
                    let header = Slice::read_header(source, start, size)?;
                    let landmark = start - container.blocks_offset();
                    if header.is_multi_ref() {
                        let slice = Slice::read(source, start, size)?;
                        // §12: "the exception to this is with multi-reference
                        // containers, where the RI data series must be read".
                        // Filed per reference with its real span, which is what
                        // `samtools index` writes and what stops every query
                        // having to consider every one of them. `samtools`
                        // puts one at each chromosome boundary, so a sorted
                        // file has one per chromosome and a locus that took
                        // one slice was taking all of them.
                        if compression.is_none() {
                            compression = Some(container_compression(source, &container, path)?);
                        }
                        let header = compression.as_ref().expect("just read");
                        for (ref_id, from, to) in super::record::slice_spans(&slice, header, path)?
                        {
                            index.push(IndexEntry {
                                ref_id,
                                start: from,
                                end: to,
                                container_offset: container.offset,
                                landmark,
                                size,
                                max_end: i64::MIN,
                            });
                        }
                        continue;
                    }
                    let (from, to) = header.range().unwrap_or((0, 0));
                    index.push(IndexEntry {
                        ref_id: header.ref_id,
                        start: from,
                        end: to,
                        container_offset: container.offset,
                        landmark,
                        size,
                        // Replaced by `sort`, which is the only thing that can know it.
                        max_end: i64::MIN,
                    });
                }
            }
            let next = container.end_offset();
            if next <= offset {
                return Err(Error::corrupt(
                    path,
                    offset,
                    "a container that does not advance, so the file cannot be walked",
                ));
            }
            offset = next;
        }
        index.sort();
        Ok(index)
    }

    fn push(&mut self, entry: IndexEntry) {
        if entry.ref_id == MULTI_REF {
            // A `.crai` lists a multi-reference slice once per reference it
            // touches, with the real reference id rather than -2 — so the two
            // ways one gets here need the same de-duplication.
            self.multi.push(entry);
        } else {
            self.by_ref.entry(entry.ref_id).or_default().push(entry);
        }
    }

    fn sort(&mut self) {
        for entries in self.by_ref.values_mut() {
            entries.sort_by_key(|e| (e.start, e.end, e.container_offset, e.landmark));
            let mut running = i64::MIN;
            for entry in entries.iter_mut() {
                running = running.max(entry.end);
                entry.max_end = running;
            }
        }
        self.multi.sort_by_key(|e| (e.container_offset, e.landmark));
        self.multi
            .dedup_by_key(|e| (e.container_offset, e.landmark));
    }

    pub fn is_empty(&self) -> bool {
        self.by_ref.is_empty() && self.multi.is_empty()
    }

    /// How many slices are filed as "could be on any reference".
    ///
    /// Zero for an index `build` made, since it reads the `RI` of every
    /// multi-reference slice and files it properly. A `.crai` from
    /// `samtools index` is also zero, for the same reason. Anything above zero
    /// is a slice every query has to consider.
    pub fn multi_len(&self) -> usize {
        self.multi.len()
    }

    /// How many slices are listed.
    pub fn len(&self) -> usize {
        self.by_ref.values().map(Vec::len).sum::<usize>() + self.multi.len()
    }

    /// Every reference the index mentions.
    pub fn references(&self) -> Vec<i32> {
        let mut out: Vec<i32> = self.by_ref.keys().copied().collect();
        out.sort_unstable();
        out
    }

    /// The slices that could hold alignments overlapping `[start, end)`.
    ///
    /// De-duplicated and in file order, so a caller reading several loci at
    /// once does not decode one slice twice. Multi-reference slices are always
    /// included: only their records say what they hold, and reading them is how
    /// one finds out.
    pub fn slices(&self, ref_id: i32, start: i64, end: i64) -> Vec<IndexEntry> {
        let mut out: Vec<IndexEntry> = Vec::new();
        if let Some(entries) = self.by_ref.get(&ref_id) {
            // Slices of one reference may overlap, so the scan starts from the
            // first that could reach the region rather than from the first that
            // starts in it.
            //
            // The key is `max_end` and not `end`, and that is not a
            // refinement. `partition_point` requires a predicate that is true
            // for a prefix and false for the rest; over a list sorted by
            // `start`, `end <= start` is not, as soon as one slice ends after a
            // slice that begins later. Its result is then *unspecified* — not
            // approximate — so a query could skip any number of overlapping
            // slices. A running maximum restores the partition, and a slice
            // whose `max_end` is past the query start is the earliest that can
            // reach it.
            //
            // One spliced read is enough to produce the shape: a slice's end is
            // the end of its longest record, and a 100 kb `N` skip is ordinary
            // RNA-seq.
            let first = entries.partition_point(|e| e.max_end <= start);
            for entry in &entries[first..] {
                if entry.start >= end {
                    break;
                }
                if entry.end > start {
                    out.push(*entry);
                }
            }
        }
        out.extend(self.multi.iter().copied());
        out.sort_by_key(|e| (e.container_offset, e.landmark));
        out.dedup_by_key(|e| (e.container_offset, e.landmark));
        out
    }

    /// Every slice, in file order.
    pub fn all(&self) -> Vec<IndexEntry> {
        let mut out: Vec<IndexEntry> = self
            .by_ref
            .values()
            .flatten()
            .chain(self.multi.iter())
            .copied()
            .collect();
        out.sort_by_key(|e| (e.container_offset, e.landmark));
        out.dedup_by_key(|e| (e.container_offset, e.landmark));
        out
    }
}

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

    const LINES: &[u8] =
        b"0\t101\t50\t1000\t200\t300\n0\t201\t50\t1000\t500\t300\n1\t1\t10\t2000\t200\t100\n";

    #[test]
    fn a_crai_parses_into_slices_by_reference() {
        let index = CramIndex::parse(LINES, "test.crai").expect("an index");
        assert_eq!(index.len(), 3);
        assert_eq!(index.references(), vec![0, 1]);
        let first = index.slices(0, 100, 110);
        assert_eq!(first.len(), 1);
        assert_eq!(first[0].container_offset, 1000);
        assert_eq!(first[0].landmark, 200);
        assert_eq!(first[0].size, 300);
        // 1-based in the file, 0-based out.
        assert_eq!(first[0].start, 100);
        assert_eq!(first[0].end, 150);
    }

    /// A slice whose span reaches past the slices after it.
    ///
    /// One spliced read makes this shape — a slice's end is the end of its
    /// longest record, and a 100 kb `N` skip is ordinary RNA-seq — so it is
    /// not a hostile-file case. Searching on `end` rather than a running
    /// maximum leaves `partition_point` with an unpartitioned predicate, whose
    /// result is unspecified: here it skipped the long slice entirely and
    /// every locus inside its span but past the next slice's end silently lost
    /// every read in it.
    const NESTED: &[u8] = b"0\t1000\t199000\t1000\t200\t300\n\
                            0\t1200\t200\t2000\t200\t300\n\
                            0\t2200\t200\t3000\t200\t300\n";

    #[test]
    fn a_slice_spanning_past_later_slices_is_still_found() {
        let index = CramIndex::parse(NESTED, "nested.crai").expect("an index");
        let long = |v: &Vec<IndexEntry>| v.iter().any(|e| e.container_offset == 1000);

        // Deep inside the long slice, past everything else.
        let far = index.slices(0, 150_000, 150_100);
        assert_eq!(far.len(), 1, "the long slice covers this locus alone");
        assert!(long(&far));

        // Past the second slice's end, still inside the first.
        let middle = index.slices(0, 5_000, 5_100);
        assert_eq!(middle.len(), 1);
        assert!(long(&middle));

        // Overlapping the third slice: both it and the long one. The third
        // covers 1-based 2200..2399, so 0-based [2199, 2399).
        let both = index.slices(0, 2_300, 2_600);
        assert_eq!(both.len(), 2, "the long slice and the third");
        assert!(long(&both));
        assert!(both.iter().any(|e| e.container_offset == 3000));

        // And the ordinary case still answers one slice, not three.
        let near = index.slices(0, 1_250, 1_260);
        assert_eq!(near.len(), 2, "the long slice and the second");
    }

    /// The same list reached through `build` rather than `parse`.
    ///
    /// `sort` is shared by both routes, and it is the only thing that can set
    /// the running maximum, so this pins the invariant rather than the parser.
    #[test]
    fn the_extended_end_is_a_running_maximum_in_start_order() {
        let index = CramIndex::parse(NESTED, "nested.crai").expect("an index");
        let entries = index.by_ref.get(&0).expect("reference 0");
        let ends: Vec<i64> = entries.iter().map(|e| e.max_end).collect();
        // 1-based 1000 + span 199 000, half-open and 0-based, is 199 999.
        assert_eq!(ends, vec![199_999, 199_999, 199_999]);
        assert!(
            ends.windows(2).all(|w| w[0] <= w[1]),
            "the search key must not decrease"
        );
    }

    #[test]
    fn a_query_returns_only_the_slices_that_overlap_it() {
        let index = CramIndex::parse(LINES, "test.crai").expect("an index");
        assert_eq!(index.slices(0, 0, 100).len(), 0);
        assert_eq!(index.slices(0, 100, 101).len(), 1);
        assert_eq!(index.slices(0, 149, 151).len(), 1);
        assert_eq!(index.slices(0, 100, 300).len(), 2);
        assert_eq!(index.slices(0, 250, 300).len(), 0);
        assert_eq!(index.slices(2, 0, 1000).len(), 0);
    }

    /// A `.crai` lists a multi-reference slice once per reference it touches.
    /// Two queries must not decode it twice, and neither must one.
    #[test]
    fn a_multi_reference_slice_is_listed_once_however_often_it_appears() {
        let lines =
            b"-2\t0\t0\t1000\t200\t300\n-2\t0\t0\t1000\t200\t300\n0\t101\t50\t2000\t200\t300\n";
        let index = CramIndex::parse(lines, "test.crai").expect("an index");
        let slices = index.slices(0, 100, 110);
        assert_eq!(slices.len(), 2, "{slices:?}");
        assert_eq!(index.all().len(), 2);
        // And a query on a reference with no slices of its own still gets it.
        assert_eq!(index.slices(5, 0, 1000).len(), 1);
    }

    #[test]
    fn a_gzipped_index_reads_the_same_as_a_plain_one() {
        use std::io::Write as _;
        let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::fast());
        encoder.write_all(LINES).expect("gzip");
        let gzipped = encoder.finish().expect("gzip");
        let index = CramIndex::parse(&gzipped, "test.crai").expect("an index");
        assert_eq!(index.len(), 3);
    }

    #[test]
    fn an_unreadable_index_is_refused_rather_than_half_read() {
        assert!(CramIndex::parse(b"", "test.crai").is_err());
        assert!(CramIndex::parse(b"0\t101\t50\n", "test.crai").is_err());
        assert!(CramIndex::parse(b"not an index at all\n", "test.crai").is_err());
    }

    #[test]
    fn slices_come_back_in_file_order_whatever_order_the_index_lists_them() {
        let lines =
            b"0\t201\t50\t3000\t200\t300\n0\t101\t50\t1000\t200\t300\n0\t151\t50\t2000\t200\t300\n";
        let index = CramIndex::parse(lines, "test.crai").expect("an index");
        let offsets: Vec<u64> = index
            .slices(0, 0, 1000)
            .iter()
            .map(|e| e.container_offset)
            .collect();
        assert_eq!(offsets, vec![1000, 2000, 3000]);
    }
}