rsomics_seq_grep/
lib.rs

1mod matcher;
2mod pattern;
3mod region;
4mod revcomp;
5mod writer;
6
7use std::io::{BufWriter, Write};
8
9use needletail::parser::FastxReader;
10use rayon::prelude::*;
11use rsomics_common::{Result, RsomicsError};
12
13pub use pattern::load_pattern_file;
14pub use region::Region;
15
16use matcher::{MatchOpts, is_hit, split_id};
17use pattern::{BuildOpts, PatternSet};
18
19pub struct GrepOptions {
20    pub patterns: Vec<String>,
21    pub by_name: bool,
22    pub use_regexp: bool,
23    pub ignore_case: bool,
24    pub by_seq: bool,
25    pub degenerate: bool,
26    pub invert_match: bool,
27    pub region: Option<Region>,
28    pub only_positive: bool,
29    pub line_width: usize,
30    pub count_only: bool,
31}
32
33pub struct GrepStats {
34    pub total_records: u64,
35    pub matched: u64,
36}
37
38/// Filter FASTA/FASTQ records from `input` (a file path, or `"-"` for
39/// stdin) into `out`. Sequence search (`-s`/`-d`/`-R`) buffers records and
40/// matches them across a rayon pool since revcomp + regex/substring search
41/// scale with sequence length; ID/name search stays a single streaming
42/// pass — it's O(1) per record and buffering would only add overhead.
43pub fn grep(input: &str, opts: &GrepOptions, out: &mut dyn Write) -> Result<GrepStats> {
44    let effective_by_seq =
45        MatchOpts::effective_by_seq(opts.by_seq, opts.degenerate, opts.region.is_some());
46
47    let patterns = pattern::build(
48        &opts.patterns,
49        &BuildOpts {
50            use_regexp: opts.use_regexp,
51            degenerate: opts.degenerate,
52            by_seq: effective_by_seq,
53            ignore_case: opts.ignore_case,
54        },
55    )?;
56
57    let match_opts = MatchOpts {
58        by_name: opts.by_name,
59        by_seq: effective_by_seq,
60        only_positive: opts.only_positive,
61        ignore_case: opts.ignore_case,
62        region: opts.region,
63    };
64
65    let reader = open(input)?;
66
67    if effective_by_seq {
68        run_parallel(reader, &match_opts, &patterns, opts, out)
69    } else {
70        run_streaming(reader, &match_opts, &patterns, opts, out)
71    }
72}
73
74fn open(input: &str) -> Result<Box<dyn FastxReader>> {
75    if input == "-" {
76        needletail::parse_fastx_stdin()
77            .map_err(|e| RsomicsError::InvalidInput(format!("stdin: {e}")))
78    } else {
79        needletail::parse_fastx_file(input)
80            .map_err(|e| RsomicsError::InvalidInput(format!("{input}: {e}")))
81    }
82}
83
84fn run_streaming(
85    mut reader: Box<dyn FastxReader>,
86    match_opts: &MatchOpts,
87    patterns: &PatternSet,
88    opts: &GrepOptions,
89    out: &mut dyn Write,
90) -> Result<GrepStats> {
91    let mut writer = BufWriter::with_capacity(1 << 20, out);
92    let mut total = 0u64;
93    let mut matched = 0u64;
94
95    while let Some(rec) = reader.next() {
96        let rec = rec.map_err(|e| RsomicsError::InvalidInput(format!("reading record: {e}")))?;
97        total += 1;
98
99        let name = rec.id();
100        let id = split_id(name);
101        let hit = is_hit(id, name, &[], match_opts, patterns);
102        if hit == opts.invert_match {
103            continue;
104        }
105        matched += 1;
106        if opts.count_only {
107            continue;
108        }
109
110        let seq = rec.seq();
111        if let Some(qual) = rec.qual() {
112            writer::write_fastq(name, &seq, qual, &mut writer)?;
113        } else {
114            writer::write_fasta(name, &seq, opts.line_width, &mut writer)?;
115        }
116    }
117
118    if opts.count_only {
119        writeln!(writer, "{matched}").map_err(RsomicsError::Io)?;
120    }
121    writer.flush().map_err(RsomicsError::Io)?;
122    Ok(GrepStats {
123        total_records: total,
124        matched,
125    })
126}
127
128struct OwnedRecord {
129    name: Vec<u8>,
130    seq: Vec<u8>,
131    qual: Option<Vec<u8>>,
132}
133
134fn run_parallel(
135    mut reader: Box<dyn FastxReader>,
136    match_opts: &MatchOpts,
137    patterns: &PatternSet,
138    opts: &GrepOptions,
139    out: &mut dyn Write,
140) -> Result<GrepStats> {
141    let mut records = Vec::new();
142    while let Some(rec) = reader.next() {
143        let rec = rec.map_err(|e| RsomicsError::InvalidInput(format!("reading record: {e}")))?;
144        records.push(OwnedRecord {
145            name: rec.id().to_vec(),
146            seq: rec.seq().into_owned(),
147            qual: rec.qual().map(<[u8]>::to_vec),
148        });
149    }
150    let total = records.len() as u64;
151
152    let results: Vec<(bool, Vec<u8>)> = records
153        .par_iter()
154        .map(|r| {
155            let id = split_id(&r.name);
156            let hit = is_hit(id, &r.name, &r.seq, match_opts, patterns);
157            let keep = hit != opts.invert_match;
158            let mut buf = Vec::new();
159            if keep && !opts.count_only {
160                if let Some(qual) = &r.qual {
161                    writer::format_fastq(&r.name, &r.seq, qual, &mut buf);
162                } else {
163                    writer::format_fasta(&r.name, &r.seq, opts.line_width, &mut buf);
164                }
165            }
166            (keep, buf)
167        })
168        .collect();
169
170    let mut writer = BufWriter::with_capacity(1 << 20, out);
171    let mut matched = 0u64;
172    for (keep, buf) in results {
173        if !keep {
174            continue;
175        }
176        matched += 1;
177        if !opts.count_only {
178            writer.write_all(&buf).map_err(RsomicsError::Io)?;
179        }
180    }
181    if opts.count_only {
182        writeln!(writer, "{matched}").map_err(RsomicsError::Io)?;
183    }
184    writer.flush().map_err(RsomicsError::Io)?;
185    Ok(GrepStats {
186        total_records: total,
187        matched,
188    })
189}
rsomics_seq_grep/lib.rs

rsomics_seq_grep/
lib.rs
