fastx-io 0.3.0

//! Sequence utilities that operate on raw `&[u8]` slices.
//!
//! Everything here is case-preserving and allocation-free unless the return type
//! says otherwise, so these helpers can be used inside hot loops.

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

/// IUPAC pairs shared by the DNA and RNA complement tables.
const AMBIGUITY_PAIRS: &[(u8, u8)] = &[
    (b'Y', b'R'),
    (b'R', b'Y'),
    (b'S', b'S'),
    (b'W', b'W'),
    (b'K', b'M'),
    (b'M', b'K'),
    (b'B', b'V'),
    (b'V', b'B'),
    (b'D', b'H'),
    (b'H', b'D'),
    (b'N', b'N'),
];

/// Complement table for DNA output: `A` becomes `T`, and `U` becomes `A`.
static COMPLEMENT: [u8; 256] = build_complement_table(&[
    (b'A', b'T'),
    (b'T', b'A'),
    (b'U', b'A'),
    (b'G', b'C'),
    (b'C', b'G'),
]);

/// Complement table for RNA output: `A` becomes `U`, and both `T` and `U`
/// become `A`.
static COMPLEMENT_RNA: [u8; 256] = build_complement_table(&[
    (b'A', b'U'),
    (b'U', b'A'),
    (b'T', b'A'),
    (b'G', b'C'),
    (b'C', b'G'),
]);

const fn build_complement_table(bases: &[(u8, u8)]) -> [u8; 256] {
    // Bytes that are not nucleotide codes map to themselves.
    let mut table = [0u8; 256];
    let mut i = 0;
    while i < 256 {
        table[i] = i as u8;
        i += 1;
    }
    let mut p = 0;
    while p < bases.len() {
        let (from, to) = bases[p];
        table[from as usize] = to;
        table[(from + 32) as usize] = to + 32; // lowercase
        p += 1;
    }
    let mut p = 0;
    while p < AMBIGUITY_PAIRS.len() {
        let (from, to) = AMBIGUITY_PAIRS[p];
        table[from as usize] = to;
        table[(from + 32) as usize] = to + 32;
        p += 1;
    }
    table
}

/// Complement a single nucleotide as DNA, preserving case.
///
/// Bytes that are not nucleotide codes pass through unchanged, so gaps and
/// padding survive.
///
/// `U` complements to `A`, and `A` complements to `T` — so complementing twice
/// turns RNA into DNA rather than returning the original byte. That is the only
/// sensible reading of a DNA complement, but it does mean the operation is not
/// an involution on RNA. Use [`complement_rna`] to keep uracil.
///
/// ```
/// # use fastx::seq::complement;
/// assert_eq!(complement(b'A'), b'T');
/// assert_eq!(complement(b'g'), b'c');
/// assert_eq!(complement(b'-'), b'-');
/// assert_eq!(complement(b'U'), b'A');
/// assert_eq!(complement(complement(b'U')), b'T'); // not 'U'
/// ```
#[inline]
pub fn complement(base: u8) -> u8 {
    COMPLEMENT[base as usize]
}

/// Complement a single nucleotide as RNA: `A` becomes `U`, not `T`.
///
/// ```
/// # use fastx::seq::complement_rna;
/// assert_eq!(complement_rna(b'A'), b'U');
/// assert_eq!(complement_rna(b'U'), b'A');
/// assert_eq!(complement_rna(b'T'), b'A');
/// assert_eq!(complement_rna(complement_rna(b'u')), b'u');
/// ```
#[inline]
pub fn complement_rna(base: u8) -> u8 {
    COMPLEMENT_RNA[base as usize]
}

/// Reverse complement of a nucleotide sequence, as DNA.
///
/// ```
/// # use fastx::seq::reverse_complement;
/// assert_eq!(reverse_complement(b"ACGTn"), b"nACGT");
/// assert_eq!(reverse_complement(b"ACGU"), b"ACGT"); // uracil becomes thymine
/// ```
pub fn reverse_complement(seq: &[u8]) -> Vec<u8> {
    seq.iter().rev().map(|&b| complement(b)).collect()
}

/// Reverse complement of a nucleotide sequence, as RNA.
///
/// ```
/// # use fastx::seq::reverse_complement_rna;
/// assert_eq!(reverse_complement_rna(b"ACGU"), b"ACGU");
/// assert_eq!(reverse_complement_rna(b"AAGG"), b"CCUU");
/// ```
pub fn reverse_complement_rna(seq: &[u8]) -> Vec<u8> {
    seq.iter().rev().map(|&b| complement_rna(b)).collect()
}

