fastx/

record.rs

//! The [`Sequence`] record type.

use std::fmt;
use std::io::Write;
use std::ops::Range;

use crate::error::{Error, Result};
use crate::format::Format;
use crate::qual::{self, PHRED33};
use crate::seq::{self, Alphabet, BaseCounts};

/// One FASTA or FASTQ record.
///
/// A record with `quality == None` is a FASTA record; a record with quality is a
/// FASTQ record whose quality string always has the same length as `seq`.
///
/// Fields are public so that pipelines can build records without ceremony, but
/// the constructors and [`Sequence::validate`] exist to keep the length
/// invariant intact.
///
/// ```
/// use fastx::Sequence;
///
/// let read = Sequence::fastq("read1", b"ACGT", b"IIII")?;
/// assert_eq!(read.len(), 4);
/// assert_eq!(read.mean_quality(), Some(40.0));
/// # Ok::<(), fastx::Error>(())
/// ```
#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
pub struct Sequence {
    /// Identifier: the header up to the first whitespace, without `>`/`@`.
    pub id: String,
    /// Everything after the first whitespace in the header, if any.
    pub description: Option<String>,
    /// Residues as ASCII bytes, without line breaks.
    pub seq: Vec<u8>,
    /// Phred quality characters (not scores), FASTQ only.
    pub quality: Option<Vec<u8>>,
}

impl Sequence {
    /// A FASTA record.
    pub fn fasta<I: Into<String>>(id: I, seq: impl Into<Vec<u8>>) -> Sequence {
        Sequence {
            id: id.into(),
            description: None,
            seq: seq.into(),
            quality: None,
        }
    }

    /// A FASTQ record. Fails if sequence and quality lengths differ.
    pub fn fastq<I: Into<String>>(
        id: I,
        seq: impl Into<Vec<u8>>,
        quality: impl Into<Vec<u8>>,
    ) -> Result<Sequence> {
        let id = id.into();
        let seq = seq.into();
        let quality = quality.into();
        if seq.len() != quality.len() {
            return Err(Error::LengthMismatch {
                id,
                seq: seq.len(),
                quality: quality.len(),
            });
        }
        Ok(Sequence {
            id,
            description: None,
            seq,
            quality: Some(quality),
        })
    }

    /// Builder-style setter for the description.
    pub fn with_description<D: Into<String>>(mut self, description: D) -> Sequence {
        let description = description.into();
        self.description = if description.is_empty() {
            None
        } else {
            Some(description)
        };
        self
    }

    /// Number of residues.
    pub fn len(&self) -> usize {
        self.seq.len()
    }

    /// True when the record has no residues.
    pub fn is_empty(&self) -> bool {
        self.seq.is_empty()
    }

    /// True when the record carries quality scores.
    pub fn has_quality(&self) -> bool {
        self.quality.is_some()
    }

    /// The format this record can be written as losslessly.
    pub fn format(&self) -> Format {
        if self.has_quality() {
            Format::Fastq
        } else {
            Format::Fasta
        }
    }

    /// The header line without its leading `>`/`@`.
    ///
    /// ```
    /// # use fastx::Sequence;
    /// let s = Sequence::fasta("chr1", b"ACGT".to_vec()).with_description("human chromosome 1");
    /// assert_eq!(s.header(), "chr1 human chromosome 1");
    /// ```
    pub fn header(&self) -> String {
        match &self.description {
            Some(d) => format!("{} {}", self.id, d),
            None => self.id.clone(),
        }
    }

    /// The residues as `&str`, if they are valid UTF-8 (ASCII in practice).
    pub fn seq_str(&self) -> Result<&str> {
        std::str::from_utf8(&self.seq).map_err(|e| Error::InvalidByte {
            id: self.id.clone(),
            pos: e.valid_up_to(),
            byte: self.seq[e.valid_up_to()],
        })
    }

    /// Reset the record to an empty state while keeping its allocations.
    ///
    /// This is what makes [`crate::FastxReader::read_into`] allocation-free.
    pub fn clear(&mut self) {
        self.id.clear();
        self.description = None;
        self.seq.clear();
        if let Some(q) = self.quality.as_mut() {
            q.clear();
        }
    }

