fastx-io 0.3.0

//! FASTA/FASTQ writer with optional gzip compression and validation.

use std::fmt;
use std::fs::File;
use std::io::{self, BufWriter, Write};
use std::path::Path;

use crate::error::{Error, Result};
use crate::format::{Compression, CompressionLevel, Format};
use crate::record::Sequence;
use crate::seq::Alphabet;

/// Default FASTA line width. NCBI uses 70, UCSC and samtools use 60.
pub const DEFAULT_LINE_WIDTH: usize = 60;

/// A writer for FASTA and FASTQ records.
///
/// FASTA sequence lines are wrapped at [`DEFAULT_LINE_WIDTH`] unless configured
/// otherwise. Writing a FASTQ record without quality scores is an error rather
/// than a silent fabrication of quality values; writing a FASTQ record to a FASTA
/// writer simply drops the quality, which is the usual `fq2fa` conversion.
///
/// ```
/// use fastx::{FastxWriter, Format, Sequence};
///
/// let mut out = Vec::new();
/// {
///     let mut writer = FastxWriter::new(&mut out, Format::Fasta).line_width(4);
///     writer.write_record(&Sequence::fasta("a", b"ACGTAC"))?;
///     writer.flush()?;
/// }
/// assert_eq!(out, b">a\nACGT\nAC\n");
/// # Ok::<(), fastx::Error>(())
/// ```
pub struct FastxWriter<W: Write> {
    inner: W,
    format: Format,
    line_width: Option<usize>,
    validate: Option<Alphabet>,
    written: u64,
}

impl<W: Write> FastxWriter<W> {
    /// A writer for `format`, wrapping FASTA at [`DEFAULT_LINE_WIDTH`].
    ///
    /// Wrap `inner` in a [`BufWriter`] yourself for best throughput, or use
    /// [`create`] which does it for you.
    pub fn new(inner: W, format: Format) -> FastxWriter<W> {
        FastxWriter {
            inner,
            format,
            line_width: Some(DEFAULT_LINE_WIDTH),
            validate: None,
            written: 0,
        }
    }

    /// Set the FASTA line width. `0` disables wrapping.
    pub fn line_width(mut self, width: usize) -> Self {
        self.line_width = if width == 0 { None } else { Some(width) };
        self
    }

    /// Validate every record against `alphabet` before writing it.
    pub fn validate(mut self, alphabet: Alphabet) -> Self {
        self.validate = Some(alphabet);
        self
    }

    /// The output format.
    pub fn format(&self) -> Format {
        self.format
    }

    /// Number of records written so far.
    pub fn records_written(&self) -> u64 {
        self.written
    }

    /// Write one record in the writer's format.
    pub fn write_record(&mut self, record: &Sequence) -> Result<()> {
        if let Some(alphabet) = self.validate {
            record.validate(alphabet)?;
        }
        match self.format {
            Format::Fasta => record.write_fasta(&mut self.inner, self.line_width)?,
            Format::Fastq => record.write_fastq(&mut self.inner)?,
        }
        self.written += 1;
        Ok(())
    }

    /// Write every record of an iterator.
    pub fn write_all<'a, I>(&mut self, records: I) -> Result<()>
    where
        I: IntoIterator<Item = &'a Sequence>,
    {
        for record in records {
            self.write_record(record)?;
        }
        Ok(())
    }

    /// Write a FASTA record from parts, without building a [`Sequence`].
    pub fn write_fasta(&mut self, id: &str, description: Option<&str>, seq: &[u8]) -> Result<()> {
        crate::record::check_writable_residues(seq, id)?;
        self.write_header(b'>', id, description)?;
        match self.line_width {
            None => {
                self.inner.write_all(seq)?;
                self.inner.write_all(b"\n")?;
            }
            Some(width) => {
                if seq.is_empty() {
                    self.inner.write_all(b"\n")?;
                }
                for chunk in seq.chunks(width) {
                    self.inner.write_all(chunk)?;
                    self.inner.write_all(b"\n")?;
                }
            }
        }
        self.written += 1;
        Ok(())
    }

