fastx-io 0.3.0

//! Paired-end reads, from two files or one interleaved stream.
//!
//! Paired data is where sequencing pipelines quietly go wrong: R1 and R2 must
//! stay in lockstep, and a tool that filters one file without the other, or
//! reads two files that were sorted differently, produces output that looks
//! perfectly well-formed and is silently mispaired. [`PairedReader`] checks that
//! the mate names line up on every pair, so that goes from a subtle wrong answer
//! to an error on the first record.
//!
//! ```
//! use fastx::paired::PairedReader;
//!
//! let r1 = b"@read1/1\nACGT\n+\nIIII\n@read2/1\nTTTT\n+\nIIII\n";
//! let r2 = b"@read1/2\nCCCC\n+\nIIII\n@read2/2\nGGGG\n+\nIIII\n";
//!
//! let mut reader = PairedReader::from_readers(&r1[..], &r2[..]);
//! let pair = reader.read_pair()?.unwrap();
//! assert_eq!(pair.first.id, "read1/1");
//! assert_eq!(pair.second.id, "read1/2");
//! assert_eq!(reader.count_pairs()?, 1); // one pair left
//! # Ok::<(), fastx::Error>(())
//! ```

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

use crate::error::{Error, Result};
use crate::format::Format;
use crate::reader::{self, FastxReader};
use crate::record::Sequence;
use crate::writer::FastxWriter;

/// Two mates of the same fragment.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Pair {
    /// The forward read, from R1.
    pub first: Sequence,
    /// The reverse read, from R2.
    pub second: Sequence,
}

impl Pair {
    /// A pair from two records.
    pub fn new(first: Sequence, second: Sequence) -> Pair {
        Pair { first, second }
    }

    /// Combined length of both mates.
    pub fn len(&self) -> usize {
        self.first.len() + self.second.len()
    }

    /// True when neither mate has any residues.
    pub fn is_empty(&self) -> bool {
        self.first.is_empty() && self.second.is_empty()
    }

    /// Reset both records, keeping their allocations.
    pub fn clear(&mut self) {
        self.first.clear();
        self.second.clear();
    }

    /// Whether the two ids name the same fragment.
    pub fn names_match(&self) -> bool {
        mate_stem(&self.first.id) == mate_stem(&self.second.id)
    }
}

/// The part of a read name shared by both mates.
///
/// Illumina writes the mate number either as a `/1` suffix on the name or as a
/// separate field after a space — and this crate's ids stop at the first
/// whitespace, so the second style already gives identical ids. Only the suffix
/// style needs stripping, in the few spellings that occur in the wild.
///
/// ```
/// # use fastx::paired::mate_stem;
/// assert_eq!(mate_stem("read1/1"), "read1");
/// assert_eq!(mate_stem("read1.2"), "read1");
/// assert_eq!(mate_stem("read1_1"), "read1");
/// assert_eq!(mate_stem("read1"), "read1");
/// // A name that merely ends in a digit is left alone.
/// assert_eq!(mate_stem("read12"), "read12");
/// ```
pub fn mate_stem(id: &str) -> &str {
    let bytes = id.as_bytes();
    if bytes.len() >= 2 {
        let last = bytes[bytes.len() - 1];
        let separator = bytes[bytes.len() - 2];
        if (last == b'1' || last == b'2') && matches!(separator, b'/' | b'.' | b'_') {
            return &id[..id.len() - 2];
        }
    }
    id
}

/// Where a [`PairedReader`] takes its records from.
// One variant holds two readers and the other holds one, so a size difference is
// inherent. Boxing would move a couple of hundred bytes to the heap and add an
// indirection to every record read, to save nothing: exactly one of these exists
// per `PairedReader`, never a collection of them.
#[allow(clippy::large_enum_variant)]
enum Source<R: Read> {
    /// Separate R1 and R2 files.
    Split {
        first: FastxReader<R>,
        second: FastxReader<R>,
    },
    /// One stream with the mates alternating.
    Interleaved(FastxReader<R>),
}

/// A reader that yields both mates at once.
///
/// By default it verifies that the mate names agree; call
/// [`PairedReader::check_names`] to turn that off for data with naming that does
/// not follow any convention.
pub struct PairedReader<R: Read> {
    source: Source<R>,
    check_names: bool,
    pairs_read: u64,
}

