gwseq-io 0.2.0

//! The BAI index.
//!
//! The binning index, the 16 kbp linear index, and the pseudo-bin 37450 that
//! carries per-reference metadata rather than alignments.

use indexmap::IndexMap;

use crate::bam::bgzf::{Chunk, VirtualOffset};
use crate::bytes::LeCursor;
use crate::error::{Error, Result};
use crate::source::ByteSource;

/// Pseudo-bin holding a reference's metadata rather than its alignments.
pub const MAGIC_BIN: u32 = 37450;

/// BAI bins cover 2^29 bases, so nothing past that is indexed.
///
/// Bounding the request against it is not tidiness: `ends=[10**12]` used to push
/// about 61 M bins through the lookup before finding that almost none of them
/// exist.
const BAI_MAX_POSITION: u64 = 1 << 29;

/// Largest run of merely *adjacent* chunks to join into one, in compressed
/// bytes. Chunks that truly overlap are joined whatever this says, since leaving
/// them apart would hand the same records out twice.
///
/// It is a request covering most of a reference this is for: its chunks are
/// contiguous and would otherwise come back as one, to be decompressed whole
/// before a single record of it is read.
pub const MAX_MERGE_SPAN: u64 = 64 * 1024 * 1024;

#[derive(Debug, Clone, Default)]
pub struct RefIndex {
    /// Chunks by bin number.
    pub bins: IndexMap<u32, Vec<Chunk>>,
    /// The linear index: the earliest virtual offset an alignment overlapping
    /// each 16 kbp window can start at.
    pub linear: Vec<VirtualOffset>,
    /// From the metadata pseudo-bin, when the file carries one.
    pub metadata: Option<RefMetadata>,
}

#[derive(Debug, Clone, Copy)]
pub struct RefMetadata {
    pub ref_start: VirtualOffset,
    pub ref_end: VirtualOffset,
    pub mapped: u64,
    pub unmapped: u64,
}

#[derive(Debug, Clone, Default)]
pub struct BamIndex {
    pub refs: Vec<RefIndex>,
    /// Unplaced unmapped reads, when the file records the count.
    pub unplaced_count: Option<u64>,
}

impl BamIndex {
    /// Read a `.bai`. The index itself is not compressed.
    pub fn read(source: &dyn ByteSource) -> Result<Self> {
        let path = source.path();
        let all = source.read_to_end(0)?;
        let mut c = LeCursor::new(&all, 0, path);

        if c.take(4)? != b"BAI\x01" {
            return Err(Error::format(path, "invalid bam index magic"));
        }
        let n_ref = c.read_u32()? as usize;
        let mut refs = Vec::with_capacity(n_ref.min(1 << 16));
        for _ in 0..n_ref {
            let mut index = RefIndex::default();
            let n_bin = c.read_u32()? as usize;
            for _ in 0..n_bin {
                let bin = c.read_u32()?;
                let n_chunk = c.read_u32()? as usize;
                if bin == MAGIC_BIN {
                    if n_chunk != 2 {
                        return Err(Error::corrupt(
                            path,
                            c.file_offset(),
                            "invalid metadata pseudo-bin",
                        ));
                    }
                    index.metadata = Some(RefMetadata {
                        ref_start: VirtualOffset(c.read_u64()?),
                        ref_end: VirtualOffset(c.read_u64()?),
                        mapped: c.read_u64()?,
                        unmapped: c.read_u64()?,
                    });
                    continue;
                }
                let mut chunks = Vec::with_capacity(n_chunk.min(1 << 16));
                for _ in 0..n_chunk {
                    chunks.push(Chunk {
                        begin: VirtualOffset(c.read_u64()?),
                        end: VirtualOffset(c.read_u64()?),
                    });
                }
                index.bins.insert(bin, chunks);
            }
            let n_intv = c.read_u32()? as usize;
            index.linear = Vec::with_capacity(n_intv.min(1 << 20));
            for _ in 0..n_intv {
                index.linear.push(VirtualOffset(c.read_u64()?));
            }
            refs.push(index);
        }
        // Optional, at the very end, and absent from plenty of real files.
        let unplaced_count = if c.remaining() >= 8 {
            Some(c.read_u64()?)
        } else {
            None
        };
        Ok(Self {
            refs,
            unplaced_count,
        })
    }