    /// Write a FASTQ record from parts, without building a [`Sequence`].
    pub fn write_fastq(
        &mut self,
        id: &str,
        description: Option<&str>,
        seq: &[u8],
        quality: &[u8],
    ) -> Result<()> {
        if seq.len() != quality.len() {
            return Err(Error::LengthMismatch {
                id: id.to_string(),
                seq: seq.len(),
                quality: quality.len(),
            });
        }
        crate::record::check_writable_residues(seq, id)?;
        crate::record::check_writable_fastq_sequence(seq, id)?;
        crate::record::check_writable_quality(quality, id)?;
        self.write_header(b'@', id, description)?;
        self.inner.write_all(seq)?;
        self.inner.write_all(b"\n+\n")?;
        self.inner.write_all(quality)?;
        self.inner.write_all(b"\n")?;
        self.written += 1;
        Ok(())
    }

    fn write_header(&mut self, prefix: u8, id: &str, description: Option<&str>) -> Result<()> {
        self.inner.write_all(&[prefix])?;
        self.inner.write_all(id.as_bytes())?;
        if let Some(description) = description.filter(|d| !d.is_empty()) {
            self.inner.write_all(b" ")?;
            self.inner.write_all(description.as_bytes())?;
        }
        self.inner.write_all(b"\n")?;
        Ok(())
    }

    /// Flush buffered bytes to the underlying writer.
    pub fn flush(&mut self) -> Result<()> {
        self.inner.flush()?;
        Ok(())
    }

    /// Flush and return the underlying writer.
    ///
    /// Always prefer this over dropping the writer: for gzip output it is the
    /// only way to learn that finishing the stream failed.
    pub fn finish(mut self) -> Result<W> {
        self.inner.flush()?;
        Ok(self.inner)
    }

    /// Borrow the underlying writer.
    pub fn get_ref(&self) -> &W {
        &self.inner
    }
}

impl<W: Write> fmt::Debug for FastxWriter<W> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("FastxWriter")
            .field("format", &self.format)
            .field("line_width", &self.line_width)
            .field("validate", &self.validate)
            .field("records_written", &self.written)
            .finish_non_exhaustive()
    }
}

/// Configuration for a [`FastxWriter`].
///
/// ```no_run
/// use fastx::{Alphabet, WriterBuilder};
///
/// // Format and gzip are inferred from the path: this writes gzipped FASTQ.
/// let mut writer = WriterBuilder::new()
///     .line_width(80)
///     .validate(Alphabet::Dna)
///     .create("out.fq.gz")?;
/// writer.finish()?;
/// # Ok::<(), fastx::Error>(())
/// ```
#[derive(Debug, Clone)]
pub struct WriterBuilder {
    format: Option<Format>,
    line_width: Option<usize>,
    compression: Option<Compression>,
    level: CompressionLevel,
    validate: Option<Alphabet>,
    buffer_size: usize,
    blocks_per_batch: Option<usize>,
}

impl Default for WriterBuilder {
    fn default() -> Self {
        WriterBuilder {
            format: None,
            line_width: Some(DEFAULT_LINE_WIDTH),
            compression: None,
            level: CompressionLevel::default(),
            validate: None,
            buffer_size: 128 * 1024,
            blocks_per_batch: None,
        }
    }
}

impl WriterBuilder {
    /// A builder with default settings.
    pub fn new() -> WriterBuilder {
        WriterBuilder::default()
    }

    /// Force the output format instead of inferring it from the path.
    pub fn format(mut self, format: Format) -> Self {
        self.format = Some(format);
        self
    }

    /// FASTA line width; `0` disables wrapping.
    pub fn line_width(mut self, width: usize) -> Self {
        self.line_width = if width == 0 { None } else { Some(width) };
        self
    }

    /// Force compression instead of inferring it from the path.
    pub fn compression(mut self, compression: Compression) -> Self {
        self.compression = Some(compression);
        self
    }