/// The type [`PairedReader::open`] and [`PairedReader::open_interleaved`] return.
pub type BoxedPairedReader = PairedReader<Box<dyn Read + Send>>;

impl PairedReader<Box<dyn Read + Send>> {
    /// Open an R1/R2 pair of files, each transparently decompressed.
    pub fn open<P: AsRef<Path>, Q: AsRef<Path>>(first: P, second: Q) -> Result<BoxedPairedReader> {
        Ok(PairedReader {
            source: Source::Split {
                first: reader::open(first)?,
                second: reader::open(second)?,
            },
            check_names: true,
            pairs_read: 0,
        })
    }

    /// Open one interleaved file, where R1 and R2 alternate.
    pub fn open_interleaved<P: AsRef<Path>>(path: P) -> Result<BoxedPairedReader> {
        Ok(PairedReader {
            source: Source::Interleaved(reader::open(path)?),
            check_names: true,
            pairs_read: 0,
        })
    }
}

impl<R: Read> PairedReader<R> {
    /// Read a pair from two separate streams.
    pub fn from_readers(first: R, second: R) -> PairedReader<R> {
        PairedReader::from_split(FastxReader::new(first), FastxReader::new(second))
    }

    /// Read pairs from one interleaved stream.
    pub fn interleaved(inner: R) -> PairedReader<R> {
        PairedReader::from_interleaved(FastxReader::new(inner))
    }

    /// Read pairs from a reader that is already configured — from
    /// [`crate::from_stdin`], say, or with a chosen buffer size.
    pub fn from_interleaved(reader: FastxReader<R>) -> PairedReader<R> {
        PairedReader {
            source: Source::Interleaved(reader),
            check_names: true,
            pairs_read: 0,
        }
    }

    /// Read pairs from two readers that are already configured.
    pub fn from_split(first: FastxReader<R>, second: FastxReader<R>) -> PairedReader<R> {
        PairedReader {
            source: Source::Split { first, second },
            check_names: true,
            pairs_read: 0,
        }
    }

    /// Whether to verify that mate names agree. On by default.
    pub fn check_names(mut self, check: bool) -> Self {
        self.check_names = check;
        self
    }

    /// Number of pairs read so far.
    pub fn pairs_read(&self) -> u64 {
        self.pairs_read
    }

    /// The format being read, once known.
    pub fn format(&self) -> Option<Format> {
        match &self.source {
            Source::Split { first, .. } => first.format(),
            Source::Interleaved(reader) => reader.format(),
        }
    }

    /// Read the next pair into `pair`, reusing its allocations.
    ///
    /// Returns `Ok(false)` when both inputs are exhausted together, and an error
    /// when only one of them is — a truncated mate file is a real problem, not
    /// an early stop.
    pub fn read_into(&mut self, pair: &mut Pair) -> Result<bool> {
        let present = match &mut self.source {
            Source::Split { first, second } => {
                let got_first = first.read_into(&mut pair.first)?;
                let got_second = second.read_into(&mut pair.second)?;
                match (got_first, got_second) {
                    (false, false) => return Ok(false),
                    (true, true) => true,
                    (true, false) => {
                        return Err(Error::PairTruncated {
                            pairs_read: self.pairs_read,
                            missing: "R2",
                        })
                    }
                    (false, true) => {
                        return Err(Error::PairTruncated {
                            pairs_read: self.pairs_read,
                            missing: "R1",
                        })
                    }
                }
            }
            Source::Interleaved(reader) => {
                if !reader.read_into(&mut pair.first)? {
                    return Ok(false);
                }
                if !reader.read_into(&mut pair.second)? {
                    return Err(Error::PairTruncated {
                        pairs_read: self.pairs_read,
                        missing: "the second mate",
                    });
                }
                true
            }
        };

        if self.check_names && !pair.names_match() {
            return Err(Error::PairMismatch {
                pairs_read: self.pairs_read,
                first: pair.first.id.clone(),
                second: pair.second.id.clone(),
            });
        }
        self.pairs_read += 1;
        Ok(present)
    }

    /// Read the next pair into a fresh [`Pair`].
    pub fn read_pair(&mut self) -> Result<Option<Pair>> {
        let mut pair = Pair::default();
        if self.read_into(&mut pair)? {
            Ok(Some(pair))
        } else {
            Ok(None)
        }
    }