    /// Per-base counts.
    pub fn base_counts(&self) -> BaseCounts {
        BaseCounts::of(&self.seq)
    }

    /// GC fraction over unambiguous bases, `None` when there are none.
    pub fn gc_content(&self) -> Option<f64> {
        seq::gc_content(&self.seq)
    }

    /// Number of `N`/ambiguity bases.
    pub fn ambiguous_count(&self) -> u64 {
        self.base_counts().n
    }

    /// Reverse complement, reversing the quality string as well.
    ///
    /// ```
    /// # use fastx::Sequence;
    /// let read = Sequence::fastq("r", b"ACGT", b"ABCD")?.reverse_complement();
    /// assert_eq!(read.seq, b"ACGT");
    /// assert_eq!(read.quality.unwrap(), b"DCBA");
    /// # Ok::<(), fastx::Error>(())
    /// ```
    pub fn reverse_complement(&self) -> Sequence {
        Sequence {
            id: self.id.clone(),
            description: self.description.clone(),
            seq: seq::reverse_complement(&self.seq),
            quality: self
                .quality
                .as_ref()
                .map(|q| q.iter().rev().copied().collect()),
        }
    }

    /// Reverse complement in place.
    pub fn reverse_complement_in_place(&mut self) {
        seq::reverse_complement_in_place(&mut self.seq);
        if let Some(q) = self.quality.as_mut() {
            q.reverse();
        }
    }

    /// Uppercase the residues in place (undoes soft masking).
    pub fn make_uppercase(&mut self) {
        self.seq.make_ascii_uppercase();
    }

    /// A sub-record covering `range`, carrying the matching quality slice.
    ///
    /// Returns [`Error::OutOfBounds`] if the range does not fit.
    pub fn subseq(&self, range: Range<usize>) -> Result<Sequence> {
        if range.start > range.end || range.end > self.seq.len() {
            return Err(Error::OutOfBounds {
                id: self.id.clone(),
                start: range.start as u64,
                end: range.end as u64,
                length: self.seq.len() as u64,
            });
        }
        Ok(Sequence {
            id: self.id.clone(),
            description: self.description.clone(),
            seq: self.seq[range.clone()].to_vec(),
            quality: self.quality.as_ref().map(|q| q[range].to_vec()),
        })
    }

    /// Keep only `range`, discarding the rest, in place.
    pub fn trim_to(&mut self, range: Range<usize>) -> Result<()> {
        if range.start > range.end || range.end > self.seq.len() {
            return Err(Error::OutOfBounds {
                id: self.id.clone(),
                start: range.start as u64,
                end: range.end as u64,
                length: self.seq.len() as u64,
            });
        }
        self.seq.truncate(range.end);
        self.seq.drain(..range.start);
        if let Some(q) = self.quality.as_mut() {
            q.truncate(range.end);
            q.drain(..range.start);
        }
        Ok(())
    }

    /// Translate the residues into protein using the standard genetic code.
    pub fn translate(&self, frame: usize, stop_at_stop: bool) -> Sequence {
        Sequence {
            id: self.id.clone(),
            description: self.description.clone(),
            seq: seq::translate(&self.seq, frame, stop_at_stop),
            quality: None,
        }
    }

    /// Iterator over overlapping k-mers.
    pub fn kmers(&self, k: usize) -> impl Iterator<Item = &[u8]> {
        seq::kmers(&self.seq, k)
    }

    /// Decoded Phred scores assuming the Phred+33 offset.
    pub fn quality_scores(&self) -> Option<Vec<u8>> {
        self.quality.as_ref().map(|q| qual::scores(q, PHRED33))
    }

    /// Decoded Phred scores for an explicit offset.
    pub fn quality_scores_with(&self, offset: u8) -> Option<Vec<u8>> {
        self.quality.as_ref().map(|q| qual::scores(q, offset))
    }

    /// Error-probability-weighted mean quality (Phred+33).
    pub fn mean_quality(&self) -> Option<f64> {
        qual::mean_quality(self.quality.as_deref()?, PHRED33)
    }

    /// Expected number of sequencing errors in the read (Phred+33).
    pub fn expected_errors(&self) -> Option<f64> {
        Some(qual::expected_errors(self.quality.as_deref()?, PHRED33))
    }

