dtxt-detect 1.0.0

use aho_corasick::AhoCorasick;
use deunicode::deunicode;

pub struct Engine {
    pub tier1_keys: Vec<String>,
    pub tier2_keys: Vec<String>,
    pub tier3_keys: Vec<String>,
    t1_corasick: Option<AhoCorasick>,
    t2_corasick: Option<AhoCorasick>,
    t3_corasick: Option<AhoCorasick>,
    pub keep_tier2_entries: bool,
    pub keep_tier3_entries: bool,
    pub fail_on_tier2: bool,
    pub fail_on_tier3: bool,
    pub fail_on_tier1_and_2: bool,
    pub normalize_unicode: bool,
    pub reset_string_on_fail: bool,
}
pub struct DtxtOutput {
    pub string: Option<String>,
    pub warnings: u32,
    pub fails: u32,
    /// Needs `keep_tier2_entries` to be true in the Engine
    pub tier2_entries: Vec<String>,
    /// Needs `keep_tier3_entires` to be true in the Engine
    pub tier3_entries: Vec<String>,
}

impl Default for DtxtOutput {
    fn default() -> Self {
        Self::new()
    }
}

impl DtxtOutput {
    pub fn new() -> DtxtOutput {
        DtxtOutput {
            string: None,
            warnings: 0,
            fails: 0,
            tier2_entries: vec![],
            tier3_entries: vec![],
        }
    }
}
impl Default for Engine {
    fn default() -> Self {
        Self::new()
    }
}

impl Engine {
    pub fn new() -> Engine {
        Engine {
            tier1_keys: vec![],
            tier2_keys: vec![],
            tier3_keys: vec![],
            t1_corasick: None,
            t2_corasick: None,
            t3_corasick: None,
            keep_tier2_entries: false,
            keep_tier3_entries: false,
            fail_on_tier2: false,
            fail_on_tier3: true,
            fail_on_tier1_and_2: true,
            normalize_unicode: true,
            reset_string_on_fail: true,
        }
    }
    /// Rebuild Engine, used when updating
    /// Normally it is done by the `Engine::load` function
    pub fn rebuild(&mut self) {
        self.t1_corasick = Some(AhoCorasick::builder()
            .ascii_case_insensitive(true)
            .build(&self.tier1_keys)
            .unwrap());
        self.t2_corasick = Some(AhoCorasick::builder()
            .ascii_case_insensitive(true)
            .build(&self.tier2_keys)
            .unwrap());
        self.t3_corasick = Some(AhoCorasick::builder()
            .ascii_case_insensitive(true)
            .build(&self.tier3_keys)
            .unwrap());
    }
    /// Load key datasets inside of the Engine
    pub fn load(
        &mut self,
        tier1_keys: Vec<String>,
        tier2_keys: Vec<String>,
        tier3_keys: Vec<String>,
        concat: bool,
        rebuild: bool,
    ) {
        if concat {
            self.tier1_keys.extend(tier1_keys);
            self.tier2_keys.extend(tier2_keys);
            self.tier3_keys.extend(tier3_keys);
        } else {
            self.tier1_keys = tier1_keys;
            self.tier2_keys = tier2_keys;
            self.tier3_keys = tier3_keys;
        }

        if rebuild {
            // Needed as the Corasick's objects need to be rebuilt  with the new keys
            self.rebuild();
        }
    }
    /// Process a String using the Engine
    pub fn process(&self, input: String) -> DtxtOutput {
        let mut work_string = input;
        let mut dtxt_output = DtxtOutput::new();

        // Normalize unicode
        if self.normalize_unicode {
            work_string = deunicode(&work_string);
        }

        // Detect tier-1 words
        let ac1 = self.t1_corasick
            .as_ref()
            .expect("[dtxt-detect] Engine not built. Call `.rebuild()` after loading keys.");
        let mut tier1_entries = vec![];
        for mat in ac1.find_iter(&work_string) {
            let pattern_id = mat.pattern().as_usize();
            tier1_entries.push(self.tier1_keys[pattern_id].clone());
        }

        // Detect tier-2 words
        let ac2 = self.t2_corasick
            .as_ref()
            .expect("[dtxt-detect] Engine not built. Call `.rebuild()` after loading keys.");
        let mut tier2_entries = vec![];
        for mat in ac2.find_iter(&work_string) {
            let pattern_id = mat.pattern().as_usize();
            tier2_entries.push(self.tier2_keys[pattern_id].clone());
        }

        // Detect tier-3 words
        let ac3 = self.t3_corasick
            .as_ref()
            .expect("[dtxt-detect] Engine not built. Call `.rebuild()` after loading keys.");
        let mut tier3_entries = vec![];
        for mat in ac3.find_iter(&work_string) {
            let pattern_id = mat.pattern().as_usize();
            tier3_entries.push(self.tier3_keys[pattern_id].clone());
        }

        // Check for warnings and fails
        if self.fail_on_tier2 {
            if !tier2_entries.is_empty() {
                dtxt_output.fails += tier2_entries.len() as u32;
            }
        } else {
            dtxt_output.warnings += tier2_entries.len() as u32;
        }
        if self.fail_on_tier3 {
            if !tier3_entries.is_empty() {
                dtxt_output.fails += tier3_entries.len() as u32;
            }
        } else {
            dtxt_output.warnings += tier3_entries.len() as u32;
        }
        if self.fail_on_tier1_and_2
            && !tier1_entries.is_empty()
            && !tier2_entries.is_empty()
            && !self.fail_on_tier2
        {
            // Do not consider if already failed on tier2 in order to have recounting the same failure
            dtxt_output.fails += 0
        }

        // Register needed entries
        if self.keep_tier2_entries {
            dtxt_output.tier2_entries = tier2_entries;
        }
        if self.keep_tier3_entries {
            dtxt_output.tier3_entries = tier3_entries;
        }

        // Return empty message if configured this way
        if self.reset_string_on_fail && dtxt_output.fails > 0 {
            dtxt_output.string = None;
        } else {
            dtxt_output.string = Some(work_string);
        }
        dtxt_output
    }
}