/// Reverse complement into an existing buffer, which is cleared first.
///
/// Use this in loops to avoid one allocation per record.
pub fn reverse_complement_into(seq: &[u8], out: &mut Vec<u8>) {
    out.clear();
    out.reserve(seq.len());
    out.extend(seq.iter().rev().map(|&b| complement(b)));
}

/// Reverse complement a sequence in place.
pub fn reverse_complement_in_place(seq: &mut [u8]) {
    let n = seq.len();
    for i in 0..n / 2 {
        let j = n - 1 - i;
        let a = complement(seq[i]);
        let b = complement(seq[j]);
        seq[i] = b;
        seq[j] = a;
    }
    if n % 2 == 1 {
        seq[n / 2] = complement(seq[n / 2]);
    }
}

/// Per-base counts of a nucleotide sequence, case-insensitive.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct BaseCounts {
    /// Adenine.
    pub a: u64,
    /// Cytosine.
    pub c: u64,
    /// Guanine.
    pub g: u64,
    /// Thymine, or uracil in RNA.
    pub t: u64,
    /// `N` and other IUPAC ambiguity codes.
    pub n: u64,
    /// Anything that is not a nucleotide code (gaps, `*`, ...).
    pub other: u64,
}

/// Class assigned to each byte by [`BASE_CLASS`].
const CLASS_A: u8 = 0;
const CLASS_C: u8 = 1;
const CLASS_G: u8 = 2;
const CLASS_T: u8 = 3;
const CLASS_AMBIGUOUS: u8 = 4;
const CLASS_OTHER: u8 = 5;

/// Maps every byte to a base class, so that counting is one table lookup and one
/// increment per base rather than a chain of comparisons.
static BASE_CLASS: [u8; 256] = build_base_class_table();

const fn build_base_class_table() -> [u8; 256] {
    let mut table = [CLASS_OTHER; 256];
    table[b'A' as usize] = CLASS_A;
    table[b'a' as usize] = CLASS_A;
    table[b'C' as usize] = CLASS_C;
    table[b'c' as usize] = CLASS_C;
    table[b'G' as usize] = CLASS_G;
    table[b'g' as usize] = CLASS_G;
    table[b'T' as usize] = CLASS_T;
    table[b't' as usize] = CLASS_T;
    table[b'U' as usize] = CLASS_T;
    table[b'u' as usize] = CLASS_T;

    let ambiguous = b"NRYSWKMBDHV";
    let mut i = 0;
    while i < ambiguous.len() {
        table[ambiguous[i] as usize] = CLASS_AMBIGUOUS;
        table[(ambiguous[i] + 32) as usize] = CLASS_AMBIGUOUS;
        i += 1;
    }
    table
}

impl BaseCounts {
    /// Count the bases of `seq`.
    pub fn of(seq: &[u8]) -> BaseCounts {
        // Eight slots rather than six so that the index is provably in range and
        // the bounds check disappears.
        let mut counts = [0u64; 8];
        for &b in seq {
            counts[(BASE_CLASS[b as usize] & 7) as usize] += 1;
        }
        BaseCounts {
            a: counts[CLASS_A as usize],
            c: counts[CLASS_C as usize],
            g: counts[CLASS_G as usize],
            t: counts[CLASS_T as usize],
            n: counts[CLASS_AMBIGUOUS as usize],
            other: counts[CLASS_OTHER as usize],
        }
    }

    /// Total number of counted bytes.
    pub fn total(&self) -> u64 {
        self.a + self.c + self.g + self.t + self.n + self.other
    }

    /// Unambiguous A/C/G/T bases.
    pub fn acgt(&self) -> u64 {
        self.a + self.c + self.g + self.t
    }

    /// GC fraction over unambiguous bases, or `None` when there are none.
    pub fn gc_content(&self) -> Option<f64> {
        let acgt = self.acgt();
        if acgt == 0 {
            None
        } else {
            Some((self.g + self.c) as f64 / acgt as f64)
        }
    }

    /// Add another set of counts.
    pub fn merge(&mut self, other: &BaseCounts) {
        self.a += other.a;
        self.c += other.c;
        self.g += other.g;
        self.t += other.t;
        self.n += other.n;
        self.other += other.other;
    }
}