    /// Convert Phred+64 quality to Phred+33 in place. No-op for FASTA records.
    pub fn convert_quality_offset(&mut self, from: u8, to: u8) {
        if let Some(q) = self.quality.as_mut() {
            for c in q.iter_mut() {
                *c = qual::encode(qual::score(*c, from), to);
            }
        }
    }

    /// Drop the quality string, turning a FASTQ record into a FASTA record.
    pub fn into_fasta(mut self) -> Sequence {
        self.quality = None;
        self
    }

    /// Check the invariants that the writers rely on.
    ///
    /// * the id is non-empty and free of whitespace/newlines
    /// * the description contains no newlines
    /// * residues are printable and non-whitespace, and in `alphabet`
    /// * quality, when present, has the same length as the sequence and is
    ///   printable ASCII
    pub fn validate(&self, alphabet: Alphabet) -> Result<()> {
        if self.id.is_empty() {
            return Err(Error::Parse {
                line: 0,
                kind: crate::error::ParseError::EmptyId,
            });
        }
        if let Some(pos) = self.id.bytes().position(|b| b.is_ascii_whitespace()) {
            return Err(Error::InvalidByte {
                id: self.id.clone(),
                pos,
                byte: self.id.as_bytes()[pos],
            });
        }
        if let Some(d) = &self.description {
            if let Some(pos) = d.bytes().position(|b| b == b'\n' || b == b'\r') {
                return Err(Error::InvalidByte {
                    id: self.id.clone(),
                    pos,
                    byte: d.as_bytes()[pos],
                });
            }
        }
        alphabet.validate_named(&self.seq, &self.id)?;
        if let Some(q) = &self.quality {
            if q.len() != self.seq.len() {
                return Err(Error::LengthMismatch {
                    id: self.id.clone(),
                    seq: self.seq.len(),
                    quality: q.len(),
                });
            }
            if let Some(pos) = q.iter().position(|&b| !(33..=126).contains(&b)) {
                return Err(Error::InvalidByte {
                    id: self.id.clone(),
                    pos,
                    byte: q[pos],
                });
            }
        }
        Ok(())
    }

    /// Write this record as FASTA, wrapping sequence lines at `line_width`
    /// (`None` writes the sequence on a single line).
    ///
    /// Fails with [`Error::InvalidByte`] if the sequence contains a byte FASTA
    /// cannot represent: `\n`, `\r` or `>`. None of the three survives a
    /// write/read cycle once line wrapping can move it, so writing one would
    /// silently corrupt the record.
    pub fn write_fasta<W: Write>(&self, out: &mut W, line_width: Option<usize>) -> Result<()> {
        check_writable_residues(&self.seq, &self.id)?;
        out.write_all(b">")?;
        self.write_header(out)?;
        match line_width.filter(|w| *w > 0) {
            None => {
                out.write_all(&self.seq)?;
                out.write_all(b"\n")?;
            }
            Some(width) => {
                for chunk in self.seq.chunks(width) {
                    out.write_all(chunk)?;
                    out.write_all(b"\n")?;
                }
                if self.seq.is_empty() {
                    out.write_all(b"\n")?;
                }
            }
        }
        Ok(())
    }

    /// Write this record as FASTQ. Fails when the record has no quality, when
    /// the lengths disagree, or when a byte cannot be represented.
    pub fn write_fastq<W: Write>(&self, out: &mut W) -> Result<()> {
        let quality = self.quality.as_ref().ok_or_else(|| Error::MissingQuality {
            id: self.id.clone(),
        })?;
        if quality.len() != self.seq.len() {
            return Err(Error::LengthMismatch {
                id: self.id.clone(),
                seq: self.seq.len(),
                quality: quality.len(),
            });
        }
        check_writable_residues(&self.seq, &self.id)?;
        check_writable_fastq_sequence(&self.seq, &self.id)?;
        check_writable_quality(quality, &self.id)?;
        out.write_all(b"@")?;
        self.write_header(out)?;
        out.write_all(&self.seq)?;
        out.write_all(b"\n+\n")?;
        out.write_all(quality)?;
        out.write_all(b"\n")?;
        Ok(())
    }

