fastx-io 0.3.0

//! Parallel processing helpers (requires the `parallel` feature).
//!
//! Two strategies are available:
//!
//! * [`par_for_each`] and friends parse on one thread and hand batches of
//!   records to a rayon pool. This works with any reader, including gzip and
//!   standard input, and is the right default when the work per record is
//!   non-trivial.
//! * [`par_map_chunks_file`] splits an *uncompressed* file into byte ranges,
//!   snaps each range to a record boundary and parses the ranges in parallel.
//!   This is the only way to make parsing itself scale across cores.
//!
//! ```
//! use fastx::{FastxReader, parallel};
//!
//! let data = b">a\nACGT\n>b\nGGCC\n>c\nAAAA\n";
//! let gc: Vec<Option<f64>> =
//!     parallel::par_map(&mut FastxReader::new(&data[..]), 64, |r| r.gc_content())?;
//! assert_eq!(gc, vec![Some(0.5), Some(1.0), Some(0.0)]);
//! # Ok::<(), fastx::Error>(())
//! ```

use std::fs::File;
use std::io::{self, BufReader, Read, Seek, SeekFrom, Take};
use std::ops::Range;
use std::path::Path;

use rayon::prelude::*;

use crate::error::{Error, Result};
use crate::format::Format;
use crate::reader::FastxReader;
use crate::record::Sequence;
use crate::stats::SeqStats;

/// Default number of records handed to the pool at a time.
pub const DEFAULT_CHUNK_SIZE: usize = 4096;

/// Reader type handed to each worker by [`par_map_chunks_file`].
pub type ChunkReader = FastxReader<Take<File>>;

/// A reusable batch of records, so that a long run does not reallocate.
struct Batch {
    records: Vec<Sequence>,
    filled: usize,
}

impl Batch {
    fn new(capacity: usize) -> Batch {
        Batch {
            records: Vec::with_capacity(capacity),
            filled: 0,
        }
    }

    /// Refill from `reader`, returning the number of records read.
    fn refill<R: Read>(&mut self, reader: &mut FastxReader<R>, capacity: usize) -> Result<usize> {
        self.filled = 0;
        while self.filled < capacity {
            if self.records.len() == self.filled {
                self.records.push(Sequence::default());
            }
            if !reader.read_into(&mut self.records[self.filled])? {
                break;
            }
            self.filled += 1;
        }
        Ok(self.filled)
    }

    fn as_slice(&self) -> &[Sequence] {
        &self.records[..self.filled]
    }
}

/// Apply `f` to every record, in parallel across batches.
///
/// Records within a batch are processed in an unspecified order. `f` must be
/// `Sync` because it runs on several threads at once — use atomics or a `Mutex`
/// if it needs to accumulate.
pub fn par_for_each<R, F>(reader: &mut FastxReader<R>, chunk_size: usize, f: F) -> Result<()>
where
    R: Read,
    F: Fn(&Sequence) -> Result<()> + Send + Sync,
{
    let chunk_size = chunk_size.max(1);
    let mut batch = Batch::new(chunk_size);
    while batch.refill(reader, chunk_size)? > 0 {
        batch.as_slice().par_iter().try_for_each(&f)?;
    }
    Ok(())
}

/// Map every record through `f` in parallel, preserving input order.
pub fn par_map<R, T, F>(reader: &mut FastxReader<R>, chunk_size: usize, f: F) -> Result<Vec<T>>
where
    R: Read,
    T: Send,
    F: Fn(&Sequence) -> T + Send + Sync,
{
    let chunk_size = chunk_size.max(1);
    let mut batch = Batch::new(chunk_size);
    let mut out = Vec::new();
    while batch.refill(reader, chunk_size)? > 0 {
        let mut mapped: Vec<T> = Vec::new();
        batch
            .as_slice()
            .par_iter()
            .map(&f)
            .collect_into_vec(&mut mapped);
        out.append(&mut mapped);
    }
    Ok(out)
}

