sensitive_rs/engine/
mod.rs

1//! Multi-pattern matching engine.
2//!
3//! [`MultiPatternEngine`] automatically selects an algorithm based on vocabulary size:
4//!
5//! | Patterns | Algorithm | Why |
6//! |----------|-----------|-----|
7//! | 0–100    | [`MatchAlgorithm::WuManber`]   | Small tables, quick scan |
8//! | 101–10k  | [`MatchAlgorithm::AhoCorasick`]| O(n) automaton scan regardless of count |
9//! | 10k+     | [`MatchAlgorithm::Regex`]      | Compilation overhead amortized over many patterns |
10//!
11//! Use [`MultiPatternEngine::recommend_algorithm`] to preview the choice, or force one with
12//! [`MultiPatternEngine::rebuild_with_algorithm`].
13
14pub mod wumanber;
15use crate::WuManber;
16use aho_corasick::{AhoCorasick, AhoCorasickBuilder};
17use regex::Regex;
18use std::sync::Arc;
19
20/// Supported matching algorithm types
21///
22/// Implements [`Display`](std::fmt::Display) for human-readable names.
23///
24/// # Examples
25///
26/// ```
27/// use sensitive_rs::MatchAlgorithm;
28///
29/// assert_eq!(MatchAlgorithm::AhoCorasick.to_string(), "Aho-Corasick");
30/// assert_eq!(MatchAlgorithm::WuManber.to_string(), "Wu-Manber");
31/// assert_eq!(MatchAlgorithm::Regex.to_string(), "Regex");
32/// ```
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum MatchAlgorithm {
35    /// Best for medium-sized vocabulary (101-10,000 patterns)
36    /// Automaton-based, O(n) scan regardless of pattern count
37    AhoCorasick,
38    /// Best for small vocabulary (0-100 patterns)
39    /// Fast with few patterns: small tables, quick scan
40    WuManber,
41    /// Best for very large vocabulary (10,000+ patterns)
42    /// Pattern compilation overhead amortized over many patterns
43    Regex,
44}
45
46impl std::fmt::Display for MatchAlgorithm {
47    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        match self {
49            Self::AhoCorasick => write!(f, "Aho-Corasick"),
50            Self::WuManber => write!(f, "Wu-Manber"),
51            Self::Regex => write!(f, "Regex"),
52        }
53    }
54}
55
56/// Multi-pattern matching engine
57pub struct MultiPatternEngine {
58    algorithm: MatchAlgorithm,    // The matching algorithm currently used
59    ac: Option<Arc<AhoCorasick>>, // Aho-Corasick Engine
60    wm: Option<Arc<WuManber>>,    // Wu-Manber Engine
61    regex_set: Option<Regex>,     // Regular Expression Engine
62    patterns: Vec<String>,        // Store all modes
63}
64
65impl std::fmt::Debug for MultiPatternEngine {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        f.debug_struct("MultiPatternEngine")
68            .field("algorithm", &self.algorithm)
69            .field("pattern_count", &self.patterns.len())
70            .field("has_ac", &self.ac.is_some())
71            .field("has_wm", &self.wm.is_some())
72            .field("has_regex", &self.regex_set.is_some())
73            .finish()
74    }
75}
76
77impl Default for MultiPatternEngine {
78    fn default() -> Self {
79        Self { algorithm: MatchAlgorithm::AhoCorasick, ac: None, wm: None, regex_set: None, patterns: Vec::new() }
80    }
81}
82
83impl MultiPatternEngine {
84    /// Create a new engine and automatically select the algorithm based on the lexicon size
85    ///
86    /// Pass `None` for `algorithm` to auto-select; pass [`Some`] to force a specific algorithm.
87    ///
88    /// # Examples
89    ///
90    /// ```
91    /// use sensitive_rs::MultiPatternEngine;
92    ///
93    /// let patterns = vec!["赌博".to_string(), "色情".to_string()];
94    /// let engine = MultiPatternEngine::new(None, &patterns);
95    /// assert_eq!(engine.find_first("含有赌博"), Some("赌博".to_string()));
96    /// ```
97    pub fn new(algorithm: Option<MatchAlgorithm>, patterns: &[String]) -> Self {
98        let algorithm = algorithm.unwrap_or_else(|| Self::recommend_algorithm(patterns.len()));
99        let mut engine = Self { algorithm, ..Default::default() };
100
101        engine.rebuild(patterns);
102        engine
103    }
104
105    /// Rebuild the engine (called when the pattern is updated)
106    pub fn rebuild(&mut self, patterns: &[String]) {
107        self.patterns = patterns.to_vec();
108
109        // Reevaluate algorithm selection based on new thesaurus size
110        let recommended = Self::recommend_algorithm(patterns.len());
111        if self.algorithm != recommended {
112            self.algorithm = recommended;
113        }
114
115        self.build_engines();
116    }
117
118    /// Recommended algorithm based on the lexicon size
119    ///
120    /// - 0-100 patterns: WuManber (few patterns = small tables, quick scan)
121    /// - 101-10,000 patterns: AhoCorasick (automaton-based, O(n) scan)
122    /// - 10,000+ patterns: Regex (compilation overhead amortized)
123    pub fn recommend_algorithm(word_count: usize) -> MatchAlgorithm {
124        match word_count {
125            0..=100 => MatchAlgorithm::WuManber,
126            101..=10_000 => MatchAlgorithm::AhoCorasick,
127            _ => MatchAlgorithm::Regex,
128        }
129    }
130
131    /// Force rebuild using the specified algorithm
132    pub fn rebuild_with_algorithm(&mut self, patterns: &[String], algorithm: MatchAlgorithm) {
133        self.patterns = patterns.to_vec();
134        self.algorithm = algorithm;
135        self.build_engines();
136    }
137
138    /// Build the corresponding engine according to the current algorithm
139    fn build_engines(&mut self) {
140        // Clear all engines
141        self.ac = None;
142        self.wm = None;
143        self.regex_set = None;
144
145        // Build the corresponding engine according to the selected algorithm
146        match self.algorithm {
147            MatchAlgorithm::AhoCorasick => {
148                if !self.patterns.is_empty() {
149                    match AhoCorasickBuilder::new()
150                        .match_kind(aho_corasick::MatchKind::LeftmostLongest)
151                        .build(&self.patterns)
152                    {
153                        Ok(ac) => self.ac = Some(Arc::new(ac)),
154                        Err(_) => {
155                            // Fallback to WuManber if AhoCorasick build fails
156                            self.algorithm = MatchAlgorithm::WuManber;
157                            self.wm = Some(Arc::new(WuManber::new_chinese(self.patterns.clone())));
158                        }
159                    }
160                }
161            }
162            MatchAlgorithm::WuManber => {
163                if !self.patterns.is_empty() {
164                    self.wm = Some(Arc::new(WuManber::new_chinese(self.patterns.clone())));
165                }
166            }
167            MatchAlgorithm::Regex => {
168                if !self.patterns.is_empty() {
169                    let escaped_patterns: Vec<String> = self.patterns.iter().map(|p| regex::escape(p)).collect();
170                    let pattern = escaped_patterns.join("|");
171
172                    match Regex::new(&pattern) {
173                        Ok(regex) => self.regex_set = Some(regex),
174                        Err(_) => {
175                            // Fallback to WuManber if Regex build fails
176                            self.algorithm = MatchAlgorithm::WuManber;
177                            self.wm = Some(Arc::new(WuManber::new_chinese(self.patterns.clone())));
178                        }
179                    }
180                }
181            }
182        }
183    }
184
185    /// Get the currently used algorithm
186    pub fn current_algorithm(&self) -> MatchAlgorithm {
187        self.algorithm
188    }
189
190    /// Get all modes
191    pub fn get_patterns(&self) -> &[String] {
192        &self.patterns
193    }
194
195    /// Find the first match
196    ///
197    /// # Examples
198    ///
199    /// ```
200    /// use sensitive_rs::MultiPatternEngine;
201    ///
202    /// let patterns = vec!["赌博".to_string()];
203    /// let engine = MultiPatternEngine::new(None, &patterns);
204    /// assert_eq!(engine.find_first("含有赌博内容"), Some("赌博".to_string()));
205    /// assert_eq!(engine.find_first("正常文本"), None);
206    /// ```
207    pub fn find_first(&self, text: &str) -> Option<String> {
208        match self.algorithm {
209            MatchAlgorithm::AhoCorasick => {
210                self.ac.as_ref()?.find(text).map(|mat| text[mat.start()..mat.end()].to_string())
211            }
212            MatchAlgorithm::WuManber => {
213                // Use the search_string method to return directly to String
214                self.wm.as_ref()?.search_string(text)
215            }
216            MatchAlgorithm::Regex => self.regex_set.as_ref()?.find(text).map(|mat| mat.as_str().to_string()),
217        }
218    }
219
220    /// Replace all matches with optimized performance
221    pub fn replace_all(&self, text: &str, replacement: &str) -> String {
222        match self.algorithm {
223            MatchAlgorithm::AhoCorasick => {
224                if let Some(ac) = &self.ac {
225                    ac.replace_all(text, &[replacement]).to_string()
226                } else {
227                    text.to_string()
228                }
229            }
230            MatchAlgorithm::WuManber => {
231                if let Some(wm) = &self.wm {
232                    if replacement.is_empty() {
233                        wm.remove_all(text)
234                    } else {
235                        let repl_char = replacement.chars().next().unwrap_or('*');
236                        wm.replace_all(text, repl_char)
237                    }
238                } else {
239                    text.to_string()
240                }
241            }
242            MatchAlgorithm::Regex => {
243                if let Some(regex) = &self.regex_set {
244                    regex.replace_all(text, replacement).to_string()
245                } else {
246                    text.to_string()
247                }
248            }
249        }
250    }
251
252    /// Find all matches
253    ///
254    /// # Examples
255    ///
256    /// ```
257    /// use sensitive_rs::MultiPatternEngine;
258    ///
259    /// let patterns = vec!["赌博".to_string(), "色情".to_string()];
260    /// let engine = MultiPatternEngine::new(None, &patterns);
261    /// let matches = engine.find_all("含有赌博和色情");
262    /// assert_eq!(matches.len(), 2);
263    /// ```
264    pub fn find_all(&self, text: &str) -> Vec<String> {
265        match self.algorithm {
266            MatchAlgorithm::AhoCorasick => {
267                if let Some(ac) = &self.ac {
268                    ac.find_iter(text).map(|mat| text[mat.start()..mat.end()].to_string()).collect()
269                } else {
270                    Vec::new()
271                }
272            }
273            MatchAlgorithm::WuManber => {
274                if let Some(wm) = &self.wm {
275                    wm.search_all_strings(text)
276                } else {
277                    Vec::new()
278                }
279            }
280            MatchAlgorithm::Regex => {
281                if let Some(regex) = &self.regex_set {
282                    regex.find_iter(text).map(|mat| mat.as_str().to_string()).collect()
283                } else {
284                    Vec::new()
285                }
286            }
287        }
288    }
289
290    /// Get detailed match information
291    pub fn find_matches_with_positions(&self, text: &str) -> Vec<MatchInfo> {
292        match self.algorithm {
293            MatchAlgorithm::AhoCorasick => {
294                if let Some(ac) = &self.ac {
295                    ac.find_iter(text)
296                        .map(|mat| MatchInfo {
297                            pattern: text[mat.start()..mat.end()].to_string(),
298                            start: mat.start(),
299                            end: mat.end(),
300                        })
301                        .collect()
302                } else {
303                    Vec::new()
304                }
305            }
306            MatchAlgorithm::WuManber => {
307                if let Some(wm) = &self.wm {
308                    wm.find_matches(text)
309                        .into_iter()
310                        .filter_map(|m| {
311                            let pattern = text.get(m.start..m.end)?;
312                            Some(MatchInfo { pattern: pattern.to_string(), start: m.start, end: m.end })
313                        })
314                        .collect()
315                } else {
316                    Vec::new()
317                }
318            }
319            MatchAlgorithm::Regex => {
320                if let Some(regex) = &self.regex_set {
321                    regex
322                        .find_iter(text)
323                        .map(|mat| MatchInfo { pattern: mat.as_str().to_string(), start: mat.start(), end: mat.end() })
324                        .collect()
325                } else {
326                    Vec::new()
327                }
328            }
329        }
330    }
331
332    /// Check if text contains any patterns
333    pub fn contains_any(&self, text: &str) -> bool {
334        self.find_first(text).is_some()
335    }
336
337    /// Get engine statistics
338    pub fn stats(&self) -> EngineStats {
339        EngineStats {
340            algorithm: self.algorithm,
341            pattern_count: self.patterns.len(),
342            memory_usage: self.estimate_memory_usage(),
343        }
344    }
345
346    /// Estimate memory usage
347    fn estimate_memory_usage(&self) -> usize {
348        let patterns_memory = self.patterns.iter().map(|p| p.len()).sum::<usize>();
349
350        let engine_memory = match self.algorithm {
351            MatchAlgorithm::WuManber => {
352                if let Some(wm) = &self.wm {
353                    wm.memory_stats().total_memory
354                } else {
355                    0
356                }
357            }
358            _ => patterns_memory * 2, // Rough estimate for other algorithms
359        };
360
361        patterns_memory + engine_memory
362    }
363}
364
365/// Match information with position details
366#[derive(Debug, Clone)]
367pub struct MatchInfo {
368    pub pattern: String,
369    pub start: usize,
370    pub end: usize,
371}
372
373/// Engine statistics
374#[derive(Debug, Clone)]
375pub struct EngineStats {
376    pub algorithm: MatchAlgorithm,
377    pub pattern_count: usize,
378    pub memory_usage: usize,
379}
380
381#[cfg(test)]
382mod tests {
383    use super::*;
384
385    /// Build an engine from `&str` patterns using the auto-selected algorithm.
386    fn engine_with(patterns: &[&str]) -> MultiPatternEngine {
387        let owned: Vec<String> = patterns.iter().map(|s| s.to_string()).collect();
388        MultiPatternEngine::new(None, &owned)
389    }
390
391    #[test]
392    fn test_engine_find_first() {
393        let engine = engine_with(&["赌博", "色情"]);
394        assert_eq!(engine.find_first("含有赌博"), Some("赌博".to_string()));
395        assert_eq!(engine.find_first("正常"), None);
396    }
397
398    #[test]
399    fn test_engine_find_all() {
400        let engine = engine_with(&["赌博", "色情"]);
401        let results = engine.find_all("含有赌博和色情");
402        assert_eq!(results.len(), 2);
403    }
404
405    #[test]
406    fn test_engine_replace_all_wumanber() {
407        // Default for small pattern sets is WuManber, which repeats the
408        // replacement char per matched character: "赌博"(2 chars) -> "**".
409        let engine = engine_with(&["赌博"]);
410        assert_eq!(engine.current_algorithm(), MatchAlgorithm::WuManber);
411        assert_eq!(engine.replace_all("含有赌博内容", "*"), "含有**内容");
412    }
413
414    #[test]
415    fn test_engine_replace_all_empty_is_removal() {
416        let engine = engine_with(&["赌博"]);
417        assert_eq!(engine.replace_all("含有赌博内容", ""), "含有内容");
418    }
419
420    #[test]
421    fn test_engine_contains_any() {
422        let engine = engine_with(&["赌博"]);
423        assert!(engine.contains_any("含有赌博"));
424        assert!(!engine.contains_any("正常"));
425    }
426
427    #[test]
428    fn test_engine_find_matches_with_positions() {
429        // AhoCorasick yields correct byte offsets.
430        // (WuManber's find_matches currently panics on multi-byte text — a known
431        // issue tracked in CHANGELOG [Unreleased]; AhoCorasick/Regex are correct.)
432        let mut engine = engine_with(&["赌博"]);
433        engine.rebuild_with_algorithm(&["赌博".to_string()], MatchAlgorithm::AhoCorasick);
434        let text = "含有赌博内容";
435        let matches = engine.find_matches_with_positions(text);
436        assert_eq!(matches.len(), 1);
437        assert_eq!(matches[0].pattern, "赌博");
438        assert_eq!(matches[0].start, 6); // "含有" = 6 bytes
439        assert_eq!(matches[0].end, 12); // "赌博" = 6 bytes
440        assert_eq!(&text[matches[0].start..matches[0].end], "赌博");
441    }
442
443    #[test]
444    fn test_engine_find_matches_with_positions_regex() {
445        let mut engine = engine_with(&["赌博"]);
446        engine.rebuild_with_algorithm(&["赌博".to_string()], MatchAlgorithm::Regex);
447        let text = "含有赌博内容";
448        let matches = engine.find_matches_with_positions(text);
449        assert_eq!(matches.len(), 1);
450        assert_eq!(matches[0].pattern, "赌博");
451        assert_eq!(&text[matches[0].start..matches[0].end], "赌博");
452    }
453
454    #[test]
455    fn test_engine_stats() {
456        let engine = engine_with(&["赌博", "色情"]);
457        let stats = engine.stats();
458        assert_eq!(stats.pattern_count, 2);
459        assert_eq!(stats.algorithm, MatchAlgorithm::WuManber); // 2 patterns -> WuManber
460    }
461
462    #[test]
463    fn test_engine_empty() {
464        let engine = MultiPatternEngine::default();
465        assert!(engine.find_all("任何文本").is_empty());
466        assert_eq!(engine.find_first("任何文本"), None);
467        assert!(!engine.contains_any("任何文本"));
468    }
469
470    #[test]
471    fn test_engine_get_patterns() {
472        let engine = engine_with(&["赌博", "色情"]);
473        let patterns = engine.get_patterns();
474        assert_eq!(patterns.len(), 2);
475        assert!(patterns.contains(&"赌博".to_string()));
476    }
477
478    #[test]
479    fn test_engine_algorithm_recommendation() {
480        assert_eq!(MultiPatternEngine::recommend_algorithm(0), MatchAlgorithm::WuManber);
481        assert_eq!(MultiPatternEngine::recommend_algorithm(100), MatchAlgorithm::WuManber);
482        assert_eq!(MultiPatternEngine::recommend_algorithm(101), MatchAlgorithm::AhoCorasick);
483        assert_eq!(MultiPatternEngine::recommend_algorithm(10_000), MatchAlgorithm::AhoCorasick);
484        assert_eq!(MultiPatternEngine::recommend_algorithm(10_001), MatchAlgorithm::Regex);
485    }
486
487    #[test]
488    fn test_match_algorithm_display() {
489        assert_eq!(MatchAlgorithm::AhoCorasick.to_string(), "Aho-Corasick");
490        assert_eq!(MatchAlgorithm::WuManber.to_string(), "Wu-Manber");
491        assert_eq!(MatchAlgorithm::Regex.to_string(), "Regex");
492    }
493
494    #[test]
495    fn test_engine_force_algorithm_aho_corasick() {
496        let mut engine = engine_with(&["赌博"]);
497        engine.rebuild_with_algorithm(&["赌博".to_string()], MatchAlgorithm::AhoCorasick);
498        assert_eq!(engine.current_algorithm(), MatchAlgorithm::AhoCorasick);
499        assert_eq!(engine.find_first("含有赌博"), Some("赌博".to_string()));
500        // AhoCorasick replaces the whole match with the single replacement string.
501        assert_eq!(engine.replace_all("含有赌博内容", "*"), "含有*内容");
502    }
503
504    #[test]
505    fn test_engine_force_algorithm_regex() {
506        let mut engine = engine_with(&["赌博"]);
507        engine.rebuild_with_algorithm(&["赌博".to_string()], MatchAlgorithm::Regex);
508        assert_eq!(engine.current_algorithm(), MatchAlgorithm::Regex);
509        assert!(engine.contains_any("含有赌博"));
510        // Regex also replaces the whole match with the single replacement string.
511        assert_eq!(engine.replace_all("含有赌博内容", "*"), "含有*内容");
512    }
513
514    #[test]
515    fn test_engine_force_algorithm_wumanber() {
516        let mut engine = MultiPatternEngine::default();
517        engine.rebuild_with_algorithm(&["赌博".to_string()], MatchAlgorithm::WuManber);
518        assert_eq!(engine.current_algorithm(), MatchAlgorithm::WuManber);
519        assert_eq!(engine.find_all("含有赌博").len(), 1);
520    }
521
522    #[test]
523    fn test_engine_find_matches_with_positions_wumanber() {
524        // Regression for the WuManber find_matches multi-byte panic (now fixed):
525        // default small-set algorithm is WuManber and must yield correct offsets.
526        let engine = engine_with(&["赌博"]);
527        assert_eq!(engine.current_algorithm(), MatchAlgorithm::WuManber);
528        let text = "含有赌博内容";
529        let matches = engine.find_matches_with_positions(text);
530        assert_eq!(matches.len(), 1);
531        assert_eq!(matches[0].pattern, "赌博");
532        assert_eq!(&text[matches[0].start..matches[0].end], "赌博");
533    }
534}
sensitive_rs/engine/mod.rs

sensitive_rs/engine/
mod.rs
