gwseq_io/source/

gzip.rs

//! Transparent gzip, for the text converters.
//!
//! A codec is sequential, and sequential is [`std::io::Read`] — a different
//! trait from [`ByteSource`], not a crippled implementation of it. There is no
//! forward-only "file" here that has to refuse `len()` and override every
//! method that would move its cursor: nothing needing random access ever sees a
//! gzip stream, and nothing reading a gzip stream ever seeks.
//!
//! [`ByteSource`]: crate::source::ByteSource

use std::io::{BufReader, Read};
use std::path::Path;

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

/// `MultiGzDecoder`, never `GzDecoder`: concatenated members are what a gzip
/// file is here — BGZF is nothing else, and `bgzip` writes one member per
/// block — and only the multi variant walks past the first.
pub type GzipReader<R> = flate2::read::MultiGzDecoder<R>;

/// The 0x1F 0x8B magic.
pub fn is_gzipped(first_two: &[u8]) -> bool {
    first_two.len() >= 2 && first_two[0] == 0x1F && first_two[1] == 0x8B
}

/// Open a path for sequential reading, inflating when it is gzipped.
///
/// This is what the converters read their bedGraph / WIG / BED through, and the
/// only place in the crate that reads sequentially rather than by offset.
///
/// The magic decides, not the name: a `.bedgraph` holding gzip is read, and a
/// `.gz` holding text is too. Returns the reader and the file's size on disk,
/// which is what the converters report progress against — the *compressed*
/// size, since that is what "how far through the file are we" means to a caller
/// watching it.
pub fn open_text(path: impl AsRef<Path>) -> Result<(TextInput, u64)> {
    let path = path.as_ref();
    let display = path.to_string_lossy().into_owned();
    let file = std::fs::File::open(path).map_err(|e| Error::io(&display, e))?;
    let size = file.metadata().map_err(|e| Error::io(&display, e))?.len();

    let mut head = [0u8; 2];
    let read = read_full(&file, &mut head).map_err(|e| Error::io(&display, e))?;
    // Rewound rather than kept: the two bytes belong to whichever decoder gets
    // the file, and threading them past it would mean wrapping every reader in
    // a chain for the sake of one peek.
    use std::io::Seek as _;
    let mut file = file;
    file.rewind().map_err(|e| Error::io(&display, e))?;

    let inner: Box<dyn Read + Send> = if is_gzipped(&head[..read]) {
        Box::new(GzipReader::new(BufReader::with_capacity(1 << 16, file)))
    } else {
        Box::new(file)
    };
    Ok((BufReader::with_capacity(1 << 20, inner), size))
}

fn read_full(mut file: &std::fs::File, buf: &mut [u8]) -> std::io::Result<usize> {
    let mut filled = 0;
    while filled < buf.len() {
        match file.read(&mut buf[filled..]) {
            Ok(0) => break,
            Ok(n) => filled += n,
            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
            Err(e) => return Err(e),
        }
    }
    Ok(filled)
}

/// Sequential reads over anything, buffered. Named so the converters do not
/// spell the generic out.
pub type TextInput = BufReader<Box<dyn Read + Send>>;

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{BufRead as _, Write as _};

    fn temp(name: &str) -> std::path::PathBuf {
        let dir = std::env::temp_dir().join("gwseq_gzip_tests");
        std::fs::create_dir_all(&dir).unwrap();
        dir.join(name)
    }

    /// One gzip member per slice, concatenated — which is what a BGZF file is,
    /// and what `MultiGzDecoder` is here to walk. Built by the test rather than
    /// by a writer this crate exports: nothing in the library writes gzip, and
    /// a public type kept alive only by the tests that use it is not a public
    /// type worth having.
    fn gzip_members(bodies: &[&[u8]]) -> Vec<u8> {
        let mut out = Vec::new();
        for body in bodies {
            let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::new(6));
            e.write_all(body).unwrap();
            out.extend_from_slice(&e.finish().unwrap());
        }
        out
    }

    fn lines_of(path: &std::path::Path) -> Vec<String> {
        let (reader, _) = open_text(path).unwrap();
        reader.lines().map(|l| l.unwrap()).collect()
    }

    #[test]
    fn plain_text_is_read_as_it_stands() {
        let path = temp("plain.txt");
        std::fs::write(&path, "one\ntwo\nthree\n").unwrap();
        assert_eq!(lines_of(&path), ["one", "two", "three"]);
        std::fs::remove_file(&path).ok();
    }

    #[test]
    fn a_gzipped_file_is_inflated_whatever_it_is_called() {
        let path = temp("nogz.bedgraph");
        std::fs::write(&path, gzip_members(&[b"one\ntwo\nthree\n"])).unwrap();
        assert_eq!(lines_of(&path), ["one", "two", "three"]);
        std::fs::remove_file(&path).ok();
    }

    #[test]
    fn every_member_of_a_multi_member_file_is_read() {
        // What `MultiGzDecoder` buys: `GzDecoder` stops after the first, which
        // on a BGZF file means losing all but its first block.
        let path = temp("members.gz");
        let bodies: Vec<Vec<u8>> = (0..50).map(|i| format!("line{i}\n").into_bytes()).collect();
        let members = gzip_members(&bodies.iter().map(|b| &b[..]).collect::<Vec<_>>());
        std::fs::write(&path, members).unwrap();
        let lines = lines_of(&path);
        assert_eq!(lines.len(), 50);
        assert_eq!(lines[49], "line49");
        std::fs::remove_file(&path).ok();
    }

    #[test]
    fn the_reported_size_is_the_size_on_disk() {
        let path = temp("size.txt");
        std::fs::write(&path, vec![b'x'; 1234]).unwrap();
        let (_, size) = open_text(&path).unwrap();
        assert_eq!(size, 1234);
        std::fs::remove_file(&path).ok();
    }

    #[test]
    fn an_empty_file_reads_as_no_lines_rather_than_failing() {
        let path = temp("empty.txt");
        std::fs::write(&path, b"").unwrap();
        assert!(lines_of(&path).is_empty());
        std::fs::remove_file(&path).ok();
    }
}