    /// Run `f` on every remaining pair, reusing one buffer.
    pub fn for_each_pair<F>(&mut self, mut f: F) -> Result<()>
    where
        F: FnMut(&Pair) -> Result<()>,
    {
        let mut pair = Pair::default();
        while self.read_into(&mut pair)? {
            f(&pair)?;
        }
        Ok(())
    }

    /// Count the remaining pairs.
    pub fn count_pairs(&mut self) -> Result<u64> {
        let mut pairs = 0;
        let mut pair = Pair::default();
        while self.read_into(&mut pair)? {
            pairs += 1;
        }
        Ok(pairs)
    }
}

impl<R: Read> Iterator for PairedReader<R> {
    type Item = Result<Pair>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.read_pair() {
            Ok(Some(pair)) => Some(Ok(pair)),
            Ok(None) => None,
            Err(e) => Some(Err(e)),
        }
    }
}

/// A writer for paired reads, to two files or one interleaved stream.
pub enum PairedWriter<W: Write> {
    /// Separate R1 and R2 outputs.
    Split {
        /// Where forward reads go.
        first: FastxWriter<W>,
        /// Where reverse reads go.
        second: FastxWriter<W>,
    },
    /// One output with the mates alternating.
    Interleaved(FastxWriter<W>),
}

impl<W: Write> PairedWriter<W> {
    /// Write one pair.
    pub fn write_pair(&mut self, pair: &Pair) -> Result<()> {
        match self {
            PairedWriter::Split { first, second } => {
                first.write_record(&pair.first)?;
                second.write_record(&pair.second)?;
            }
            PairedWriter::Interleaved(writer) => {
                writer.write_record(&pair.first)?;
                writer.write_record(&pair.second)?;
            }
        }
        Ok(())
    }

    /// Flush every underlying writer.
    pub fn flush(&mut self) -> Result<()> {
        match self {
            PairedWriter::Split { first, second } => {
                first.flush()?;
                second.flush()?;
            }
            PairedWriter::Interleaved(writer) => writer.flush()?,
        }
        Ok(())
    }

