rust_censure/
censor.rs

1use once_cell::sync::Lazy;
2use std::collections::HashMap;
3use std::sync::RwLock;
4
5use super::structs::*;
6
7use crate::lang::common::{
8    NORMALIZATION_PATTERNS, PAT_SPACE, PAT_PUNCT1, PAT_PUNCT2, PAT_PUNCT3, PAT_PREP
9};
10use crate::lang::LangProvider;
11use crate::util::{remove_duplicates, is_pi_or_e_word};
12use fancy_regex;
13
14
15impl Censor {
16    pub fn new(lang: CensorLang) -> Result<Self, CensorError> {
17        let data = match lang {
18            CensorLang::Ru => crate::lang::ru::Ru::data(),
19            CensorLang::En => crate::lang::en::En::data(),
20        };
21        Ok(Self { lang, data, re_cache: Lazy::new(|| RwLock::new(HashMap::new())) })
22    }
23
24    fn is_match_cached(&self, pat: &str, text: &str) -> bool {
25        // Check cache
26        {
27            let cache = self.re_cache.read().unwrap();
28            if let Some(r) = cache.get(pat) {
29                return r.is_match(text).unwrap_or(false)
30            }
31        }
32
33        // Compile and cache
34        let mut cache = self.re_cache.write().unwrap();
35        let r = fancy_regex::Regex::new(pat).expect("invalid regex");
36        let res = r.is_match(text).unwrap_or(false);
37        self.cache_pattern(pat, r, &mut cache); // cache this pattern
38        res
39    }
40
41    fn cache_pattern(&self, pat: &str, r: fancy_regex::Regex, cache: &mut std::sync::RwLockWriteGuard<HashMap<String, fancy_regex::Regex>>) {
42        // Check cache
43        if let Some(_) = cache.get(pat) {
44            return // already cached
45        }
46
47        cache.insert(pat.to_string(), r);
48    }
49
50    fn compile_and_cache_pattern(&self, pat: &str, cache: &mut std::sync::RwLockWriteGuard<HashMap<String, fancy_regex::Regex>>) {
51        let r = fancy_regex::Regex::new(pat).expect("invalid regex");
52        self.cache_pattern(pat, r, cache);
53    }
54
55    pub fn precompile_all_patterns(&self) {
56        self.precompile_foul_data();
57        self.precompile_foul_core();
58        self.precompile_bad_phrases();
59        self.precompile_bad_semi_phrases();
60        self.precompile_excludes_core();
61        self.precompile_excludes_data();
62    }
63
64    pub fn precompile_foul_data(&self) {
65        let mut cache = self.re_cache.write().unwrap();
66
67        for (_, pats) in self.data.foul_data {
68            for &pat in pats {
69                self.compile_and_cache_pattern(pat, &mut cache);
70            }
71        }
72    }
73
74    pub fn precompile_foul_core(&self) {
75        let mut cache = self.re_cache.write().unwrap();
76
77        for (pat, _) in self.data.foul_core {
78            self.compile_and_cache_pattern(pat, &mut cache);
79        }
80    }
81
82    pub fn precompile_bad_phrases(&self) {
83        let mut cache = self.re_cache.write().unwrap();
84
85        for &pat in self.data.bad_phrases {
86            self.compile_and_cache_pattern(pat, &mut cache);
87        }
88    }
89
90    pub fn precompile_bad_semi_phrases(&self) {
91        let mut cache = self.re_cache.write().unwrap();
92
93        for &pat in self.data.bad_semi_phrases {
94            self.compile_and_cache_pattern(pat, &mut cache);
95        }
96    }
97
98    pub fn precompile_excludes_core(&self) {
99        let mut cache = self.re_cache.write().unwrap();
100
101        for (pat, _) in self.data.excludes_core {
102            self.compile_and_cache_pattern(pat, &mut cache);
103        }
104    }
105
106    pub fn precompile_excludes_data(&self) {
107        let mut cache = self.re_cache.write().unwrap();
108
109        for (_, pats) in self.data.excludes_data {
110            for &pat in pats {
111                self.compile_and_cache_pattern(pat, &mut cache);
112            }
113        }
114    }
115
116    fn replace_all_cached<'a>(&self, pat: &str, text: &'a str, repl: &str) -> Option<String> {
117        // Quick negative guard: if it doesn't match, skip compiling/allocating a String for replace.
118        if !self.is_match_cached(pat, text) {
119            return None;
120        }
121
122        // read from cache
123        let cache = self.re_cache.read().unwrap();
124        let compiled = cache.get(pat).unwrap();
125
126        // replace
127        let replaced = compiled.replace_all(text, repl).into_owned();
128        if replaced == text { None } else { Some(replaced) }
129    }
130
131    fn split_line_ru(&self, line: &str) -> Vec<String> {
132        // port of CensorRu._split_line: remove punctuation1, then punctuation2 -> space
133        let step1 = PAT_PUNCT1.replace_all(line, "");
134        let step2 = PAT_PUNCT2.replace_all(&step1, " ");
135        let mut buf = String::new();
136        let mut out = Vec::new();
137
138        for w in PAT_SPACE.split(&step2) {
139            let w = w.unwrap();
140
141            if w.is_empty() { continue; }
142            if w.chars().count() < 3 && !PAT_PREP.is_match(w).unwrap_or(false) {
143                buf.push_str(w);
144            } else {
145                if !buf.is_empty() {
146                    out.push(std::mem::take(&mut buf));
147                }
148                out.push(w.to_string());
149            }
150        }
151        if !buf.is_empty() { out.push(buf); }
152        out
153    }
154
155    fn split_line_en(&self, line: &str) -> Vec<String> {
156        // behave like CensorEn._split_line
157        let step1 = PAT_PUNCT1.replace_all(line, "");
158        let step2 = PAT_PUNCT2.replace_all(&step1, " ");
159        let mut buf = String::new();
160        let mut out = Vec::new();
161
162        for w in PAT_SPACE.split(&step2) {
163            let w = w.unwrap();
164
165            if w.is_empty() { continue; }
166            if w.chars().count() < 3 {
167                buf.push_str(w);
168            } else {
169                if !buf.is_empty() {
170                    out.push(std::mem::take(&mut buf));
171                }
172                out.push(w.to_string());
173            }
174        }
175        if !buf.is_empty() { out.push(buf); }
176        out
177    }
178
179    fn split_line(&self, s: &str) -> Vec<String> {
180        match self.lang {
181            CensorLang::Ru => self.split_line_ru(s),
182            CensorLang::En => self.split_line_en(s),
183        }
184    }
185
186    fn prepare_word(&self, mut w: String) -> String {
187        if !is_pi_or_e_word(&w) {
188            // trim punctuation edges
189            w = PAT_PUNCT3.replace_all(&w, "").into_owned();
190        }
191        let mut w = w.to_lowercase();
192
193        // apply normalization patterns in order
194        for (pat, rep) in NORMALIZATION_PATTERNS.iter() {
195            w = pat.replace_all(&w, *rep).into_owned();
196        }
197
198        // transliteration of similar chars
199        w = crate::lang::common::translate_similar_chars(&w, self.data.trans_tab);
200
201        // deduplicate (AAA -> AA)
202        remove_duplicates(&w)
203    }
204
205    pub fn is_word_good(&self, raw: &str) -> bool {
206        let w = self.prepare_word(raw.to_string());
207        self.check_word_impl(&w).is_good
208    }
209
210    fn check_word_impl(&self, prepared: &str) -> WordInfo {
211        let mut info = WordInfo::new(prepared.to_string());
212
213        // Keys in your maps are &str, so we need to build a string from the first character
214        let fl_str = prepared.chars().next().map(|c| {
215            // Important: make a one-character string
216            // (if you want grapheme-cluster support, use unicode-segmentation instead)
217            let mut s = String::new();
218            s.push(c);
219            s
220        }).unwrap_or_default();
221
222        // 1) Accuse stage: FOUL_DATA[first_letter]
223        if let Some(pats) = self.data.foul_data.get(fl_str.as_str()) {
224            for &pat in pats {
225                if self.is_match_cached(pat, prepared) {
226                    info.is_good = false;
227                    info.accuse.push(pat.to_string()); // now stored as a string rule
228                    break;
229                }
230            }
231        }
232
233        // 2) If still good → check FOUL_CORE
234        if info.is_good {
235            for (&_key, &pat) in self.data.foul_core.iter() {
236                if self.is_match_cached(pat, prepared) {
237                    info.is_good = false;
238                    info.accuse.push(pat.to_string());
239                    break;
240                }
241            }
242        }
243
244        // 3) If still good → check BAD_SEMI_PHRASES
245        if info.is_good {
246            for &pat in self.data.bad_semi_phrases.iter() {
247                if self.is_match_cached(pat, prepared) {
248                    info.is_good = false;
249                    info.accuse.push(pat.to_string());
250                    break;
251                }
252            }
253        }
254
255        // 4) Excuse stage: if already accused, check exceptions
256        if !info.is_good {
257            // EXCLUDES_CORE
258            for (&_key, &pat) in self.data.excludes_core.iter() {
259                if self.is_match_cached(pat, prepared) {
260                    info.is_good = true;
261                    info.excuse.push(pat.to_string());
262                    break;
263                }
264            }
265            // EXCLUDES_DATA[first_letter]
266            if !info.is_good {
267                if let Some(pats) = self.data.excludes_data.get(fl_str.as_str()) {
268                    for &pat in pats {
269                        if self.is_match_cached(pat, prepared) {
270                            info.is_good = true;
271                            info.excuse.push(pat.to_string());
272                            break;
273                        }
274                    }
275                }
276            }
277        }
278
279        info
280    }
281
282    /// returns replaced line plus counts
283    pub fn clean_line(&self, line: &str) -> CleanLineResult {
284        // Mutable working buffer that accumulates changes
285        let mut out = line.to_string();
286
287        // Counters and diagnostics
288        let mut bad_words = 0usize;
289        let mut bad_phrases = 0usize;
290        let mut detected_words = Vec::new();
291        let mut detected_pats = Vec::new();
292
293        // 1) Word-by-word replacement (first hit per surface word):
294        //
295        // - Split the *original* line into tokens according to language rules.
296        // - For each token, normalize and check with accuse/excuse logic.
297        // - If bad, replace the *first* occurrence of the exact surface token in `out`.
298        //   This preserves original casing/punctuation and mirrors your Python behavior.
299        for word in self.split_line(line) {
300            let prepared = self.prepare_word(word.clone());
301            let info = self.check_word_impl(&prepared);
302            if !info.is_good {
303                bad_words += 1;
304                out = out.replacen(&word, self.data.beep, 1);
305                detected_words.push(word);
306                if let Some(p) = info.accuse.get(0) {
307                    detected_pats.push(p.clone());
308                }
309            }
310        }
311
312        // 2) Phrase-level replacements:
313        //
314        // - BAD_SEMI_PHRASES are broad patterns that run over the whole string.
315        // - We first check via `is_match_cached` to avoid unnecessary work,
316        //   then call `replace_all_cached` which compiles via the same cache.
317        for &pat in self.data.bad_semi_phrases.iter() {
318            if let Some(new_out) = self.replace_all_cached(pat, &out, self.data.beep) {
319                bad_phrases += 1;
320                detected_pats.push(pat.to_string());
321                out = new_out;
322            }
323        }
324
325        // If you also maintain BAD_PHRASES, process them the same way:
326        for &pat in self.data.bad_phrases.iter() {
327            if let Some(new_out) = self.replace_all_cached(pat, &out, self.data.beep) {
328                bad_phrases += 1;
329                detected_pats.push(pat.to_string());
330                out = new_out;
331            }
332        }
333
334        CleanLineResult {
335            line: out,
336            bad_words_count: bad_words,
337            bad_phrases_count: bad_phrases,
338            detected_bad_words: detected_words,
339            detected_patterns: detected_pats,
340        }
341    }
342
343    /// Clean an HTML string while preserving tags and replacing bad words with `beep_html`.
344    /// @TODO: Rewrite the implementation, so it'll work with any HTML tags (incl broken etc).
345    pub fn clean_html_line(&self, line: &str) -> CleanHtmlResult {
346        use crate::html::{tokenize_html, TokType, Token};
347
348        let tokens = tokenize_html(line);
349
350        let mut current_word = String::new();            // plain word (no tags)
351        let mut current_tagged = String::new();          // word with tags as text
352        let mut tagged_list: Vec<&Token> = Vec::new();   // token objects for pre/post reconstruction
353
354        let mut out = String::new();
355        let mut bad_count = 0usize;
356
357        let beep_html = self.data.beep_html; // HTML replacement for a bad word
358
359        // Compute "pre" (opening + self-closing tags) and "post" (closing tags)
360        // from the tokens collected for the current word.
361        fn get_remained_tokens(tagged: &[&Token]) -> (String, String) {
362            let mut pre = String::new();
363            let mut post = String::new();
364
365            for t in tagged {
366                match t.kind {
367                    TokType::TagOpen | TokType::TagSelf => {
368                        // opening/self tags should remain before the censored placeholder
369                        pre.push_str(&t.value);
370                    }
371                    TokType::TagClose => {
372                        // closing tags should remain after the censored placeholder
373                        post.push_str(&t.value);
374                    }
375                    _ => {}
376                }
377            }
378            (pre, post)
379        }
380
381        // Flush the currently accumulated word (and its tag list) into `out`.
382        // If the word is bad, we output `pre + beep_html + post`. Otherwise, we output the original tagged text.
383        // Optionally append a trailing literal (space/spacer) after flushing.
384        let process_spacer = |cw: &mut String,
385                                  ctw: &mut String,
386                                  twl: &mut Vec<&Token>,
387                                  r: &mut String,
388                                  bwc: &mut usize,
389                                  tok: Option<&Token>| {
390            if !cw.is_empty() {
391                // println!("{}", cw);
392                if !self.is_word_good(cw) {
393                    let (pre, post) = get_remained_tokens(twl);
394                    *r += &pre;
395                    *r += beep_html;
396                    *r += &post;
397                    *bwc += 1;
398                } else {
399                    // Good word: emit the original tagged fragment unchanged
400                    *r += ctw;
401                }
402            }
403            // Reset per-word buffers
404            twl.clear();
405            cw.clear();
406            ctw.clear();
407
408            // Append trailing boundary (space/spacer) if provided
409            if let Some(t) = tok {
410                *r += &t.value;
411            }
412        };
413
414        // Iterate over tokens exactly like the Python version
415        for tok in &tokens {
416            match tok.kind {
417                TokType::TagOpen | TokType::TagClose | TokType::TagSelf => {
418                    // Tags are part of the current "tagged word"; they do NOT trigger a flush
419                    tagged_list.push(tok);
420                    current_tagged.push_str(&tok.value);
421                }
422                TokType::Word => {
423                    // Word fragments are appended to both plain and tagged buffers
424                    // println!("current_word: {}", current_word);
425                    if !self.is_word_good(&current_word) {
426                        process_spacer(
427                            &mut current_word,
428                            &mut current_tagged,
429                            &mut tagged_list,
430                            &mut out,
431                            &mut bad_count,
432                            Some(tok),
433                        );
434                    } else {
435                        tagged_list.push(tok);
436                        current_tagged.push_str(&tok.value);
437                        current_word.push_str(&tok.value);
438                    }
439                }
440                TokType::Space |  TokType::Spacer => {
441                    // Boundary: process the current word and then append the space/spacer
442                    process_spacer(
443                        &mut current_word,
444                        &mut current_tagged,
445                        &mut tagged_list,
446                        &mut out,
447                        &mut bad_count,
448                        Some(tok),
449                    );
450                }
451            }
452        }
453
454        // Final flush if the line ended without a trailing space
455        if !current_word.is_empty() || !current_tagged.is_empty() {
456            process_spacer(
457                &mut current_word,
458                &mut current_tagged,
459                &mut tagged_list,
460                &mut out,
461                &mut bad_count,
462                None,
463            );
464        }
465
466        CleanHtmlResult { line: out, bad_words_count: bad_count }
467    }
468}