rsomics-seq-grep 0.1.1

mod matcher;
mod pattern;
mod region;
mod revcomp;
mod writer;

use std::io::{BufWriter, Write};

use needletail::parser::FastxReader;
use rayon::prelude::*;
use rsomics_common::{Result, RsomicsError};

pub use pattern::load_pattern_file;
pub use region::Region;

use matcher::{MatchOpts, is_hit, split_id};
use pattern::{BuildOpts, PatternSet};

pub struct GrepOptions {
    pub patterns: Vec<String>,
    pub by_name: bool,
    pub use_regexp: bool,
    pub ignore_case: bool,
    pub by_seq: bool,
    pub degenerate: bool,
    pub invert_match: bool,
    pub region: Option<Region>,
    pub only_positive: bool,
    pub line_width: usize,
    pub count_only: bool,
}

pub struct GrepStats {
    pub total_records: u64,
    pub matched: u64,
}

/// Filter FASTA/FASTQ records from `input` (a file path, or `"-"` for
/// stdin) into `out`. Sequence search (`-s`/`-d`/`-R`) buffers records and
/// matches them across a rayon pool since revcomp + regex/substring search
/// scale with sequence length; ID/name search stays a single streaming
/// pass — it's O(1) per record and buffering would only add overhead.
pub fn grep(input: &str, opts: &GrepOptions, out: &mut dyn Write) -> Result<GrepStats> {
    let effective_by_seq =
        MatchOpts::effective_by_seq(opts.by_seq, opts.degenerate, opts.region.is_some());

    let patterns = pattern::build(
        &opts.patterns,
        &BuildOpts {
            use_regexp: opts.use_regexp,
            degenerate: opts.degenerate,
            by_seq: effective_by_seq,
            ignore_case: opts.ignore_case,
        },
    )?;

    let match_opts = MatchOpts {
        by_name: opts.by_name,
        by_seq: effective_by_seq,
        only_positive: opts.only_positive,
        ignore_case: opts.ignore_case,
        region: opts.region,
    };

    let reader = open(input)?;

    if effective_by_seq {
        run_parallel(reader, &match_opts, &patterns, opts, out)
    } else {
        run_streaming(reader, &match_opts, &patterns, opts, out)
    }
}

fn open(input: &str) -> Result<Box<dyn FastxReader>> {
    if input == "-" {
        needletail::parse_fastx_stdin()
            .map_err(|e| RsomicsError::InvalidInput(format!("stdin: {e}")))
    } else {
        needletail::parse_fastx_file(input)
            .map_err(|e| RsomicsError::InvalidInput(format!("{input}: {e}")))
    }
}

fn run_streaming(
    mut reader: Box<dyn FastxReader>,
    match_opts: &MatchOpts,
    patterns: &PatternSet,
    opts: &GrepOptions,
    out: &mut dyn Write,
) -> Result<GrepStats> {
    let mut writer = BufWriter::with_capacity(1 << 20, out);
    let mut total = 0u64;
    let mut matched = 0u64;

    while let Some(rec) = reader.next() {
        let rec = rec.map_err(|e| RsomicsError::InvalidInput(format!("reading record: {e}")))?;
        total += 1;

        let name = rec.id();
        let id = split_id(name);
        let hit = is_hit(id, name, &[], match_opts, patterns);
        if hit == opts.invert_match {
            continue;
        }
        matched += 1;
        if opts.count_only {
            continue;
        }

        let seq = rec.seq();
        if let Some(qual) = rec.qual() {
            writer::write_fastq(name, &seq, qual, &mut writer)?;
        } else {
            writer::write_fasta(name, &seq, opts.line_width, &mut writer)?;
        }
    }

    if opts.count_only {
        writeln!(writer, "{matched}").map_err(RsomicsError::Io)?;
    }
    writer.flush().map_err(RsomicsError::Io)?;
    Ok(GrepStats {
        total_records: total,
        matched,
    })
}

struct OwnedRecord {
    name: Vec<u8>,
    seq: Vec<u8>,
    qual: Option<Vec<u8>>,
}

fn run_parallel(
    mut reader: Box<dyn FastxReader>,
    match_opts: &MatchOpts,
    patterns: &PatternSet,
    opts: &GrepOptions,
    out: &mut dyn Write,
) -> Result<GrepStats> {
    let mut records = Vec::new();
    while let Some(rec) = reader.next() {
        let rec = rec.map_err(|e| RsomicsError::InvalidInput(format!("reading record: {e}")))?;
        records.push(OwnedRecord {
            name: rec.id().to_vec(),
            seq: rec.seq().into_owned(),
            qual: rec.qual().map(<[u8]>::to_vec),
        });
    }
    let total = records.len() as u64;

    let results: Vec<(bool, Vec<u8>)> = records
        .par_iter()
        .map(|r| {
            let id = split_id(&r.name);
            let hit = is_hit(id, &r.name, &r.seq, match_opts, patterns);
            let keep = hit != opts.invert_match;
            let mut buf = Vec::new();
            if keep && !opts.count_only {
                if let Some(qual) = &r.qual {
                    writer::format_fastq(&r.name, &r.seq, qual, &mut buf);
                } else {
                    writer::format_fasta(&r.name, &r.seq, opts.line_width, &mut buf);
                }
            }
            (keep, buf)
        })
        .collect();

    let mut writer = BufWriter::with_capacity(1 << 20, out);
    let mut matched = 0u64;
    for (keep, buf) in results {
        if !keep {
            continue;
        }
        matched += 1;
        if !opts.count_only {
            writer.write_all(&buf).map_err(RsomicsError::Io)?;
        }
    }
    if opts.count_only {
        writeln!(writer, "{matched}").map_err(RsomicsError::Io)?;
    }
    writer.flush().map_err(RsomicsError::Io)?;
    Ok(GrepStats {
        total_records: total,
        matched,
    })
}