    /// gzip compression level.
    pub fn level(mut self, level: CompressionLevel) -> Self {
        self.level = level;
        self
    }

    /// Validate records against an alphabet before writing.
    pub fn validate(mut self, alphabet: Alphabet) -> Self {
        self.validate = Some(alphabet);
        self
    }

    /// Output buffer size in bytes.
    pub fn buffer_size(mut self, bytes: usize) -> Self {
        self.buffer_size = bytes;
        self
    }

    /// How many BGZF blocks to compress at a time.
    ///
    /// Only affects BGZF output. Blocks are independent, so a batch is
    /// compressed across cores and the bytes written are identical either way —
    /// this trades `blocks × 64 KiB` of memory for parallelism. The default is
    /// several blocks per core; pass 1 to force single-threaded compression.
    pub fn blocks_per_batch(mut self, blocks: usize) -> Self {
        self.blocks_per_batch = Some(blocks.max(1));
        self
    }

    /// Build a writer around any [`Write`]. Requires a format to be set.
    pub fn build<W: Write>(&self, inner: W) -> Result<FastxWriter<W>> {
        let format = self.format.ok_or(Error::UnknownFormat {
            hint: "no format given and no path to infer it from".to_string(),
        })?;
        Ok(FastxWriter {
            inner,
            format,
            line_width: self.line_width,
            validate: self.validate,
            written: 0,
        })
    }

    /// Create a file, inferring format and gzip compression from its extension.
    pub fn create<P: AsRef<Path>>(&self, path: P) -> Result<BoxedWriter> {
        let path = path.as_ref();
        let format = match self.format.or_else(|| Format::from_path(path)) {
            Some(format) => format,
            None => {
                return Err(Error::UnknownFormat {
                    hint: format!("unrecognised extension in {}", path.display()),
                })
            }
        };
        // Compressed output defaults to BGZF rather than plain gzip. It is valid
        // gzip either way, so nothing downstream breaks, but BGZF can later be
        // indexed and randomly accessed, which plain gzip never can. Ask for
        // `Compression::Gzip` explicitly to get a single deflate stream.
        let compression = self
            .compression
            .unwrap_or(match Compression::from_path(path) {
                Compression::Gzip => Compression::Bgzf,
                other => other,
            });
        let file = File::create(path)
            .map_err(|e| Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display()))))?;
        let buffered = BufWriter::with_capacity(self.buffer_size, file);
        let sink: Box<dyn Write + Send> = match compression {
            Compression::None => Box::new(buffered),
            Compression::Gzip => gzip_writer(buffered, self.level)?,
            Compression::Bgzf => bgzf_writer(buffered, self.level, self.blocks_per_batch)?,
            Compression::Zstd => zstd_writer(buffered, self.level)?,
        };
        Ok(FastxWriter {
            inner: sink,
            format,
            line_width: self.line_width,
            validate: self.validate,
            written: 0,
        })
    }

    /// Write to standard output.
    pub fn stdout(&self) -> Result<BoxedWriter> {
        let format = self.format.unwrap_or(Format::Fasta);
        let sink: Box<dyn Write + Send> = match self.compression.unwrap_or(Compression::None) {
            Compression::None => Box::new(BufWriter::with_capacity(self.buffer_size, io::stdout())),
            Compression::Gzip => gzip_writer(
                BufWriter::with_capacity(self.buffer_size, io::stdout()),
                self.level,
            )?,
            Compression::Bgzf => bgzf_writer(
                BufWriter::with_capacity(self.buffer_size, io::stdout()),
                self.level,
                self.blocks_per_batch,
            )?,
            Compression::Zstd => zstd_writer(
                BufWriter::with_capacity(self.buffer_size, io::stdout()),
                self.level,
            )?,
        };
        Ok(FastxWriter {
            inner: sink,
            format,
            line_width: self.line_width,
            validate: self.validate,
            written: 0,
        })
    }
}