    /// The chunks that may hold alignments overlapping `[start, end)` on
    /// `ref_index`, merged and in file order.
    pub fn chunks(
        &self,
        ref_index: usize,
        start: i64,
        end: i64,
        max_merge_span: Option<u64>,
    ) -> Result<Vec<Chunk>> {
        let reference = self
            .refs
            .get(ref_index)
            .ok_or_else(|| Error::invalid(format!("ref index {ref_index} out of range")))?;

        // A locus that ends before it starts is a request no answer fits, and
        // an empty list is an answer — the one a caller reads as "nothing
        // aligned here". bbi refuses the same shape in `IndexedLocs::build`,
        // and the two formats should not disagree about what a locus is.
        if end < start {
            return Err(Error::invalid(format!(
                "Locus {start}-{end} ends before it starts"
            )));
        }
        // Clamped at 0, since a window may be asked for around a position near
        // the start of a reference, and bounded above by what the index can
        // address at all. An empty locus past that point genuinely has nothing
        // in it, which is not the same as a reversed one.
        let bounded_start = start.max(0) as u64;
        let bounded_end = (end.max(0) as u64).min(BAI_MAX_POSITION);
        if bounded_start >= BAI_MAX_POSITION || bounded_end <= bounded_start {
            return Ok(Vec::new());
        }

        // No alignment overlapping the region starts before this.
        let min_offset = if reference.linear.is_empty() {
            VirtualOffset(0)
        } else {
            let window = (bounded_start >> 14) as usize;
            // A window past the last one the linear index covers means no
            // alignment starts that late, so the last entry is as far as the
            // file needs to be read from.
            *reference
                .linear
                .get(window)
                .unwrap_or_else(|| reference.linear.last().expect("checked non-empty"))
        };

        let mut chunks: Vec<Chunk> = Vec::new();
        for bin in reg2bins(bounded_start, bounded_end) {
            let Some(bin_chunks) = reference.bins.get(&bin) else {
                continue;
            };
            chunks.extend(bin_chunks.iter().copied().filter(|c| c.end >= min_offset));
        }
        chunks.sort_by_key(|c| c.begin);
        Ok(merge(chunks, max_merge_span))
    }
}

/// Merge overlapping chunks always, and adjacent ones while they stay under
/// `max_merge_span` compressed bytes.
fn merge(chunks: Vec<Chunk>, max_merge_span: Option<u64>) -> Vec<Chunk> {
    if chunks.len() <= 1 {
        return chunks;
    }
    let mut merged: Vec<Chunk> = Vec::with_capacity(chunks.len());
    merged.push(chunks[0]);
    for current in &chunks[1..] {
        let last = merged.last_mut().expect("pushed one above");
        if current.begin < last.end {
            // Overlapping, so they have records in common; merging is what keeps
            // those records from being read twice.
            last.end = last.end.max(current.end);
            continue;
        }
        // Compressed bytes the two would span together, which is what the reader
        // has to hold to decompress the result in one piece.
        let span = last.end.max(current.end).block_offset() - last.begin.block_offset();
        let within_budget = max_merge_span.is_none_or(|budget| span <= budget);
        if current.begin == last.end && within_budget {
            // Adjacent, and joining them saves a read; they share no record, so
            // leaving them apart would also have been right.
            last.end = last.end.max(current.end);
        } else {
            merged.push(*current);
        }
    }
    merged
}