/// GC fraction of a sequence, ignoring ambiguity codes.
///
/// Returns `None` when the sequence has no unambiguous bases.
///
/// ```
/// # use fastx::seq::gc_content;
/// assert_eq!(gc_content(b"GGCC"), Some(1.0));
/// assert_eq!(gc_content(b"ATAT"), Some(0.0));
/// assert_eq!(gc_content(b"GCATNNNN"), Some(0.5));
/// assert_eq!(gc_content(b"NNNN"), None);
/// ```
pub fn gc_content(seq: &[u8]) -> Option<f64> {
    BaseCounts::of(seq).gc_content()
}

/// A residue alphabet used for validation.
///
/// Marked `#[non_exhaustive]`: alphabets are a judgement call, not a closed set,
/// so adding one should not break callers. Alphabets are normally passed in
/// rather than matched on, which makes the cost of the annotation close to zero.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum Alphabet {
    /// `ACGT` plus `N`, either case.
    Dna,
    /// `ACGU` plus `N`, either case.
    Rna,
    /// The full IUPAC nucleotide set plus `-` and `.` gaps.
    Iupac,
    /// The 20 amino acids plus `BXZJUO*-.`, either case.
    Protein,
    /// Any printable, non-whitespace byte.
    Any,
}

impl Alphabet {
    /// Whether `byte` is a member of the alphabet.
    pub fn contains(self, byte: u8) -> bool {
        let upper = byte.to_ascii_uppercase();
        match self {
            Alphabet::Dna => matches!(upper, b'A' | b'C' | b'G' | b'T' | b'N'),
            Alphabet::Rna => matches!(upper, b'A' | b'C' | b'G' | b'U' | b'N'),
            Alphabet::Iupac => {
                matches!(
                    upper,
                    b'A' | b'C'
                        | b'G'
                        | b'T'
                        | b'U'
                        | b'R'
                        | b'Y'
                        | b'S'
                        | b'W'
                        | b'K'
                        | b'M'
                        | b'B'
                        | b'D'
                        | b'H'
                        | b'V'
                        | b'N'
                        | b'-'
                        | b'.'
                )
            }
            Alphabet::Protein => upper.is_ascii_uppercase() || matches!(upper, b'*' | b'-' | b'.'),
            Alphabet::Any => byte.is_ascii_graphic(),
        }
    }

    /// Check every byte of `seq`, reporting the first violation.
    pub fn validate(self, seq: &[u8]) -> Result<()> {
        self.validate_named(seq, "")
    }

    /// Like [`Alphabet::validate`] but attaches a record id to the error.
    pub fn validate_named(self, seq: &[u8], id: &str) -> Result<()> {
        match seq.iter().position(|&b| !self.contains(b)) {
            None => Ok(()),
            Some(pos) => Err(Error::InvalidByte {
                id: id.to_string(),
                pos,
                byte: seq[pos],
            }),
        }
    }
}

/// Uppercase a sequence in place (soft-masked genomes use lowercase for repeats).
pub fn make_uppercase(seq: &mut [u8]) {
    seq.make_ascii_uppercase();
}

/// Replace `U`/`u` with `T`/`t`, turning RNA into DNA in place.
pub fn rna_to_dna(seq: &mut [u8]) {
    for b in seq.iter_mut() {
        match *b {
            b'U' => *b = b'T',
            b'u' => *b = b't',
            _ => {}
        }
    }
}

/// The standard genetic code (NCBI translation table 1), in `TCAG` codon order.
const CODON_TABLE: &[u8; 64] = b"FFLLSSSSYY**CC*WLLLLPPPPHHQQRRRRIIIMTTTTNNKKSSRRVVVVAAAADDEEGGGG";

/// Sentinel stored in [`CODON_INDEX`] for bytes that are not `T`/`U`, `C`, `A` or `G`.
const NOT_A_BASE: u8 = u8::MAX;

/// Maps every byte to its position in the `TCAG` codon ordering.
static CODON_INDEX: [u8; 256] = build_codon_index_table();

const fn build_codon_index_table() -> [u8; 256] {
    let mut table = [NOT_A_BASE; 256];
    table[b'T' as usize] = 0;
    table[b't' as usize] = 0;
    table[b'U' as usize] = 0;
    table[b'u' as usize] = 0;
    table[b'C' as usize] = 1;
    table[b'c' as usize] = 1;
    table[b'A' as usize] = 2;
    table[b'a' as usize] = 2;
    table[b'G' as usize] = 3;
    table[b'g' as usize] = 3;
    table
}