/// Fold every record into a single value in parallel.
///
/// `identity` builds a fresh accumulator per thread, `fold` adds one record and
/// `merge` combines two accumulators. This is the pattern for parallel counting
/// without a lock.
///
/// ```
/// use fastx::{FastxReader, parallel};
///
/// let data = b">a\nACGT\n>b\nGGCCGG\n";
/// let bases = parallel::par_fold(
///     &mut FastxReader::new(&data[..]),
///     1024,
///     || 0usize,
///     |acc, record| acc + record.len(),
///     |a, b| a + b,
/// )?;
/// assert_eq!(bases, 10);
/// # Ok::<(), fastx::Error>(())
/// ```
pub fn par_fold<R, T, I, F, M>(
    reader: &mut FastxReader<R>,
    chunk_size: usize,
    identity: I,
    fold: F,
    merge: M,
) -> Result<T>
where
    R: Read,
    T: Send,
    I: Fn() -> T + Send + Sync,
    F: Fn(T, &Sequence) -> T + Send + Sync,
    M: Fn(T, T) -> T + Send + Sync,
{
    let chunk_size = chunk_size.max(1);
    let mut batch = Batch::new(chunk_size);
    let mut accumulator = identity();
    while batch.refill(reader, chunk_size)? > 0 {
        let folded = batch
            .as_slice()
            .par_iter()
            .fold(&identity, &fold)
            .reduce(&identity, &merge);
        accumulator = merge(accumulator, folded);
    }
    Ok(accumulator)
}

/// Compute [`SeqStats`] using all cores.
pub fn par_stats<R: Read>(reader: &mut FastxReader<R>, chunk_size: usize) -> Result<SeqStats> {
    par_fold(
        reader,
        chunk_size,
        SeqStats::new,
        |mut stats, record| {
            stats.push(record);
            stats
        },
        |mut a, b| {
            a.merge(&b);
            a
        },
    )
}

/// Split an uncompressed FASTA/FASTQ file into `parts` byte ranges that each
/// start exactly on a record boundary.
///
/// The returned ranges cover the whole file, are non-overlapping and in order.
/// Fewer ranges than requested may come back when the file is small or its
/// records are large.
///
/// ```no_run
/// let ranges = fastx::parallel::split_file("reads.fq", 8)?;
/// assert!(ranges.windows(2).all(|w| w[0].end == w[1].start));
/// # Ok::<(), fastx::Error>(())
/// ```
pub fn split_file<P: AsRef<Path>>(path: P, parts: usize) -> Result<Vec<Range<u64>>> {
    let path = path.as_ref();
    if crate::format::Compression::from_path(path) != crate::format::Compression::None {
        return Err(Error::Unsupported(
            "compressed files cannot be split by byte range; use par_for_each instead",
        ));
    }
    let mut file = File::open(path)
        .map_err(|e| Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display()))))?;
    let size = file.metadata()?.len();
    if size == 0 {
        return Ok(Vec::new());
    }
    let format = sniff_format(&mut file)?;
    let parts = parts.max(1);

    let mut boundaries = vec![0u64];
    for i in 1..parts {
        let candidate = size * i as u64 / parts as u64;
        if candidate <= *boundaries.last().unwrap() {
            continue;
        }
        if let Some(start) = find_record_start(&mut file, candidate, size, format)? {
            if start > *boundaries.last().unwrap() && start < size {
                boundaries.push(start);
            }
        }
    }
    boundaries.push(size);
    Ok(boundaries.windows(2).map(|w| w[0]..w[1]).collect())
}

/// Run `f` on a reader for each chunk of an uncompressed file, in parallel.
///
/// Each worker gets its own file handle limited to its byte range, so parsing
/// itself is parallel. Results come back in file order.
///
/// ```no_run
/// use fastx::parallel;
///
/// // Count records with real parallel parsing.
/// let counts = parallel::par_map_chunks_file("reads.fq", 8, |reader| reader.count_records())?;
/// let total: u64 = counts.iter().sum();
/// # Ok::<(), fastx::Error>(())
/// ```
pub fn par_map_chunks_file<P, T, F>(path: P, parts: usize, f: F) -> Result<Vec<T>>
where
    P: AsRef<Path>,
    T: Send,
    F: Fn(&mut ChunkReader) -> Result<T> + Send + Sync,
{
    let path = path.as_ref();
    let ranges = split_file(path, parts)?;
    if ranges.is_empty() {
        return Ok(Vec::new());
    }
    let format = sniff_format(&mut File::open(path)?)?;
    ranges
        .into_par_iter()
        .map(|range| {
            let mut file = File::open(path)?;
            file.seek(SeekFrom::Start(range.start))?;
            let mut reader = FastxReader::with_format(file.take(range.end - range.start), format);
            f(&mut reader)
        })
        .collect()
}