    fn write_header<W: Write>(&self, out: &mut W) -> Result<()> {
        out.write_all(self.id.as_bytes())?;
        if let Some(d) = &self.description {
            out.write_all(b" ")?;
            out.write_all(d.as_bytes())?;
        }
        out.write_all(b"\n")?;
        Ok(())
    }

    /// Render the record in its native format as a `String`.
    pub fn to_string_in(&self, format: Format) -> Result<String> {
        let mut buf = Vec::with_capacity(self.seq.len() * 2 + 64);
        match format {
            Format::Fasta => self.write_fasta(&mut buf, Some(60))?,
            Format::Fastq => self.write_fastq(&mut buf)?,
        }
        Ok(String::from_utf8_lossy(&buf).into_owned())
    }

    /// Split a raw header (without the `>`/`@`) into id and description.
    ///
    /// Splitting happens on ASCII whitespace over the raw bytes. Doing it on
    /// `char::is_whitespace` over a `str` would be wrong twice over: a header
    /// containing, say, a non-breaking space would be split in a place other
    /// tools do not split, and slicing after a multi-byte whitespace character
    /// would land inside it and panic.
    pub(crate) fn set_header(&mut self, header: &[u8]) {
        let header = trim_ascii_end(header);
        let id = header_id(header);
        self.id.push_str(&String::from_utf8_lossy(id));
        let description = trim_ascii_start(&header[id.len()..]);
        self.description = if description.is_empty() {
            None
        } else {
            Some(String::from_utf8_lossy(description).into_owned())
        };
    }
}

/// The identifier part of a raw header: everything up to the first ASCII
/// whitespace byte.
///
/// The FASTA indexer uses this too. Both must agree on where an id ends, or an
/// index will not contain the names the reader produces — which is precisely the
/// bug that appears when one of them splits on Unicode whitespace and the other
/// does not.
pub(crate) fn header_id(header: &[u8]) -> &[u8] {
    let end = header
        .iter()
        .position(|b| b.is_ascii_whitespace())
        .unwrap_or(header.len());
    &header[..end]
}

/// Reject sequence bytes that FASTA and FASTQ cannot represent.
///
/// Neither format has an escape mechanism, so some bytes simply cannot survive a
/// write/read cycle and writing them would silently corrupt data:
///
/// * `\n` ends a line, so a sequence containing one would be read back as two.
/// * `\r` is stripped when it precedes a line break, and line wrapping can move
///   any `\r` to the end of a line — so `"AC\r"` would come back as `"AC"`.
/// * `>` starts a header, and wrapping can move any `>` to the start of a line,
///   splitting one record into two.
///
/// The check is a single SIMD pass, so it costs almost nothing next to the write
/// itself. Refusing to write is the only honest option here: the alternative is
/// output that this crate cannot read back.
pub(crate) fn check_writable_residues(seq: &[u8], id: &str) -> Result<()> {
    match memchr::memchr3(b'\n', b'\r', b'>', seq) {
        None => Ok(()),
        Some(pos) => Err(Error::InvalidByte {
            id: id.to_string(),
            pos,
            byte: seq[pos],
        }),
    }
}

/// A FASTQ sequence may not begin with `+`, which would be read as the separator.
pub(crate) fn check_writable_fastq_sequence(seq: &[u8], id: &str) -> Result<()> {
    if seq.first() == Some(&b'+') {
        return Err(Error::InvalidByte {
            id: id.to_string(),
            pos: 0,
            byte: b'+',
        });
    }
    Ok(())
}

/// Quality strings are length-delimited, so only line breaks are unrepresentable.
pub(crate) fn check_writable_quality(quality: &[u8], id: &str) -> Result<()> {
    match memchr::memchr2(b'\n', b'\r', quality) {
        None => Ok(()),
        Some(pos) => Err(Error::InvalidByte {
            id: id.to_string(),
            pos,
            byte: quality[pos],
        }),
    }
}

/// `[u8]::trim_ascii_start` is stable only from Rust 1.80; our MSRV is 1.74.
fn trim_ascii_start(mut bytes: &[u8]) -> &[u8] {
    while let [first, rest @ ..] = bytes {
        if first.is_ascii_whitespace() {
            bytes = rest;
        } else {
            break;
        }
    }
    bytes
}