#[inline]
fn base_index(base: u8) -> Option<usize> {
    match CODON_INDEX[base as usize] {
        NOT_A_BASE => None,
        index => Some(index as usize),
    }
}

/// Translate a single codon using the standard genetic code.
///
/// Stop codons become `*` and codons containing an ambiguous base become `X`.
///
/// ```
/// # use fastx::seq::translate_codon;
/// assert_eq!(translate_codon(b"ATG"), b'M');
/// assert_eq!(translate_codon(b"tga"), b'*');
/// assert_eq!(translate_codon(b"ANG"), b'X');
/// ```
pub fn translate_codon(codon: &[u8]) -> u8 {
    if codon.len() < 3 {
        return b'X';
    }
    match (
        base_index(codon[0]),
        base_index(codon[1]),
        base_index(codon[2]),
    ) {
        (Some(a), Some(b), Some(c)) => CODON_TABLE[a * 16 + b * 4 + c],
        _ => b'X',
    }
}

/// Translate a nucleotide sequence into a protein sequence.
///
/// `frame` is 0, 1 or 2 and skips that many leading bases. A trailing partial
/// codon is dropped. When `stop_at_stop` is true, translation ends before the
/// first stop codon.
///
/// ```
/// # use fastx::seq::translate;
/// assert_eq!(translate(b"ATGGCCTGA", 0, false), b"MA*");
/// assert_eq!(translate(b"ATGGCCTGA", 0, true), b"MA");
/// assert_eq!(translate(b"AATGGCC", 1, false), b"MA");
/// ```
pub fn translate(seq: &[u8], frame: usize, stop_at_stop: bool) -> Vec<u8> {
    let seq = if frame < seq.len() {
        &seq[frame..]
    } else {
        &[][..]
    };
    let mut out = Vec::with_capacity(seq.len() / 3);
    for codon in seq.chunks_exact(3) {
        let aa = translate_codon(codon);
        if stop_at_stop && aa == b'*' {
            break;
        }
        out.push(aa);
    }
    out
}

/// Iterator over the overlapping k-mers of a sequence.
///
/// ```
/// # use fastx::seq::kmers;
/// let all: Vec<&[u8]> = kmers(b"ACGTA", 3).collect();
/// assert_eq!(all, vec![&b"ACG"[..], &b"CGT"[..], &b"GTA"[..]]);
/// assert_eq!(kmers(b"AC", 3).count(), 0);
/// ```
pub fn kmers(seq: &[u8], k: usize) -> impl Iterator<Item = &[u8]> {
    let n = if k == 0 || seq.len() < k {
        0
    } else {
        seq.len() - k + 1
    };
    (0..n).map(move |i| &seq[i..i + k])
}

/// Canonical k-mer: the lexicographically smaller of a k-mer and its reverse
/// complement, so that both strands hash to the same value.
pub fn canonical_kmer(kmer: &[u8]) -> Vec<u8> {
    let rc = reverse_complement(kmer);
    if rc.as_slice() < kmer {
        rc
    } else {
        kmer.to_vec()
    }
}

/// Hamming distance between two equal-length sequences, case-insensitive.
///
/// Returns `None` if the lengths differ.
pub fn hamming_distance(a: &[u8], b: &[u8]) -> Option<usize> {
    if a.len() != b.len() {
        return None;
    }
    Some(
        a.iter()
            .zip(b)
            .filter(|(x, y)| !x.eq_ignore_ascii_case(y))
            .count(),
    )
}

/// N50 of a set of lengths: the length `L` such that contigs of at least `L`
/// cover half of the total assembly length.
///
/// ```
/// # use fastx::seq::n50;
/// assert_eq!(n50(&mut [2, 3, 4, 5, 6]), Some(5));
/// assert_eq!(n50(&mut []), None);
/// ```
pub fn n50(lengths: &mut [u64]) -> Option<u64> {
    nx(lengths, 0.5)
}