    /// Flush and finish, surfacing any error from closing a compressed stream.
    pub fn finish(self) -> Result<()> {
        match self {
            PairedWriter::Split { first, second } => {
                first.finish()?;
                second.finish()?;
            }
            PairedWriter::Interleaved(writer) => {
                writer.finish()?;
            }
        }
        Ok(())
    }
}

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

    const R1: &[u8] = b"@read1/1\nACGT\n+\nIIII\n@read2/1\nTTTT\n+\nJJJJ\n";
    const R2: &[u8] = b"@read1/2\nCCCC\n+\nIIII\n@read2/2\nGGGG\n+\nJJJJ\n";

    #[test]
    fn reads_pairs_from_two_streams() {
        let mut reader = PairedReader::from_readers(R1, R2);
        let pairs: Vec<Pair> = reader.by_ref().collect::<Result<Vec<_>>>().unwrap();
        assert_eq!(pairs.len(), 2);
        assert_eq!(pairs[0].first.seq, b"ACGT");
        assert_eq!(pairs[0].second.seq, b"CCCC");
        assert_eq!(pairs[1].first.id, "read2/1");
        assert_eq!(reader.pairs_read(), 2);
    }

    #[test]
    fn reads_interleaved() {
        let interleaved = b"@r/1\nAC\n+\nII\n@r/2\nGT\n+\nII\n@s/1\nAA\n+\nII\n@s/2\nTT\n+\nII\n";
        let pairs: Vec<Pair> = PairedReader::interleaved(&interleaved[..])
            .collect::<Result<Vec<_>>>()
            .unwrap();
        assert_eq!(pairs.len(), 2);
        assert_eq!(pairs[0].second.seq, b"GT");
        assert_eq!(pairs[1].first.id, "s/1");
    }

    #[test]
    fn catches_mispaired_names() {
        // R2 sorted differently: the classic silent pipeline bug.
        let shuffled = b"@read2/2\nGGGG\n+\nJJJJ\n@read1/2\nCCCC\n+\nIIII\n";
        let error = PairedReader::from_readers(R1, &shuffled[..])
            .read_pair()
            .unwrap_err();
        match error {
            Error::PairMismatch { first, second, .. } => {
                assert_eq!((first.as_str(), second.as_str()), ("read1/1", "read2/2"));
            }
            other => panic!("expected PairMismatch, got {other}"),
        }

        // The check can be waived for data with unconventional names.
        let pairs = PairedReader::from_readers(R1, &shuffled[..])
            .check_names(false)
            .count_pairs()
            .unwrap();
        assert_eq!(pairs, 2);
    }

    #[test]
    fn catches_a_truncated_mate_file() {
        let short = b"@read1/2\nCCCC\n+\nIIII\n";
        let error = PairedReader::from_readers(R1, &short[..])
            .count_pairs()
            .unwrap_err();
        assert!(
            matches!(
                error,
                Error::PairTruncated {
                    missing: "R2",
                    pairs_read: 1
                }
            ),
            "{error}"
        );

        let error = PairedReader::from_readers(&short[..], R2)
            .count_pairs()
            .unwrap_err();
        assert!(
            matches!(error, Error::PairTruncated { missing: "R1", .. }),
            "{error}"
        );

        // An odd number of interleaved records is the same problem.
        let odd = b"@r/1\nAC\n+\nII\n@r/2\nGT\n+\nII\n@s/1\nAA\n+\nII\n";
        let error = PairedReader::interleaved(&odd[..])
            .count_pairs()
            .unwrap_err();
        assert!(
            matches!(error, Error::PairTruncated { pairs_read: 1, .. }),
            "{error}"
        );
    }

    #[test]
    fn accepts_illumina_style_names() {
        // Modern Illumina puts the mate number after a space, so the ids match
        // outright once the header is split.
        let r1 = b"@A00123:1:HXX:1:1101:1000:1000 1:N:0:ATCG\nAC\n+\nII\n";
        let r2 = b"@A00123:1:HXX:1:1101:1000:1000 2:N:0:ATCG\nGT\n+\nII\n";
        let pair = PairedReader::from_readers(&r1[..], &r2[..])
            .read_pair()
            .unwrap()
            .unwrap();
        assert!(pair.names_match());
        assert_eq!(pair.first.description.as_deref(), Some("1:N:0:ATCG"));
    }

    #[test]
    fn mate_stems() {
        for (id, stem) in [
            ("read/1", "read"),
            ("read/2", "read"),
            ("read.1", "read"),
            ("read_2", "read"),
            ("read", "read"),
            ("read1", "read1"),
            ("r", "r"),
            ("", ""),
            ("/1", ""),
            ("read/3", "read/3"),
        ] {
            assert_eq!(mate_stem(id), stem, "{id}");
        }
    }

    #[test]
    fn empty_input_yields_no_pairs() {
        assert_eq!(
            PairedReader::from_readers(&b""[..], &b""[..])
                .count_pairs()
                .unwrap(),
            0
        );
        assert_eq!(
            PairedReader::interleaved(&b""[..]).count_pairs().unwrap(),
            0
        );
    }

    #[test]
    fn round_trips_through_the_writer() {
        let mut interleaved = Vec::new();
        {
            let mut writer =
                PairedWriter::Interleaved(FastxWriter::new(&mut interleaved, Format::Fastq));
            let mut reader = PairedReader::from_readers(R1, R2);
            reader
                .for_each_pair(|pair| writer.write_pair(pair))
                .unwrap();
            writer.finish().unwrap();
        }

        let pairs: Vec<Pair> = PairedReader::interleaved(&interleaved[..])
            .collect::<Result<Vec<_>>>()
            .unwrap();
        let original: Vec<Pair> = PairedReader::from_readers(R1, R2)
            .collect::<Result<Vec<_>>>()
            .unwrap();
        assert_eq!(pairs, original);
    }

    #[test]
    fn split_writer_separates_the_mates() {
        let mut first = Vec::new();
        let mut second = Vec::new();
        {
            let mut writer = PairedWriter::Split {
                first: FastxWriter::new(&mut first, Format::Fastq),
                second: FastxWriter::new(&mut second, Format::Fastq),
            };
            PairedReader::from_readers(R1, R2)
                .for_each_pair(|pair| writer.write_pair(pair))
                .unwrap();
            writer.finish().unwrap();
        }
        assert_eq!(first, R1);
        assert_eq!(second, R2);
    }
}