gwseq-io 0.2.0

//! BGZF: gzip members with an extra field carrying each block's compressed
//! size, so a virtual offset can address a record inside one.
//!
//! Kept out of `reader.rs`, because it is a distinct format concern — a 48-bit block offset with a 16-bit offset inside it, an
//! EOF marker block, a header whose magic has to be checked before its size
//! field is trusted — and off-by-one here reads plausible garbage.

use std::io::Read;

use bytes::{Bytes, BytesMut};

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

pub const HEADER_SIZE: usize = 18;
pub const EOF_SIZE: usize = 28;
/// The largest a BGZF block can be: `BSIZE` is a `u16` holding size − 1.
pub const MAX_BLOCK_SIZE: usize = 65536;

/// The 28-byte empty block every BAM ends with.
///
/// A file without it is truncated, and the reader says so at open rather than
/// after walking to the last record.
pub static BGZF_EOF_BLOCK: [u8; EOF_SIZE] = [
    0x1f, 0x8b, 0x08, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0x06, 0x00, 0x42, 0x43, 0x02, 0x00,
    0x1b, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
];

/// A BAI virtual offset: the block's file offset in the top 48 bits, the offset
/// within the decompressed block in the low 16.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
pub struct VirtualOffset(pub u64);

impl VirtualOffset {
    #[inline]
    pub fn block_offset(self) -> u64 {
        self.0 >> 16
    }

    #[inline]
    pub fn within_block(self) -> usize {
        (self.0 & 0xFFFF) as usize
    }

    #[inline]
    pub fn new(block_offset: u64, within_block: u16) -> Self {
        VirtualOffset((block_offset << 16) | within_block as u64)
    }
}

/// Check that a block header is one, before its size field is believed.
///
/// The block size drives the whole walk and is read out of these very bytes, so
/// on a file that is not BGZF — or one whose blocks do not start where the index
/// says — whatever sits there would be taken for a length.
fn check_block_header(head: &[u8], path: &str, at: u64) -> Result<()> {
    let ok = head[0] == 0x1f
        && head[1] == 0x8b
        && head[2] == 8
        && (head[3] & 0x04) != 0
        && head[12] == b'B'
        && head[13] == b'C'
        && u16::from_le_bytes([head[14], head[15]]) == 2;
    if ok {
        Ok(())
    } else {
        Err(Error::corrupt(
            path,
            at,
            format!(
                "no bgzf block header at {at}, so the file is corrupt or its index \
                 does not belong to it"
            ),
        ))
    }
}

/// Decompressed size of a BGZF block, from its `BSIZE` field.
///
/// Widened before the increment: `BSIZE` holds the size less one, so a
/// spec-legal 64 KiB block stores 65535 and a `u16` sum would wrap it to 0 and
/// stall the walk.
#[inline]
fn block_size(head: &[u8]) -> usize {
    u16::from_le_bytes([head[16], head[17]]) as usize + 1
}

/// A half-open run of the file, as the BAI describes one.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Chunk {
    pub begin: VirtualOffset,
    pub end: VirtualOffset,
}

impl Chunk {
    /// Compressed bytes the chunk spans, which is what a run of them weighs.
    #[inline]
    pub fn compressed_size(&self) -> u64 {
        self.end
            .block_offset()
            .saturating_sub(self.begin.block_offset())
    }
}

/// The alignment records one index chunk holds, decompressed and trimmed to the
/// two virtual offsets it lies between.
///
/// Returns whole records starting at a record boundary — which is what a virtual
/// offset names — so the result can be handed straight to the record decoder.
pub fn decompress_chunk(source: &dyn ByteSource, chunk: Chunk, path: &str) -> Result<Bytes> {
    let first_block = chunk.begin.block_offset();
    let last_block = chunk.end.block_offset();
    let begin_within = chunk.begin.within_block();
    let end_within = chunk.end.within_block();

    // A whole block past the chunk's own end, so the last one arrives complete
    // however far into it the chunk stops. The end offset is never before the
    // start, so this is at least one block on its own.
    let wanted = (last_block - first_block) as usize + MAX_BLOCK_SIZE;
    // Short at the end of the file, which is where a read deliberately asking
    // for one block more than the chunk spans has nothing behind it.
    let raw = source.read_at(first_block, wanted)?;

    let mut out = BytesMut::new();
    let mut index = 0usize;
    // Every way out of this loop but the two breaks below is an error: a silent
    // break on "not enough data" would let a truncated file come back as
    // however many alignments survived the cut.
    loop {
        let at = first_block + index as u64;
        // Past the last block the chunk names: a virtual offset points at a
        // record *inside* the block it names, so nothing beyond it is ours.
        if at > last_block {
            break;
        }
        // The chunk ends on a block boundary, so the block starting there holds
        // none of it and need not even be present.
        if at == last_block && end_within == 0 {
            break;
        }

        if index + HEADER_SIZE > raw.len() {
            return Err(Error::corrupt(
                path,
                at,
                format!(
                    "the bgzf block at {at} is cut short (its {HEADER_SIZE}-byte header \
                     does not fit what is left of the file)"
                ),
            ));
        }
        let head = &raw[index..index + HEADER_SIZE];
        check_block_header(head, path, at)?;
        let size = block_size(head);
        if size < HEADER_SIZE {
            return Err(Error::corrupt(
                path,
                at,
                format!(
                    "the bgzf block at {at} declares {size} bytes, which is less than \
                     its own header"
                ),
            ));
        }
        if index + size > raw.len() {
            return Err(Error::corrupt(
                path,
                at,
                format!(
                    "the bgzf block at {at} declares {size} bytes and only {} are left, \
                     so the file is truncated",
                    raw.len() - index
                ),
            ));
        }

        let block = inflate(&raw[index..index + size], path, at)?;
        index += size;

        // The chunk starts partway into its first block and ends partway into
        // its last; every block between is taken whole. Both can be the same
        // block, in which case both trims apply to it.
        let from = if at == first_block { begin_within } else { 0 };
        let to = if at == last_block {
            end_within.min(block.len())
        } else {
            block.len()
        };
        if to > from {
            out.extend_from_slice(&block[from..to]);
        }
    }
    Ok(out.freeze())
}