/// Apply `f` to every record of an uncompressed file with parallel parsing.
pub fn par_for_each_file<P, F>(path: P, parts: usize, f: F) -> Result<()>
where
    P: AsRef<Path>,
    F: Fn(&Sequence) -> Result<()> + Send + Sync,
{
    par_map_chunks_file(path, parts, |reader| reader.for_each_record(&f))?;
    Ok(())
}

/// Compute [`SeqStats`] for an uncompressed file with parallel parsing.
pub fn par_stats_file<P: AsRef<Path>>(path: P, parts: usize) -> Result<SeqStats> {
    let per_chunk = par_map_chunks_file(path, parts, |reader| {
        let mut stats = SeqStats::new();
        reader.for_each_record(|record| {
            stats.push(record);
            Ok(())
        })?;
        Ok(stats)
    })?;
    let mut total = SeqStats::new();
    for stats in &per_chunk {
        total.merge(stats);
    }
    Ok(total)
}

/// Read the first meaningful byte of a file to determine its format.
fn sniff_format(file: &mut File) -> Result<Format> {
    file.seek(SeekFrom::Start(0))?;
    let mut byte = [0u8; 1];
    loop {
        let read = file.read(&mut byte)?;
        if read == 0 {
            return Err(Error::UnknownFormat {
                hint: "file is empty".to_string(),
            });
        }
        if byte[0].is_ascii_whitespace() {
            continue;
        }
        return Format::from_first_byte(byte[0]).ok_or_else(|| Error::UnknownFormat {
            hint: format!("first byte is {:?}", byte[0]),
        });
    }
}

/// Find the offset of the first record that starts at or after `from`.
///
/// For FASTQ a `@` at the start of a line is ambiguous, because `@` is also a
/// quality character, so a candidate is only accepted when the following three
/// lines look like a complete record.
fn find_record_start(file: &mut File, from: u64, size: u64, format: Format) -> Result<Option<u64>> {
    if from >= size {
        return Ok(None);
    }
    file.seek(SeekFrom::Start(from))?;
    let mut reader = BufReader::with_capacity(64 * 1024, file);
    let mut offset = from;
    let mut line = Vec::new();

    // Unless we happen to start at byte 0, the first line is a partial one.
    if from > 0 {
        let read = read_line(&mut reader, &mut line)?;
        if read == 0 {
            return Ok(None);
        }
        offset += read as u64;
    }

    match format {
        Format::Fasta => loop {
            let read = read_line(&mut reader, &mut line)?;
            if read == 0 {
                return Ok(None);
            }
            if line.first() == Some(&b'>') {
                return Ok(Some(offset));
            }
            offset += read as u64;
        },
        Format::Fastq => {
            // Sliding window of four consecutive lines.
            let mut window: Vec<(u64, Vec<u8>)> = Vec::with_capacity(4);
            loop {
                let read = read_line(&mut reader, &mut line)?;
                if read == 0 {
                    return Ok(None);
                }
                if window.len() == 4 {
                    window.remove(0);
                }
                window.push((offset, trim_newline(&line).to_vec()));
                offset += read as u64;
                if window.len() == 4
                    && window[0].1.first() == Some(&b'@')
                    && window[2].1.first() == Some(&b'+')
                    && window[1].1.len() == window[3].1.len()
                {
                    return Ok(Some(window[0].0));
                }
            }
        }
    }
}

fn read_line<R: Read>(reader: &mut BufReader<R>, line: &mut Vec<u8>) -> Result<usize> {
    use std::io::BufRead;
    line.clear();
    Ok(reader.read_until(b'\n', line)?)
}