#[cfg(feature = "gzip")]
fn gzip_writer<W: Write + Send + 'static>(
    sink: W,
    level: CompressionLevel,
) -> Result<Box<dyn Write + Send>> {
    Ok(Box::new(flate2::write::GzEncoder::new(
        sink,
        flate2::Compression::new(level.0.min(9)),
    )))
}

#[cfg(not(feature = "gzip"))]
fn gzip_writer<W: Write + Send + 'static>(
    _sink: W,
    _level: CompressionLevel,
) -> Result<Box<dyn Write + Send>> {
    Err(Error::FeatureDisabled("gzip"))
}

#[cfg(feature = "gzip")]
fn bgzf_writer<W: Write + Send + 'static>(
    sink: W,
    level: CompressionLevel,
    blocks_per_batch: Option<usize>,
) -> Result<Box<dyn Write + Send>> {
    let mut writer = crate::bgzf::BgzfWriter::with_level(sink, level);
    if let Some(blocks) = blocks_per_batch {
        writer = writer.blocks_per_batch(blocks);
    }
    Ok(Box::new(writer))
}

#[cfg(not(feature = "gzip"))]
fn bgzf_writer<W: Write + Send + 'static>(
    _sink: W,
    _level: CompressionLevel,
    _blocks_per_batch: Option<usize>,
) -> Result<Box<dyn Write + Send>> {
    Err(Error::FeatureDisabled("gzip"))
}

/// Zstandard output.
///
/// `auto_finish` matters: a zstd frame needs an epilogue, and without it the
/// file is truncated. This mirrors how flate2's encoder finishes on drop, so
/// both compressed paths behave the same whether the caller calls `finish` or
/// simply drops the writer.
#[cfg(feature = "zstd")]
fn zstd_writer<W: Write + Send + 'static>(
    sink: W,
    level: CompressionLevel,
) -> Result<Box<dyn Write + Send>> {
    // gzip levels run 0-9, zstd 1-22. Map the level across rather than passing
    // a gzip number to zstd, where 9 would be a middling setting.
    let zstd_level = match level.0 {
        0 => 1,
        level => (level.min(9) as i32 - 1) * 21 / 8 + 1,
    };
    let encoder = zstd::stream::write::Encoder::new(sink, zstd_level)?;
    Ok(Box::new(encoder.auto_finish()))
}

#[cfg(not(feature = "zstd"))]
fn zstd_writer<W: Write + Send + 'static>(
    _sink: W,
    _level: CompressionLevel,
) -> Result<Box<dyn Write + Send>> {
    Err(Error::FeatureDisabled("zstd"))
}

/// The boxed writer type produced by [`create`] and [`WriterBuilder::stdout`].
pub type BoxedWriter = FastxWriter<Box<dyn Write + Send>>;

/// Create a sequence file, inferring format and gzip compression from the path.
///
/// ```no_run
/// let mut writer = fastx::create("contigs.fasta")?;
/// writer.write_record(&fastx::Sequence::fasta("contig1", b"ACGT"))?;
/// writer.finish()?;
/// # Ok::<(), fastx::Error>(())
/// ```
pub fn create<P: AsRef<Path>>(path: P) -> Result<BoxedWriter> {
    WriterBuilder::default().create(path)
}