/// Inflate one BGZF block, which is a gzip member.
fn inflate(block: &[u8], path: &str, at: u64) -> Result<Vec<u8>> {
    let mut out = Vec::with_capacity(MAX_BLOCK_SIZE);
    flate2::read::GzDecoder::new(block)
        .read_to_end(&mut out)
        .map_err(|e| {
            Error::corrupt(
                path,
                at,
                format!("could not inflate the bgzf block at {at}: {e}"),
            )
        })?;
    Ok(out)
}

/// Check that the file ends with the EOF marker block.
///
/// Done before the header, so a truncated file is refused rather than opened
/// and read short.
pub fn check_eof(source: &dyn ByteSource) -> Result<()> {
    let path = source.path();
    let len = source.len()?;
    if len < EOF_SIZE as u64 {
        return Err(Error::format(
            path,
            "file is too short to be a bam (it does not hold even the bgzf end-of-file block)",
        ));
    }
    let tail = source.read_exact_at(len - EOF_SIZE as u64, EOF_SIZE)?;
    if tail[..] != BGZF_EOF_BLOCK[..] {
        return Err(Error::format(
            path,
            "bam file is truncated (it does not end with the bgzf end-of-file block)",
        ));
    }
    Ok(())
}

/// Sequential inflated bytes from the start of a BGZF file.
///
/// What the header reader walks: the header is at the front and is not addressed
/// by any virtual offset, so it is read as an ordinary multi-member gzip stream
/// rather than block by block.
pub fn header_reader(source: &dyn ByteSource) -> impl Read + '_ {
    flate2::read::MultiGzDecoder::new(SourceReader {
        source,
        offset: 0,
        buffer: Bytes::new(),
    })
}

/// A [`ByteSource`] as a sequential [`Read`], reading a block at a time.
struct SourceReader<'a> {
    source: &'a dyn ByteSource,
    offset: u64,
    buffer: Bytes,
}