fn trim_newline(line: &[u8]) -> &[u8] {
    let mut end = line.len();
    if end > 0 && line[end - 1] == b'\n' {
        end -= 1;
    }
    if end > 0 && line[end - 1] == b'\r' {
        end -= 1;
    }
    &line[..end]
}

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

    fn temp_dir(tag: &str) -> std::path::PathBuf {
        let dir = std::env::temp_dir().join(format!("fastx-par-{}-{tag}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        dir
    }

    fn write_fastq(path: &Path, records: usize) {
        let mut file = std::fs::File::create(path).unwrap();
        for i in 0..records {
            // Quality deliberately contains '@' and '+' to confuse naive splitters.
            writeln!(file, "@read{i} sample\nACGTACGTAC\n+\n@@++IIIIII").unwrap();
        }
    }

    fn write_fasta(path: &Path, records: usize) {
        let mut file = std::fs::File::create(path).unwrap();
        for i in 0..records {
            writeln!(file, ">contig{i}\nACGTACGTAC\nGGGG").unwrap();
        }
    }

    #[test]
    fn parallel_batches_preserve_order() {
        let data = b">a\nAC\n>b\nACGT\n>c\nACGTAC\n";
        let lengths = par_map(&mut FastxReader::new(&data[..]), 2, |r| r.len()).unwrap();
        assert_eq!(lengths, vec![2, 4, 6]);
    }

    #[test]
    fn parallel_fold_counts_bases() {
        let data = b"@a\nACGT\n+\nIIII\n@b\nAC\n+\nII\n";
        let total = par_fold(
            &mut FastxReader::new(&data[..]),
            1,
            || 0usize,
            |acc, r| acc + r.len(),
            |a, b| a + b,
        )
        .unwrap();
        assert_eq!(total, 6);
    }

    #[test]
    fn parallel_for_each_propagates_errors() {
        let data = b">a\nACGT\n";
        let err = par_for_each(&mut FastxReader::new(&data[..]), 4, |_| {
            Err(Error::Index("boom".into()))
        });
        assert!(err.is_err());
    }

    #[test]
    fn parallel_stats_match_sequential() {
        let data = b">a\nACGTACGTAC\n>b\nGGCC\n>c\nNNNN\n";
        let parallel = par_stats(&mut FastxReader::new(&data[..]), 2).unwrap();
        let mut sequential = SeqStats::new();
        FastxReader::new(&data[..])
            .for_each_record(|r| {
                sequential.push(r);
                Ok(())
            })
            .unwrap();
        assert_eq!(parallel.count, sequential.count);
        assert_eq!(parallel.total_length, sequential.total_length);
        assert_eq!(parallel.n50(), sequential.n50());
        assert_eq!(parallel.gc_content(), sequential.gc_content());
    }

    #[test]
    fn splits_fastq_on_record_boundaries() {
        let dir = temp_dir("fq");
        let path = dir.join("reads.fq");
        write_fastq(&path, 500);

        for parts in [1, 2, 3, 7, 64, 4096] {
            let ranges = split_file(&path, parts).unwrap();
            assert_eq!(ranges.first().unwrap().start, 0);
            assert_eq!(
                ranges.last().unwrap().end,
                std::fs::metadata(&path).unwrap().len()
            );
            assert!(
                ranges.windows(2).all(|w| w[0].end == w[1].start),
                "{ranges:?}"
            );

            let counts =
                par_map_chunks_file(&path, parts, |reader| reader.count_records()).unwrap();
            assert_eq!(counts.iter().sum::<u64>(), 500, "parts={parts}");
        }
        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn splits_fasta_on_record_boundaries() {
        let dir = temp_dir("fa");
        let path = dir.join("contigs.fa");
        write_fasta(&path, 300);

        for parts in [1, 4, 9, 128] {
            let ids = par_map_chunks_file(&path, parts, |reader| {
                let mut ids = Vec::new();
                reader.for_each_record(|r| {
                    ids.push(r.id.clone());
                    Ok(())
                })?;
                Ok(ids)
            })
            .unwrap();
            let flat: Vec<String> = ids.into_iter().flatten().collect();
            assert_eq!(flat.len(), 300, "parts={parts}");
            // Chunks come back in file order, so the ids do too.
            assert_eq!(flat[0], "contig0");
            assert_eq!(flat[299], "contig299");
        }

        let stats = par_stats_file(&path, 8).unwrap();
        assert_eq!(stats.count, 300);
        assert_eq!(stats.total_length, 300 * 14);
        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn parallel_file_processing_sees_every_record() {
        let dir = temp_dir("each");
        let path = dir.join("reads.fq");
        write_fastq(&path, 200);
        let count = std::sync::atomic::AtomicUsize::new(0);
        par_for_each_file(&path, 8, |_| {
            count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            Ok(())
        })
        .unwrap();
        assert_eq!(count.load(std::sync::atomic::Ordering::Relaxed), 200);
        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn empty_file_splits_into_nothing() {
        let dir = temp_dir("empty");
        let path = dir.join("empty.fa");
        std::fs::write(&path, b"").unwrap();
        assert!(split_file(&path, 4).unwrap().is_empty());
        assert!(par_map_chunks_file(&path, 4, |r| r.count_records())
            .unwrap()
            .is_empty());
        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn refuses_to_split_gzip() {
        assert!(split_file("x.fq.gz", 4).is_err());
    }
}