/// Write records to standard output in `format`.
pub fn stdout(format: Format) -> Result<BoxedWriter> {
    WriterBuilder::default().format(format).stdout()
}

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

    fn write_to_string(format: Format, records: &[Sequence], width: usize) -> String {
        let mut out = Vec::new();
        let mut writer = FastxWriter::new(&mut out, format).line_width(width);
        writer.write_all(records).unwrap();
        writer.flush().unwrap();
        String::from_utf8(out).unwrap()
    }

    #[test]
    fn writes_wrapped_fasta() {
        let records = vec![Sequence::fasta("a", b"ACGTACGTAC").with_description("d")];
        assert_eq!(
            write_to_string(Format::Fasta, &records, 4),
            ">a d\nACGT\nACGT\nAC\n"
        );
        assert_eq!(
            write_to_string(Format::Fasta, &records, 0),
            ">a d\nACGTACGTAC\n"
        );
    }

    #[test]
    fn drops_quality_when_writing_fasta() {
        let records = vec![Sequence::fastq("a", b"ACGT", b"IIII").unwrap()];
        assert_eq!(write_to_string(Format::Fasta, &records, 60), ">a\nACGT\n");
    }

    #[test]
    fn refuses_fastq_without_quality() {
        let mut out = Vec::new();
        let mut writer = FastxWriter::new(&mut out, Format::Fastq);
        let err = writer
            .write_record(&Sequence::fasta("a", b"ACGT"))
            .unwrap_err();
        assert!(matches!(err, Error::MissingQuality { .. }));
    }

    #[test]
    fn validates_on_demand() {
        let mut out = Vec::new();
        let mut writer = FastxWriter::new(&mut out, Format::Fasta).validate(Alphabet::Dna);
        assert!(writer.write_record(&Sequence::fasta("a", b"ACGT")).is_ok());
        let err = writer
            .write_record(&Sequence::fasta("b", b"ACGX"))
            .unwrap_err();
        assert!(matches!(err, Error::InvalidByte { byte: b'X', .. }));
        assert_eq!(writer.records_written(), 1);
    }

    #[test]
    fn writes_parts_without_records() {
        let mut out = Vec::new();
        let mut writer = FastxWriter::new(&mut out, Format::Fastq);
        writer
            .write_fastq("r1", Some("desc"), b"ACGT", b"IIII")
            .unwrap();
        writer.write_fasta("r2", None, b"TTTT").unwrap();
        writer.flush().unwrap();
        assert_eq!(out, b"@r1 desc\nACGT\n+\nIIII\n>r2\nTTTT\n");

        let mut out = Vec::new();
        let mut writer = FastxWriter::new(&mut out, Format::Fastq);
        assert!(writer.write_fastq("r", None, b"ACGT", b"II").is_err());
    }

    #[test]
    fn round_trips_through_reader() {
        let original: Vec<Sequence> = (0..50)
            .map(|i| {
                Sequence::fastq(
                    format!("read{i}"),
                    "ACGTN".repeat(i + 1).into_bytes(),
                    "IIIII".repeat(i + 1).into_bytes(),
                )
                .unwrap()
                .with_description(format!("record number {i}"))
            })
            .collect();

        let text = write_to_string(Format::Fastq, &original, 60);
        let parsed: Vec<Sequence> = FastxReader::new(text.as_bytes())
            .collect::<Result<Vec<_>>>()
            .unwrap();
        assert_eq!(parsed, original);

        // FASTA round trip re-joins wrapped lines.
        let fasta_input: Vec<Sequence> =
            original.iter().cloned().map(Sequence::into_fasta).collect();
        let text = write_to_string(Format::Fasta, &fasta_input, 7);
        let parsed: Vec<Sequence> = FastxReader::new(text.as_bytes())
            .collect::<Result<Vec<_>>>()
            .unwrap();
        assert_eq!(parsed, fasta_input);
    }

    #[test]
    fn empty_sequence_round_trips() {
        let records = vec![
            Sequence::fasta("a", Vec::new()),
            Sequence::fasta("b", b"AC"),
        ];
        let text = write_to_string(Format::Fasta, &records, 60);
        assert_eq!(text, ">a\n\n>b\nAC\n");
        let parsed: Vec<Sequence> = FastxReader::new(text.as_bytes())
            .collect::<Result<Vec<_>>>()
            .unwrap();
        assert_eq!(parsed, records);
    }

    #[test]
    fn builder_requires_a_format() {
        let mut out = Vec::new();
        assert!(matches!(
            WriterBuilder::default().build(&mut out),
            Err(Error::UnknownFormat { .. })
        ));
        assert!(WriterBuilder::default()
            .format(Format::Fasta)
            .build(&mut out)
            .is_ok());
    }
}