gwseq-io 0.2.0

//! The BAM header: the SAM text block and the binary reference list.

use std::io::Read;

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

#[derive(Debug, Clone)]
pub struct HeaderField {
    pub tag: String,
    pub value: String,
}

/// One `@`-prefixed line: its two-letter type and its fields.
#[derive(Debug, Clone)]
pub struct HeaderLine {
    pub kind: String,
    pub fields: Vec<HeaderField>,
}

#[derive(Debug, Clone, Default)]
pub struct SamHeader {
    pub lines: Vec<HeaderLine>,
}

impl SamHeader {
    /// Parse the text block: `@HD\tVN:1.6\tSO:coordinate`, one line each.
    ///
    /// A field without a colon is kept with an empty tag rather than dropped —
    /// the header is the file's, and a reader that silently discards part of it
    /// is worse than one that hands it over as it stands.
    pub fn parse(text: &str) -> Self {
        let lines = text
            .lines()
            .filter_map(|line| {
                let line = line.strip_prefix('@')?;
                let mut parts = line.split('\t');
                let kind = parts.next()?.to_string();
                let fields = parts
                    .filter(|f| !f.is_empty())
                    .map(|field| match field.split_once(':') {
                        Some((tag, value)) => HeaderField {
                            tag: tag.to_string(),
                            value: value.to_string(),
                        },
                        None => HeaderField {
                            tag: String::new(),
                            value: field.to_string(),
                        },
                    })
                    .collect();
                Some(HeaderLine { kind, fields })
            })
            .collect();
        Self { lines }
    }
}

/// Read the header and the reference list.
///
/// Reference names come from the **binary list**, not from the `@SQ` lines: the
/// two can disagree, and the binary list is what record `ref_id`s index into.
pub fn read(source: &dyn ByteSource) -> Result<(SamHeader, ChrMap)> {
    let path = source.path();
    let mut reader = super::bgzf::header_reader(source);

    let mut magic = [0u8; 4];
    read_exact(&mut reader, &mut magic, path)?;
    if &magic != b"BAM\x01" {
        return Err(Error::format(path, "invalid bam magic string"));
    }

    let l_text = read_u32(&mut reader, path)? as usize;
    let text = read_sized(&mut reader, l_text, MAX_TEXT, "header text", path)?;
    let header = SamHeader::parse(&String::from_utf8_lossy(&text));

    let n_ref = read_u32(&mut reader, path)? as usize;
    let mut entries = Vec::with_capacity(n_ref.min(1 << 16));
    for index in 0..n_ref {
        let l_name = read_u32(&mut reader, path)? as usize;
        if l_name == 0 {
            return Err(Error::corrupt(path, 0, "a bam reference has an empty name"));
        }
        let name = read_sized(&mut reader, l_name, MAX_NAME, "reference name", path)?;
        // NUL-terminated in the file; the terminator is not part of the name.
        let name = String::from_utf8_lossy(name.strip_suffix(&[0]).unwrap_or(&name)).into_owned();
        let size = read_u32(&mut reader, path)? as i64;
        entries.push((name, size, index));
    }
    Ok((header, ChrMap::from_indexed_entries(entries)))
}

/// A bam header longer than this is not a header. Real ones run to a few MB on
/// a scaffolded assembly with a long `@PG` chain; a quarter of a gibibyte is
/// past anything a tool writes and short of anything that hurts to allocate.
const MAX_TEXT: usize = 256 << 20;

/// A reference name longer than this is not a name. The SAM spec's own regex
/// puts no bound on it, but a `u32` field does, and a corrupt one asks for
/// 4 GiB.
const MAX_NAME: usize = 64 << 10;

/// Read `len` bytes, refusing an implausible `len` before anything is
/// allocated and growing into what the file really holds rather than reserving
/// the number it named.
///
/// The `take` is what makes the difference: `vec![0u8; len]` allocates and
/// zeroes `len` bytes before a single one has been read, and a failed
/// allocation is an abort no `Result` can carry.
fn read_sized(
    reader: &mut impl Read,
    len: usize,
    max: usize,
    what: &str,
    path: &str,
) -> Result<Vec<u8>> {
    if len > max {
        return Err(Error::corrupt(
            path,
            0,
            format!("bam {what} declares {len} bytes, more than the {max} this reader allows"),
        ));
    }
    let mut buf = Vec::new();
    reader
        .take(len as u64)
        .read_to_end(&mut buf)
        .map_err(|e| Error::corrupt(path, 0, format!("bam header ended early: {e}")))?;
    if buf.len() != len {
        return Err(Error::corrupt(
            path,
            0,
            format!(
                "bam {what} declares {len} bytes and the file holds {}",
                buf.len()
            ),
        ));
    }
    Ok(buf)
}

fn read_exact(reader: &mut impl Read, buf: &mut [u8], path: &str) -> Result<()> {
    reader
        .read_exact(buf)
        .map_err(|e| Error::corrupt(path, 0, format!("bam header ended early: {e}")))
}

fn read_u32(reader: &mut impl Read, path: &str) -> Result<u32> {
    let mut buf = [0u8; 4];
    read_exact(reader, &mut buf, path)?;
    Ok(u32::from_le_bytes(buf))
}

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

    #[test]
    fn header_lines_split_into_typed_fields() {
        let header = SamHeader::parse("@HD\tVN:1.6\tSO:coordinate\n@SQ\tSN:chr1\tLN:1000\n");
        assert_eq!(header.lines.len(), 2);
        assert_eq!(header.lines[0].kind, "HD");
        assert_eq!(header.lines[0].fields[1].tag, "SO");
        assert_eq!(header.lines[0].fields[1].value, "coordinate");
        assert_eq!(header.lines[1].fields[0].value, "chr1");
    }

    #[test]
    fn a_value_carrying_colons_keeps_all_but_the_first() {
        // A @PG CL: field is a whole command line, colons and all.
        let header = SamHeader::parse("@PG\tID:x\tCL:bwa mem -R @RG\\tID:1 ref.fa\n");
        assert_eq!(header.lines[0].fields[1].tag, "CL");
        assert!(header.lines[0].fields[1].value.contains("ID:1"));
    }

    #[test]
    fn a_field_with_no_colon_is_kept_rather_than_dropped() {
        let header = SamHeader::parse("@CO\tsome free text\n");
        assert_eq!(header.lines[0].kind, "CO");
        assert_eq!(header.lines[0].fields[0].tag, "");
        assert_eq!(header.lines[0].fields[0].value, "some free text");
    }

    #[test]
    fn lines_that_are_not_header_lines_are_skipped() {
        let header = SamHeader::parse("@HD\tVN:1.6\nnot a header line\n\n@SQ\tSN:chr1\n");
        assert_eq!(header.lines.len(), 2);
        assert_eq!(header.lines[1].kind, "SQ");
    }

    #[test]
    fn an_empty_header_parses_to_nothing() {
        assert!(SamHeader::parse("").lines.is_empty());
    }
}