fastx-io 0.3.0

//! Format and compression detection.

use std::fmt;
use std::path::Path;

/// A sequence file format.
///
/// Deliberately *not* `#[non_exhaustive]`: FASTA and FASTQ are the whole domain
/// of this crate, so matching on both arms exhaustively is meant to be pleasant
/// and is a promise we intend to keep.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Format {
    /// FASTA: `>id description` followed by one or more sequence lines.
    Fasta,
    /// FASTQ: `@id description`, sequence, `+`, quality.
    Fastq,
}

impl Format {
    /// Recognised FASTA extensions (without the leading dot).
    pub const FASTA_EXTENSIONS: &'static [&'static str] = &[
        "fa", "fasta", "fna", "faa", "ffn", "frn", "fas", "mpfa", "seq",
    ];

    /// Recognised FASTQ extensions (without the leading dot).
    pub const FASTQ_EXTENSIONS: &'static [&'static str] = &["fq", "fastq"];

    /// Infer the format from a file extension, ignoring a trailing `.gz`/`.bgz`/`.z`.
    ///
    /// ```
    /// use fastx::Format;
    /// assert_eq!(Format::from_path("reads.fq.gz"), Some(Format::Fastq));
    /// assert_eq!(Format::from_path("genome.fna"), Some(Format::Fasta));
    /// assert_eq!(Format::from_path("notes.txt"), None);
    /// ```
    pub fn from_path<P: AsRef<Path>>(path: P) -> Option<Format> {
        let path = path.as_ref();
        let ext = path.extension()?.to_str()?.to_ascii_lowercase();
        if matches!(ext.as_str(), "gz" | "bgz" | "gzip" | "z" | "zst" | "bz2") {
            let stem = path.file_stem()?;
            return Format::from_path(Path::new(stem));
        }
        Format::from_extension(&ext)
    }

    /// Infer the format from a bare extension such as `"fasta"`.
    pub fn from_extension(ext: &str) -> Option<Format> {
        let ext = ext.trim_start_matches('.').to_ascii_lowercase();
        if Format::FASTA_EXTENSIONS.contains(&ext.as_str()) {
            Some(Format::Fasta)
        } else if Format::FASTQ_EXTENSIONS.contains(&ext.as_str()) {
            Some(Format::Fastq)
        } else {
            None
        }
    }

    /// Infer the format from the first meaningful byte of a stream.
    ///
    /// ```
    /// use fastx::Format;
    /// assert_eq!(Format::from_first_byte(b'>'), Some(Format::Fasta));
    /// assert_eq!(Format::from_first_byte(b'@'), Some(Format::Fastq));
    /// ```
    pub fn from_first_byte(byte: u8) -> Option<Format> {
        match byte {
            b'>' | b';' => Some(Format::Fasta),
            b'@' => Some(Format::Fastq),
            _ => None,
        }
    }

    /// The byte that starts a record header in this format.
    pub const fn header_byte(self) -> u8 {
        match self {
            Format::Fasta => b'>',
            Format::Fastq => b'@',
        }
    }

    /// The canonical extension used when creating files.
    pub const fn extension(self) -> &'static str {
        match self {
            Format::Fasta => "fasta",
            Format::Fastq => "fastq",
        }
    }

    /// Whether records in this format carry per-base quality scores.
    pub const fn has_quality(self) -> bool {
        matches!(self, Format::Fastq)
    }
}

impl fmt::Display for Format {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Format::Fasta => f.write_str("FASTA"),
            Format::Fastq => f.write_str("FASTQ"),
        }
    }
}

/// Compression applied to a sequence file.
///
/// Marked `#[non_exhaustive]`: match with a `_` arm. New containers keep
/// appearing — this list has already grown twice — and each one should be a
/// minor release rather than a breaking one.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
#[non_exhaustive]
pub enum Compression {
    /// Plain, uncompressed bytes.
    #[default]
    None,
    /// One deflate stream, as `gzip` produces. Readable start to finish only.
    Gzip,
    /// Block-compressed gzip, as `bgzip` produces.
    ///
    /// Valid gzip that any tool can decompress, but split into independent
    /// members so that a reader with a `.gzi` index can seek into it. Costs a
    /// percent or two of compression ratio and is what the samtools ecosystem
    /// expects, so it is the default for compressed output.
    Bgzf,
    /// Zstandard (requires the `zstd` feature).
    ///
    /// Compresses faster and smaller than gzip, which makes it attractive for
    /// intermediate files, but the wider bioinformatics toolchain does not read
    /// it and a plain zstd frame cannot be randomly accessed — so a reference
    /// genome still wants [`Compression::Bgzf`].
    Zstd,
}