/// Every bin that may hold an alignment overlapping `[start, end)`.
///
/// An empty interval comes back with no bins rather than being decremented:
/// `end - 1` on an unsigned zero is a coordinate near the top of the range, and
/// the loops below would then walk the whole of it — some 2^38 iterations,
/// pushing into the list the entire way.
pub fn reg2bins(start: u64, end: u64) -> Vec<u32> {
    let mut bins = Vec::new();
    if end <= start {
        return bins;
    }
    let end = end - 1;
    bins.push(0);
    for (offset, shift) in [(1u64, 26u32), (9, 23), (73, 20), (585, 17), (4681, 14)] {
        for bin in (offset + (start >> shift))..=(offset + (end >> shift)) {
            bins.push(bin as u32);
        }
    }
    bins
}

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

    fn vo(block: u64, within: u16) -> VirtualOffset {
        VirtualOffset::new(block, within)
    }

    fn chunk(a: u64, b: u64) -> Chunk {
        Chunk {
            begin: vo(a, 0),
            end: vo(b, 0),
        }
    }

    /// A one-reference index with the given bins and linear windows.
    fn bai(bins: &[(u32, &[Chunk])], linear: &[VirtualOffset], tail: bool) -> Vec<u8> {
        let mut b = b"BAI\x01".to_vec();
        b.extend_from_slice(&1u32.to_le_bytes()); // n_ref
        b.extend_from_slice(&(bins.len() as u32).to_le_bytes());
        for (bin, chunks) in bins {
            b.extend_from_slice(&bin.to_le_bytes());
            b.extend_from_slice(&(chunks.len() as u32).to_le_bytes());
            for c in *chunks {
                b.extend_from_slice(&c.begin.0.to_le_bytes());
                b.extend_from_slice(&c.end.0.to_le_bytes());
            }
        }
        b.extend_from_slice(&(linear.len() as u32).to_le_bytes());
        for offset in linear {
            b.extend_from_slice(&offset.0.to_le_bytes());
        }
        if tail {
            b.extend_from_slice(&42u64.to_le_bytes());
        }
        b
    }

    #[test]
    fn reads_bins_the_linear_index_and_the_optional_tail() {
        let bytes = bai(
            &[(4681, &[chunk(100, 200)]), (4682, &[chunk(200, 300)])],
            &[vo(0, 0), vo(100, 0)],
            true,
        );
        let index = BamIndex::read(&MemorySource::new(bytes)).unwrap();
        assert_eq!(index.refs.len(), 1);
        assert_eq!(index.refs[0].bins.len(), 2);
        assert_eq!(index.refs[0].linear.len(), 2);
        assert_eq!(index.unplaced_count, Some(42));
    }

    #[test]
    fn the_optional_tail_really_is_optional() {
        let bytes = bai(&[(4681, &[chunk(100, 200)])], &[vo(0, 0)], false);
        let index = BamIndex::read(&MemorySource::new(bytes)).unwrap();
        assert_eq!(index.unplaced_count, None);
    }

    #[test]
    fn the_metadata_pseudo_bin_is_not_a_bin_of_chunks() {
        let mut b = b"BAI\x01".to_vec();
        b.extend_from_slice(&1u32.to_le_bytes());
        b.extend_from_slice(&2u32.to_le_bytes()); // two "bins"
        b.extend_from_slice(&4681u32.to_le_bytes());
        b.extend_from_slice(&1u32.to_le_bytes());
        b.extend_from_slice(&vo(100, 0).0.to_le_bytes());
        b.extend_from_slice(&vo(200, 0).0.to_le_bytes());
        b.extend_from_slice(&MAGIC_BIN.to_le_bytes());
        b.extend_from_slice(&2u32.to_le_bytes());
        for value in [7u64, 8, 9, 10] {
            b.extend_from_slice(&value.to_le_bytes());
        }
        b.extend_from_slice(&0u32.to_le_bytes()); // n_intv

        let index = BamIndex::read(&MemorySource::new(b)).unwrap();
        assert_eq!(index.refs[0].bins.len(), 1, "the pseudo-bin is not a bin");
        let meta = index.refs[0].metadata.unwrap();
        assert_eq!((meta.mapped, meta.unmapped), (9, 10));
    }

    #[test]
    fn a_bad_magic_is_refused() {
        let err = BamIndex::read(&MemorySource::new(b"NOPE".to_vec()))
            .unwrap_err()
            .to_string();
        assert!(err.contains("invalid bam index magic"), "{err}");
    }

    #[test]
    fn a_truncated_index_is_corrupt_not_a_panic() {
        let mut bytes = bai(&[(4681, &[chunk(100, 200)])], &[vo(0, 0)], false);
        bytes.truncate(bytes.len() - 5);
        assert!(matches!(
            BamIndex::read(&MemorySource::new(bytes)),
            Err(Error::Corrupt { .. })
        ));
    }

    #[test]
    fn reg2bins_covers_every_level_and_refuses_an_empty_region() {
        assert!(reg2bins(100, 100).is_empty());
        assert!(reg2bins(100, 50).is_empty());
        // The root, and one bin at each of the five levels.
        assert_eq!(reg2bins(0, 16384), [0, 1, 9, 73, 585, 4681]);
        // A region reaching into the next 16 kbp window adds one more.
        assert_eq!(reg2bins(0, 16385), [0, 1, 9, 73, 585, 4681, 4682]);
    }

    #[test]
    fn a_region_past_what_the_index_addresses_asks_for_nothing() {
        let index = BamIndex::read(&MemorySource::new(bai(
            &[(4681, &[chunk(100, 200)])],
            &[vo(0, 0)],
            false,
        )))
        .unwrap();
        // Unclamped, a query at 10^12 enumerates 61 M bins.
        assert!(index
            .chunks(0, 1_000_000_000_000, 1_000_000_001_000, None)
            .unwrap()
            .is_empty());
        assert!(index
            .chunks(0, 1 << 30, (1 << 30) + 10, None)
            .unwrap()
            .is_empty());
    }

    #[test]
    fn the_linear_index_drops_chunks_that_end_before_the_region_can_start() {
        // Bin 4681 is the first 16 kbp window; its chunk ends at block 100.
        // The linear index says nothing overlapping window 0 starts before 500.
        let index = BamIndex::read(&MemorySource::new(bai(
            &[(4681, &[chunk(50, 100)])],
            &[vo(500, 0)],
            false,
        )))
        .unwrap();
        assert!(index.chunks(0, 0, 1000, None).unwrap().is_empty());
    }

    #[test]
    fn overlapping_chunks_merge_whatever_the_budget() {
        assert_eq!(
            merge(vec![chunk(0, 200), chunk(100, 300)], Some(0)),
            [chunk(0, 300)]
        );
    }

    #[test]
    fn adjacent_chunks_merge_only_inside_the_budget() {
        // Adjacent: the first ends exactly where the second begins.
        assert_eq!(
            merge(vec![chunk(0, 100), chunk(100, 200)], Some(1000)),
            [chunk(0, 200)]
        );
        // The same pair with a budget too small to span them stays apart.
        assert_eq!(
            merge(vec![chunk(0, 100), chunk(100, 200)], Some(50)),
            [chunk(0, 100), chunk(100, 200)]
        );
        // No budget at all means no limit.
        assert_eq!(
            merge(vec![chunk(0, 100), chunk(100, 200)], None),
            [chunk(0, 200)]
        );
    }

    #[test]
    fn chunks_with_a_gap_between_them_never_merge() {
        let apart = vec![chunk(0, 100), chunk(500, 600)];
        assert_eq!(merge(apart.clone(), None), apart);
    }

    #[test]
    fn an_out_of_range_reference_is_refused_by_name() {
        let index = BamIndex::default();
        let err = index.chunks(3, 0, 10, None).unwrap_err().to_string();
        assert!(err.contains("ref index 3 out of range"), "{err}");
    }
}