rsomics_seq_grep/
pattern.rs

1use std::collections::HashSet;
2use std::fs;
3use std::path::Path;
4
5use regex::bytes::Regex;
6use rsomics_common::{Result, RsomicsError};
7
8/// One compiled matcher per `-p`/`-f` pattern set. The three modes mirror
9/// `seqkit grep`'s three matching strategies, which are mutually exclusive
10/// per invocation (chosen once from `-r`/`-d`/`-s`, not per pattern).
11pub enum PatternSet {
12    /// `-p`/`-f` without `-r`/`-d`, against ID or full name: whole-string equality.
13    Exact(HashSet<Vec<u8>>),
14    /// `-s` without `-r`/`-d`: substring search on the sequence.
15    Substring(Vec<Vec<u8>>),
16    /// `-r` (any target) or `-d` (sequence only, IUPAC-expanded first): unanchored regex.
17    Regex(Vec<Regex>),
18}
19
20impl PatternSet {
21    pub fn is_match(&self, target: &[u8], ignore_case: bool) -> bool {
22        match self {
23            PatternSet::Exact(set) => {
24                if ignore_case {
25                    set.contains(&target.to_ascii_lowercase())
26                } else {
27                    set.contains(target)
28                }
29            }
30            PatternSet::Substring(pats) => {
31                let lowered;
32                let t: &[u8] = if ignore_case {
33                    lowered = target.to_ascii_lowercase();
34                    &lowered
35                } else {
36                    target
37                };
38                pats.iter().any(|p| {
39                    if p.is_empty() {
40                        t.is_empty()
41                    } else {
42                        memchr::memmem::find(t, p).is_some()
43                    }
44                })
45            }
46            PatternSet::Regex(res) => res.iter().any(|re| re.is_match(target)),
47        }
48    }
49}
50
51pub struct BuildOpts {
52    pub use_regexp: bool,
53    pub degenerate: bool,
54    pub by_seq: bool,
55    pub ignore_case: bool,
56}
57
58/// Build the matcher for one grep invocation. Precedence matches upstream:
59/// `-d` and `-r` both compile to regex (an empty pattern becomes `^$`, an
60/// exact-empty-string match, rather than the "matches everywhere" empty
61/// regex); otherwise sequence search is substring, ID/name search is exact.
62pub fn build(patterns: &[String], opts: &BuildOpts) -> Result<PatternSet> {
63    if opts.use_regexp && opts.degenerate {
64        return Err(RsomicsError::InvalidInput(
65            "cannot give both -d/--degenerate and -r/--use-regexp".into(),
66        ));
67    }
68
69    if opts.use_regexp || opts.degenerate {
70        let mut res = Vec::with_capacity(patterns.len());
71        for p in patterns {
72            let expanded = if opts.degenerate {
73                degenerate_to_regex(p)
74            } else if p.is_empty() {
75                "^$".to_string()
76            } else {
77                p.clone()
78            };
79            let source = if opts.ignore_case {
80                format!("(?i){expanded}")
81            } else {
82                expanded
83            };
84            let re = Regex::new(&source)
85                .map_err(|e| RsomicsError::InvalidInput(format!("invalid regex {p:?}: {e}")))?;
86            res.push(re);
87        }
88        return Ok(PatternSet::Regex(res));
89    }
90
91    if opts.by_seq {
92        for p in patterns {
93            if !is_legal_seq_pattern(p.as_bytes()) {
94                return Err(RsomicsError::InvalidInput(format!(
95                    "illegal DNA/RNA/Protein sequence: {p}"
96                )));
97            }
98        }
99        let subs = patterns
100            .iter()
101            .map(|p| {
102                let bytes = p.as_bytes().to_vec();
103                if opts.ignore_case {
104                    bytes.to_ascii_lowercase()
105                } else {
106                    bytes
107                }
108            })
109            .collect();
110        return Ok(PatternSet::Substring(subs));
111    }
112
113    let set = patterns
114        .iter()
115        .map(|p| {
116            let bytes = p.as_bytes().to_vec();
117            if opts.ignore_case {
118                bytes.to_ascii_lowercase()
119            } else {
120                bytes
121            }
122        })
123        .collect();
124    Ok(PatternSet::Exact(set))
125}
126
127const fn alphabet_set(bytes: &[u8]) -> [bool; 256] {
128    let mut set = [false; 256];
129    let mut i = 0;
130    while i < bytes.len() {
131        set[bytes[i] as usize] = true;
132        i += 1;
133    }
134    set
135}
136
137// The valid-byte set of each seqkit alphabet (bio/seq/alphabet.go): letters +
138// gap (` -.`) + ambiguous codes, all case variants listed explicitly.
139const DNA_REDUNDANT: [bool; 256] = alphabet_set(b"acgtryswkmbdhvACGTRYSWKMBDHV -.nN.");
140const RNA_REDUNDANT: [bool; 256] = alphabet_set(b"acguryswkmbdhvACGURYSWKMBDHV -.nN");
141const PROTEIN: [bool; 256] =
142    alphabet_set(b"abcdefghijklmnopqrstuvwyzABCDEFGHIJKLMNOPQRSTUVWYZ -xX*_.");
143
144/// A `-s`/`-f` literal pattern is legal only if every byte is valid in one of
145/// seqkit's DNAredundant / RNAredundant / Protein alphabets — mirroring
146/// `grep.go`'s `IsValid` triple-check, which rejects e.g. a digit in `-s -p T1A`.
147/// An empty pattern is vacuously legal, matching seqkit.
148fn is_legal_seq_pattern(p: &[u8]) -> bool {
149    p.iter().all(|&b| DNA_REDUNDANT[b as usize])
150        || p.iter().all(|&b| RNA_REDUNDANT[b as usize])
151        || p.iter().all(|&b| PROTEIN[b as usize])
152}
153
154/// `seqkit`'s `Degenerate2Regexp` over `DegenerateBaseMapNucl` (`bio/seq/seq.go`):
155/// each IUPAC nucleotide letter expands to the regex class of the bases it stands
156/// for; every other byte — a plain A/C/G/T/U or a non-IUPAC char — passes through
157/// literally, so `-d -p 'RY9SW'` runs (the `9` becomes a literal) rather than
158/// erroring. Protein degenerate codes are not covered.
159fn degenerate_to_regex(p: &str) -> String {
160    if p.is_empty() {
161        return "^$".to_string();
162    }
163    let mut out = String::with_capacity(p.len() * 2);
164    for b in p.bytes() {
165        match b {
166            b'A' => out.push('A'),
167            b'T' | b'U' => out.push_str("[TU]"),
168            b'C' => out.push('C'),
169            b'G' => out.push('G'),
170            b'R' => out.push_str("[AG]"),
171            b'Y' => out.push_str("[CTU]"),
172            b'M' => out.push_str("[AC]"),
173            b'K' => out.push_str("[GTU]"),
174            b'S' => out.push_str("[CG]"),
175            b'W' => out.push_str("[ATU]"),
176            b'H' => out.push_str("[ACTU]"),
177            b'B' => out.push_str("[CGTU]"),
178            b'V' => out.push_str("[ACG]"),
179            b'D' => out.push_str("[AGTU]"),
180            b'N' => out.push_str("[ACGTU]"),
181            b'a' => out.push('a'),
182            b't' | b'u' => out.push_str("[tu]"),
183            b'c' => out.push('c'),
184            b'g' => out.push('g'),
185            b'r' => out.push_str("[ag]"),
186            b'y' => out.push_str("[ctu]"),
187            b'm' => out.push_str("[ac]"),
188            b'k' => out.push_str("[gtu]"),
189            b's' => out.push_str("[cg]"),
190            b'w' => out.push_str("[atu]"),
191            b'h' => out.push_str("[actu]"),
192            b'b' => out.push_str("[cgtu]"),
193            b'v' => out.push_str("[acg]"),
194            b'd' => out.push_str("[agtu]"),
195            b'n' => out.push_str("[acgtu]"),
196            other => out.push(other as char),
197        }
198    }
199    out
200}
201
202/// Load patterns from a plain-text file, one per line. Trailing `\r\n`/`\n`
203/// is stripped; blank lines are skipped (matches `seqkit`'s `breader`
204/// default line reader plus its own empty-line skip in `grep.go`).
205pub fn load_pattern_file(path: &Path) -> Result<Vec<String>> {
206    let raw = fs::read(path)
207        .map_err(|e| RsomicsError::InvalidInput(format!("{}: {e}", path.display())))?;
208    let text = String::from_utf8(raw)
209        .map_err(|e| RsomicsError::InvalidInput(format!("{}: {e}", path.display())))?;
210    Ok(text
211        .split('\n')
212        .map(|line| line.strip_suffix('\r').unwrap_or(line))
213        .filter(|line| !line.is_empty())
214        .map(str::to_owned)
215        .collect())
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221
222    fn opts(use_regexp: bool, degenerate: bool, by_seq: bool, ignore_case: bool) -> BuildOpts {
223        BuildOpts {
224            use_regexp,
225            degenerate,
226            by_seq,
227            ignore_case,
228        }
229    }
230
231    #[test]
232    fn exact_match_is_whole_string() {
233        let set = build(&["seq1".to_string()], &opts(false, false, false, false)).unwrap();
234        assert!(set.is_match(b"seq1", false));
235        assert!(!set.is_match(b"seq1_dup", false));
236    }
237
238    #[test]
239    fn substring_matches_partial() {
240        let set = build(&["TTTT".to_string()], &opts(false, false, true, false)).unwrap();
241        assert!(set.is_match(b"AATTTTAA", false));
242        assert!(!set.is_match(b"AACCCCAA", false));
243    }
244
245    #[test]
246    fn empty_substring_pattern_only_matches_empty_target() {
247        let set = build(&[String::new()], &opts(false, false, true, false)).unwrap();
248        assert!(set.is_match(b"", false));
249        assert!(!set.is_match(b"ACGT", false));
250    }
251
252    #[test]
253    fn degenerate_expands_iupac() {
254        let set = build(&["RYSW".to_string()], &opts(false, true, true, false)).unwrap();
255        assert!(set.is_match(b"ACGT", false));
256        assert!(!set.is_match(b"CCCC", false));
257    }
258
259    #[test]
260    fn degenerate_passes_non_iupac_through() {
261        // A non-IUPAC byte becomes a literal regex char (no error).
262        let set = build(&["RY9SW".to_string()], &opts(false, true, true, false)).unwrap();
263        assert!(set.is_match(b"AC9CA", false));
264        assert!(!set.is_match(b"ACGCA", false));
265    }
266
267    #[test]
268    fn by_seq_rejects_illegal_alphabet() {
269        assert!(build(&["T1A".to_string()], &opts(false, false, true, false)).is_err());
270    }
271
272    #[test]
273    fn by_seq_accepts_iupac_and_protein() {
274        for p in ["ATGC", "RYSWKMBDHVN", "MEK", "AT.GC", ""] {
275            assert!(
276                build(&[p.to_string()], &opts(false, false, true, false)).is_ok(),
277                "should accept {p:?}"
278            );
279        }
280    }
281
282    #[test]
283    fn regex_is_unanchored() {
284        let set = build(&["seq1".to_string()], &opts(true, false, false, false)).unwrap();
285        assert!(set.is_match(b"seq1_dup", false));
286    }
287
288    #[test]
289    fn ignore_case_lowercases_both_sides() {
290        let set = build(&["SEQ1".to_string()], &opts(false, false, false, true)).unwrap();
291        assert!(set.is_match(b"seq1", true));
292    }
293
294    #[test]
295    fn empty_regex_pattern_is_anchored_empty() {
296        let set = build(&[String::new()], &opts(true, false, false, false)).unwrap();
297        assert!(set.is_match(b"", false));
298        assert!(!set.is_match(b"anything", false));
299    }
300}
rsomics_seq_grep/pattern.rs

rsomics_seq_grep/
pattern.rs