fn trim_ascii_end(mut bytes: &[u8]) -> &[u8] {
    while let [rest @ .., last] = bytes {
        if last.is_ascii_whitespace() {
            bytes = rest;
        } else {
            break;
        }
    }
    bytes
}

impl fmt::Display for Sequence {
    /// FASTA with 60-column wrapping, or FASTQ when quality is present.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut buf = Vec::new();
        let rendered = match self.format() {
            Format::Fasta => self.write_fasta(&mut buf, Some(60)),
            Format::Fastq => self.write_fastq(&mut buf),
        };
        rendered.map_err(|_| fmt::Error)?;
        f.write_str(&String::from_utf8_lossy(&buf))
    }
}

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

    #[test]
    fn fastq_requires_matching_lengths() {
        assert!(Sequence::fastq("r", b"ACGT", b"III").is_err());
        assert!(Sequence::fastq("r", b"ACGT", b"IIII").is_ok());
    }

    #[test]
    fn header_splitting() {
        let mut s = Sequence::default();
        s.set_header(b"chr1 human chromosome 1");
        assert_eq!(s.id, "chr1");
        assert_eq!(s.description.as_deref(), Some("human chromosome 1"));

        let mut s = Sequence::default();
        s.set_header(b"lonely");
        assert_eq!(s.id, "lonely");
        assert_eq!(s.description, None);

        let mut s = Sequence::default();
        s.set_header(b"tabbed\tdesc  with   spaces  ");
        assert_eq!(s.id, "tabbed");
        assert_eq!(s.description.as_deref(), Some("desc  with   spaces"));
    }

    #[test]
    fn header_splitting_handles_non_ascii() {
        // A non-breaking space (U+00A0) is not a separator, and slicing around
        // it must not land inside the character.
        let mut s = Sequence::default();
        s.set_header("chr\u{a0}1 description".as_bytes());
        assert_eq!(s.id, "chr\u{a0}1");
        assert_eq!(s.description.as_deref(), Some("description"));

        // Invalid UTF-8 is replaced rather than rejected.
        let mut s = Sequence::default();
        s.set_header(&[b'i', b'd', 0xff, b' ', b'd']);
        assert!(s.id.starts_with("id"));
        assert_eq!(s.description.as_deref(), Some("d"));

        // A header that is nothing but whitespace leaves an empty id, which the
        // reader turns into a parse error.
        let mut s = Sequence::default();
        s.set_header(b"   \t ");
        assert!(s.id.is_empty());
    }

    #[test]
    fn fasta_wrapping() {
        let s = Sequence::fasta("x", b"AAAAACCCCC".to_vec());
        let mut out = Vec::new();
        s.write_fasta(&mut out, Some(5)).unwrap();
        assert_eq!(out, b">x\nAAAAA\nCCCCC\n");

        let mut out = Vec::new();
        s.write_fasta(&mut out, None).unwrap();
        assert_eq!(out, b">x\nAAAAACCCCC\n");

        // An empty sequence still gets a (blank) sequence line.
        let mut out = Vec::new();
        Sequence::fasta("empty", Vec::new())
            .write_fasta(&mut out, Some(60))
            .unwrap();
        assert_eq!(out, b">empty\n\n");
    }

    #[test]
    fn fastq_output_and_missing_quality() {
        let r = Sequence::fastq("r", b"ACGT", b"IIII")
            .unwrap()
            .with_description("d");
        let mut out = Vec::new();
        r.write_fastq(&mut out).unwrap();
        assert_eq!(out, b"@r d\nACGT\n+\nIIII\n");

        let mut out = Vec::new();
        assert!(matches!(
            Sequence::fasta("r", b"ACGT".to_vec()).write_fastq(&mut out),
            Err(Error::MissingQuality { .. })
        ));
    }

    #[test]
    fn trimming_keeps_quality_aligned() {
        let mut r = Sequence::fastq("r", b"AACCGGTT", b"01234567").unwrap();
        r.trim_to(2..6).unwrap();
        assert_eq!(r.seq, b"CCGG");
        assert_eq!(r.quality.as_deref(), Some(&b"2345"[..]));
        assert!(r.trim_to(0..99).is_err());
    }

    #[test]
    fn subseq_bounds() {
        let r = Sequence::fastq("r", b"AACCGG", b"012345").unwrap();
        let sub = r.subseq(1..3).unwrap();
        assert_eq!(sub.seq, b"AC");
        assert_eq!(sub.quality.as_deref(), Some(&b"12"[..]));
        #[allow(clippy::reversed_empty_ranges)] // deliberately backwards
        {
            assert!(r.subseq(4..2).is_err());
        }
        assert!(r.subseq(0..7).is_err());
    }

    #[test]
    fn validation_catches_bad_records() {
        let mut r = Sequence::fasta("ok", b"ACGT".to_vec());
        assert!(r.validate(Alphabet::Dna).is_ok());
        r.id = "has space".into();
        assert!(r.validate(Alphabet::Dna).is_err());
        r.id = String::new();
        assert!(r.validate(Alphabet::Dna).is_err());

        let mut r = Sequence::fastq("r", b"ACGT", b"IIII").unwrap();
        r.quality = Some(b"II".to_vec());
        assert!(matches!(
            r.validate(Alphabet::Dna),
            Err(Error::LengthMismatch { .. })
        ));
        r.quality = Some(b"II\nI".to_vec());
        assert!(matches!(
            r.validate(Alphabet::Dna),
            Err(Error::InvalidByte { .. })
        ));
    }

    #[test]
    fn refuses_to_write_unrepresentable_residues() {
        // Regression: the reader accepts a bare '\r' inside a sequence line, but
        // writing it would put the '\r' before a newline, where reading strips
        // it — so ">x\nAC\r\n" would come back as "AC". Refuse instead.
        let mut out = Vec::new();
        for bad in [&b"AC\r"[..], b"AC\nGT", b"AC>GT", b">AC"] {
            let record = Sequence::fasta("x", bad.to_vec());
            assert!(
                matches!(
                    record.write_fasta(&mut out, Some(60)),
                    Err(Error::InvalidByte { .. })
                ),
                "{:?} should be rejected",
                String::from_utf8_lossy(bad)
            );
        }
        // Ordinary residues, gaps and stops are all still fine.
        for good in [&b"ACGTN"[..], b"acgt-n.", b"MEEPQSDPSV*"] {
            let record = Sequence::fasta("x", good.to_vec());
            assert!(record.write_fasta(&mut out, Some(60)).is_ok());
        }
    }

    #[test]
    fn refuses_to_write_unrepresentable_fastq() {
        let mut out = Vec::new();

        // A leading '+' in the sequence would be read back as the separator.
        let record = Sequence::fastq("x", b"+CGT", b"IIII").unwrap();
        assert!(matches!(
            record.write_fastq(&mut out),
            Err(Error::InvalidByte {
                pos: 0,
                byte: b'+',
                ..
            })
        ));
        // '+' anywhere else is harmless.
        assert!(Sequence::fastq("x", b"A+GT", b"IIII")
            .unwrap()
            .write_fastq(&mut out)
            .is_ok());

        // A line break in the quality string breaks the record.
        let record = Sequence::fastq("x", b"ACGT", b"II\nI").unwrap();
        assert!(matches!(
            record.write_fastq(&mut out),
            Err(Error::InvalidByte { byte: b'\n', .. })
        ));
        // '@' in quality is legal (Q31) and must not be rejected.
        assert!(Sequence::fastq("x", b"ACGT", b"@@@@")
            .unwrap()
            .write_fastq(&mut out)
            .is_ok());
    }

    #[test]
    fn quality_offset_conversion() {
        let mut r = Sequence::fastq("r", b"ACGT", b"hhhh").unwrap();
        r.convert_quality_offset(crate::qual::PHRED64, PHRED33);
        assert_eq!(r.quality.as_deref(), Some(&b"IIII"[..]));
    }

    #[test]
    fn clear_keeps_capacity() {
        let mut r = Sequence::fastq("r", b"ACGT", b"IIII").unwrap();
        let cap = r.seq.capacity();
        r.clear();
        assert!(r.id.is_empty() && r.seq.is_empty());
        assert_eq!(r.quality.as_deref(), Some(&b""[..]));
        assert_eq!(r.seq.capacity(), cap);
    }
}