impl Compression {
    /// The first two bytes of a gzip member.
    pub const GZIP_MAGIC: [u8; 2] = [0x1f, 0x8b];

    /// The four-byte magic number of a Zstandard frame.
    pub const ZSTD_MAGIC: [u8; 4] = [0x28, 0xb5, 0x2f, 0xfd];

    /// Infer compression from a file extension.
    ///
    /// Note that `.gz` maps to [`Compression::Gzip`] rather than
    /// [`Compression::Bgzf`]: an extension cannot tell the two apart, and only
    /// the bytes can. Writers pick BGZF for `.gz` deliberately; readers sniff.
    pub fn from_path<P: AsRef<Path>>(path: P) -> Compression {
        match path
            .as_ref()
            .extension()
            .and_then(|e| e.to_str())
            .map(|e| e.to_ascii_lowercase())
            .as_deref()
        {
            Some("gz") | Some("bgz") | Some("gzip") => Compression::Gzip,
            Some("zst") | Some("zstd") => Compression::Zstd,
            _ => Compression::None,
        }
    }

    /// Infer compression from the leading bytes of a stream.
    ///
    /// Does not distinguish plain gzip from BGZF — both report
    /// [`Compression::Gzip`], since telling them apart means parsing the extra
    /// field. Use [`crate::bgzf::is_bgzf`] for that.
    pub fn from_magic(bytes: &[u8]) -> Compression {
        if bytes.len() >= 2 && bytes[..2] == Compression::GZIP_MAGIC {
            Compression::Gzip
        } else if bytes.len() >= 4 && bytes[..4] == Compression::ZSTD_MAGIC {
            Compression::Zstd
        } else {
            Compression::None
        }
    }
}

/// gzip compression level used when writing.
///
/// The level dominates the cost of writing compressed output — far more than
/// parsing does. On an 81 MiB FASTQ, one run took 2.8 s at level 1 and 18.6 s at
/// level 6, for output of 42.2 MB versus 38.9 MB: **6.5× the time to save 8%**.
/// [`CompressionLevel::FAST`] is almost always the right choice for files that
/// another pipeline stage will immediately read back.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CompressionLevel(pub u32);

impl CompressionLevel {
    /// No compression, fastest. Still a valid gzip stream.
    pub const NONE: CompressionLevel = CompressionLevel(0);
    /// Fast, slightly larger output — what pipeline intermediates want.
    pub const FAST: CompressionLevel = CompressionLevel(1);
    /// gzip's own default, and this crate's, for consistency with other tools.
    pub const DEFAULT: CompressionLevel = CompressionLevel(6);
    /// Smallest output, slowest — for data you write once and archive.
    pub const BEST: CompressionLevel = CompressionLevel(9);
}

impl Default for CompressionLevel {
    fn default() -> Self {
        CompressionLevel::DEFAULT
    }
}

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

    #[test]
    fn detects_format_from_extensions() {
        assert_eq!(Format::from_path("a.fa"), Some(Format::Fasta));
        assert_eq!(Format::from_path("a.FASTA"), Some(Format::Fasta));
        assert_eq!(Format::from_path("a.faa"), Some(Format::Fasta));
        assert_eq!(Format::from_path("/tmp/x/a.fna.gz"), Some(Format::Fasta));
        assert_eq!(Format::from_path("a.fq"), Some(Format::Fastq));
        assert_eq!(Format::from_path("a.fastq.gz"), Some(Format::Fastq));
        assert_eq!(Format::from_path("a.gz"), None);
        assert_eq!(Format::from_path("a"), None);
    }

    #[test]
    fn detects_compression() {
        assert_eq!(Compression::from_path("a.fq.gz"), Compression::Gzip);
        assert_eq!(Compression::from_path("a.fq"), Compression::None);
        assert_eq!(
            Compression::from_magic(&[0x1f, 0x8b, 0x08]),
            Compression::Gzip
        );
        assert_eq!(Compression::from_magic(b">seq"), Compression::None);
        assert_eq!(Compression::from_magic(&[0x1f]), Compression::None);
    }

    #[test]
    fn detects_zstd() {
        assert_eq!(Compression::from_path("a.fq.zst"), Compression::Zstd);
        assert_eq!(Compression::from_path("a.fa.zstd"), Compression::Zstd);
        // The format still resolves through a zstd suffix.
        assert_eq!(Format::from_path("reads.fq.zst"), Some(Format::Fastq));

        assert_eq!(
            Compression::from_magic(&[0x28, 0xb5, 0x2f, 0xfd, 0x00]),
            Compression::Zstd
        );
        // Three bytes of the magic are not enough to claim a match.
        assert_eq!(
            Compression::from_magic(&[0x28, 0xb5, 0x2f]),
            Compression::None
        );
    }
}