/// Generalised N-statistic: `nx(lengths, 0.9)` is the N90.
///
/// `fraction` is clamped to `0.0..=1.0`. The slice is sorted as a side effect.
pub fn nx(lengths: &mut [u64], fraction: f64) -> Option<u64> {
    if lengths.is_empty() {
        return None;
    }
    let total: u64 = lengths.iter().sum();
    if total == 0 {
        return Some(0);
    }
    lengths.sort_unstable_by(|a, b| b.cmp(a));
    let target = total as f64 * fraction.clamp(0.0, 1.0);
    let mut acc = 0u64;
    for &len in lengths.iter() {
        acc += len;
        if acc as f64 >= target {
            return Some(len);
        }
    }
    lengths.last().copied()
}

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

    #[test]
    fn complements_iupac_and_preserves_case() {
        assert_eq!(reverse_complement(b"ACGTacgt"), b"acgtACGT");
        assert_eq!(reverse_complement(b"RYKM"), b"KMRY");
        assert_eq!(reverse_complement(b""), b"");
        // Non-nucleotide bytes are preserved.
        assert_eq!(reverse_complement(b"AC-GT"), b"AC-GT");
    }

    #[test]
    fn in_place_matches_allocating() {
        for s in [&b""[..], b"A", b"AC", b"ACG", b"ACGTN", b"acgtRYn"] {
            let mut owned = s.to_vec();
            reverse_complement_in_place(&mut owned);
            assert_eq!(
                owned,
                reverse_complement(s),
                "{:?}",
                String::from_utf8_lossy(s)
            );
        }
    }

    #[test]
    fn counts_bases() {
        let c = BaseCounts::of(b"AACCGGTTNNxx");
        assert_eq!((c.a, c.c, c.g, c.t, c.n, c.other), (2, 2, 2, 2, 2, 2));
        assert_eq!(c.total(), 12);
        assert_eq!(c.acgt(), 8);
        assert_eq!(c.gc_content(), Some(0.5));
    }

    #[test]
    fn validates_alphabets() {
        assert!(Alphabet::Dna.validate(b"ACGTN").is_ok());
        assert!(Alphabet::Dna.validate(b"ACGU").is_err());
        assert!(Alphabet::Rna.validate(b"ACGU").is_ok());
        assert!(Alphabet::Iupac.validate(b"ACGTRYKM-").is_ok());
        assert!(Alphabet::Protein.validate(b"MEEPQSDPSV*").is_ok());
        assert!(Alphabet::Any.validate(b"anything!").is_ok());
        assert!(Alphabet::Any.validate(b"tab\there").is_err());

        match Alphabet::Dna.validate_named(b"ACG!T", "read1") {
            Err(Error::InvalidByte { id, pos, byte }) => {
                assert_eq!((id.as_str(), pos, byte), ("read1", 3, b'!'));
            }
            other => panic!("expected InvalidByte, got {other:?}"),
        }
    }

    #[test]
    fn translates_all_codons() {
        // Spot-check every position of the table via the classic ordering.
        assert_eq!(translate(b"TTTTTCTTATTG", 0, false), b"FFLL");
        assert_eq!(translate(b"ATGCATTAA", 0, false), b"MH*");
        assert_eq!(translate(b"AUGCAU", 0, false), b"MH"); // RNA input
                                                           // Trailing partial codon dropped.
        assert_eq!(translate(b"ATGCA", 0, false), b"M");
        assert_eq!(translate(b"AT", 0, false), b"");
        assert_eq!(translate(b"ATG", 5, false), b"");
    }

    #[test]
    fn kmers_and_canonical() {
        assert_eq!(kmers(b"AAAA", 4).count(), 1);
        assert_eq!(kmers(b"AAAA", 0).count(), 0);
        assert_eq!(canonical_kmer(b"TTT"), b"AAA");
        assert_eq!(canonical_kmer(b"AAA"), b"AAA");
    }

    #[test]
    fn hamming() {
        assert_eq!(hamming_distance(b"ACGT", b"acgt"), Some(0));
        assert_eq!(hamming_distance(b"ACGT", b"ACGA"), Some(1));
        assert_eq!(hamming_distance(b"ACGT", b"ACG"), None);
    }

    #[test]
    fn n_statistics() {
        // Total 100; sorted desc 50,30,15,5 -> 50 reaches 50.
        assert_eq!(n50(&mut [5, 50, 30, 15]), Some(50));
        assert_eq!(nx(&mut [5, 50, 30, 15], 0.9), Some(15));
        assert_eq!(nx(&mut [0, 0], 0.5), Some(0));
    }

    #[test]
    fn rna_dna_conversion() {
        let mut s = b"ACGUacgu".to_vec();
        rna_to_dna(&mut s);
        assert_eq!(s, b"ACGTacgt");
        make_uppercase(&mut s);
        assert_eq!(s, b"ACGTACGT");
    }
}