impl Read for SourceReader<'_> {
    fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> {
        if self.buffer.is_empty() {
            self.buffer = self
                .source
                .read_at(self.offset, MAX_BLOCK_SIZE)
                .map_err(std::io::Error::other)?;
            if self.buffer.is_empty() {
                return Ok(0);
            }
            self.offset += self.buffer.len() as u64;
        }
        let take = out.len().min(self.buffer.len());
        out[..take].copy_from_slice(&self.buffer[..take]);
        self.buffer = self.buffer.slice(take..);
        Ok(take)
    }
}

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

    /// One BGZF block wrapping `payload`.
    fn block(payload: &[u8]) -> Vec<u8> {
        let mut encoder =
            flate2::write::DeflateEncoder::new(Vec::new(), flate2::Compression::new(6));
        encoder.write_all(payload).unwrap();
        let deflated = encoder.finish().unwrap();

        let total = HEADER_SIZE + deflated.len() + 8;
        let mut out = Vec::with_capacity(total);
        out.extend_from_slice(&[0x1f, 0x8b, 8, 4, 0, 0, 0, 0, 0, 0xff]);
        out.extend_from_slice(&6u16.to_le_bytes()); // XLEN
        out.extend_from_slice(b"BC");
        out.extend_from_slice(&2u16.to_le_bytes());
        out.extend_from_slice(&((total - 1) as u16).to_le_bytes()); // BSIZE
        out.extend_from_slice(&deflated);
        out.extend_from_slice(&crc32(payload).to_le_bytes());
        out.extend_from_slice(&(payload.len() as u32).to_le_bytes());
        assert_eq!(out.len(), total);
        out
    }

    fn crc32(data: &[u8]) -> u32 {
        let mut hasher = flate2::Crc::new();
        hasher.update(data);
        hasher.sum()
    }

    #[test]
    fn a_virtual_offset_splits_into_its_two_halves() {
        let offset = VirtualOffset::new(0x0001_2345_6789, 0xABCD);
        assert_eq!(offset.block_offset(), 0x0001_2345_6789);
        assert_eq!(offset.within_block(), 0xABCD);
        // The top 48 bits and the low 16, packed as the format packs them.
        assert_eq!(offset.0, 0x0001_2345_6789_ABCD);
    }

    #[test]
    fn a_chunk_inside_one_block_is_trimmed_at_both_ends() {
        let payload: Vec<u8> = (0..200u8).collect();
        let source = MemorySource::new(block(&payload));
        let chunk = Chunk {
            begin: VirtualOffset::new(0, 50),
            end: VirtualOffset::new(0, 120),
        };
        let got = decompress_chunk(&source, chunk, "test").unwrap();
        assert_eq!(&got[..], &payload[50..120]);
    }

    #[test]
    fn a_chunk_spanning_blocks_takes_the_middle_ones_whole() {
        let first: Vec<u8> = (0..100u8).collect();
        let second: Vec<u8> = (100..200u8).collect();
        let third: Vec<u8> = (200..250u8).collect();
        let mut bytes = block(&first);
        let second_at = bytes.len() as u64;
        bytes.extend_from_slice(&block(&second));
        let third_at = bytes.len() as u64;
        bytes.extend_from_slice(&block(&third));
        let source = MemorySource::new(bytes);

        let chunk = Chunk {
            begin: VirtualOffset::new(0, 90),
            end: VirtualOffset::new(third_at, 10),
        };
        let got = decompress_chunk(&source, chunk, "test").unwrap();
        let mut expected = first[90..].to_vec();
        expected.extend_from_slice(&second);
        expected.extend_from_slice(&third[..10]);
        assert_eq!(&got[..], &expected[..]);
        assert!(second_at > 0);
    }

    #[test]
    fn a_chunk_ending_on_a_block_boundary_does_not_need_that_block() {
        let first: Vec<u8> = (0..100u8).collect();
        let bytes = block(&first);
        let end_at = bytes.len() as u64;
        // The block at `end_at` is absent from the file entirely.
        let source = MemorySource::new(bytes);
        let chunk = Chunk {
            begin: VirtualOffset::new(0, 0),
            end: VirtualOffset::new(end_at, 0),
        };
        assert_eq!(
            &decompress_chunk(&source, chunk, "test").unwrap()[..],
            &first[..]
        );
    }

    #[test]
    fn something_that_is_not_a_bgzf_block_is_refused_by_name() {
        let source = MemorySource::new(vec![0u8; 4096]);
        let chunk = Chunk {
            begin: VirtualOffset::new(0, 0),
            end: VirtualOffset::new(0, 10),
        };
        let err = decompress_chunk(&source, chunk, "test")
            .unwrap_err()
            .to_string();
        assert!(err.contains("no bgzf block header"), "{err}");
    }

    #[test]
    fn a_truncated_block_is_corrupt_not_a_short_read() {
        let payload: Vec<u8> = (0..200u8).collect();
        let mut bytes = block(&payload);
        bytes.truncate(bytes.len() - 10);
        let source = MemorySource::new(bytes);
        let chunk = Chunk {
            begin: VirtualOffset::new(0, 0),
            end: VirtualOffset::new(0, 200),
        };
        let err = decompress_chunk(&source, chunk, "test")
            .unwrap_err()
            .to_string();
        assert!(err.contains("truncated"), "{err}");
    }

    #[test]
    fn the_eof_block_is_what_says_a_file_is_whole() {
        let mut bytes = block(b"hello");
        bytes.extend_from_slice(&BGZF_EOF_BLOCK);
        assert!(check_eof(&MemorySource::new(bytes.clone())).is_ok());

        bytes.truncate(bytes.len() - 1);
        let err = check_eof(&MemorySource::new(bytes))
            .unwrap_err()
            .to_string();
        assert!(err.contains("truncated"), "{err}");

        let err = check_eof(&MemorySource::new(vec![0u8; 4]))
            .unwrap_err()
            .to_string();
        assert!(err.contains("too short"), "{err}");
    }

    #[test]
    fn the_header_reader_walks_every_member() {
        let mut bytes = block(b"BAM\x01first ");
        bytes.extend_from_slice(&block(b"second"));
        bytes.extend_from_slice(&BGZF_EOF_BLOCK);
        let source = MemorySource::new(bytes);
        let mut text = Vec::new();
        header_reader(&source).read_to_end(&mut text).unwrap();
        assert_eq!(&text, b"BAM\x01first second");
    }

    #[test]
    fn a_chunk_compressed_size_is_its_block_span() {
        let chunk = Chunk {
            begin: VirtualOffset::new(1000, 5),
            end: VirtualOffset::new(9000, 7),
        };
        assert_eq!(chunk.compressed_size(), 8000);
    }
}