use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::LazyLock;
use unicode_normalization::UnicodeNormalization;
use regex::Regex;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum DigitPolicy {
#[serde(rename = "fa")] Fa, #[serde(rename = "ar")] Ar, #[serde(rename = "latin")] Latin, #[serde(rename = "auto")] Auto, #[serde(rename = "preserve")] Preserve, }
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum PunctPolicy {
#[serde(rename = "fa")] Fa, #[serde(rename = "latin")] Latin, #[serde(rename = "keep")] Keep, }
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum ZwnjPolicy {
#[serde(rename = "smart")] Smart, #[serde(rename = "force")] Force, #[serde(rename = "none")] None, }
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum UnicodeForm {
#[serde(rename = "NFC")] Nfc,
#[serde(rename = "NFKC")] Nfkc, }
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default)]
pub struct NormalizeOptions {
pub locale: String,
pub unicode_form: UnicodeForm,
pub remove_diacritics: bool,
pub remove_tatweel: bool,
pub digits: DigitPolicy,
pub punctuation: PunctPolicy,
pub zwnj: ZwnjPolicy,
pub trim: bool,
pub squeeze_whitespace: bool,
pub normalize_newlines: bool,
pub drop_bidi_controls: bool,
pub confusable_safe: bool,
pub protect_urls: bool,
pub protect_emails: bool,
pub protect_code: bool,
pub protect_html_tags: bool,
pub slang_map: HashMap<String, String>,
pub zwnj_compound_words: Vec<String>,
pub custom_rules: Vec<(String, String)>,
}
impl Default for NormalizeOptions {
fn default() -> Self {
NormalizeOptions {
locale: "fa-IR".to_string(),
unicode_form: UnicodeForm::Nfkc,
remove_diacritics: true,
remove_tatweel: true,
digits: DigitPolicy::Auto,
punctuation: PunctPolicy::Fa,
zwnj: ZwnjPolicy::Smart,
trim: true,
squeeze_whitespace: true,
normalize_newlines: true,
drop_bidi_controls: false,
confusable_safe: false,
protect_urls: true,
protect_emails: true,
protect_code: true,
protect_html_tags: true,
slang_map: HashMap::new(),
zwnj_compound_words: Vec::new(),
custom_rules: Vec::new(),
}
}
}
const LATIN_DIGITS: [char; 10] = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'];
const PERSIAN_DIGITS: [char; 10] = ['۰', '۱', '۲', '۳', '۴', '۵', '۶', '۷', '۸', '۹'];
const ARABIC_DIGITS: [char; 10] = ['٠', '١', '٢', '٣', '٤', '٥', '٦', '٧', '٨', '٩'];
fn arabic_to_persian(c: char) -> char {
match c {
'ي' | 'ى' => 'ی',
'إ' | 'أ' | 'ٱ' => 'ا',
'ؤ' => 'و',
'ك' | 'ﻙ' => 'ک',
'ۀ' | 'ة' => 'ه',
_ => c,
}
}
fn punct_fa(c: char) -> char {
match c {
',' => '،',
';' => '؛',
'?' => '؟',
_ => c,
}
}
fn punct_latin(c: char) -> char {
match c {
'،' => ',',
'؛' => ';',
'؟' => '?',
_ => c,
}
}
fn confusable_safe(c: char) -> char {
match c {
'۰'..='۹' => char::from_u32(c as u32 - '۰' as u32 + '0' as u32).unwrap(),
'٠'..='٩' => char::from_u32(c as u32 - '٠' as u32 + '0' as u32).unwrap(),
'ك' => 'ک',
'ي' | 'ى' => 'ی',
'ة' | 'ۀ' => 'ه',
_ => c,
}
}
static MI_NEMI_PREFIX_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"(می|نمی)\s+([\u0600-\u06FF]{2,})").unwrap());
static SUFFIXES_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"([\u0600-\u06FF]{2,})\s+(ها|های|تر|ترین|ام|ات|اش|ایم|اید|اند|گان|مند|وار)\b").unwrap());
static FORCE_ZWNJ_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"([آ-ی])\s+([آ-ی])").unwrap());
static SQUEEZE_WHITESPACE_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"[ \t\u{00A0}]+").unwrap());
static SQUEEZE_POST_PUNCT_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\s+([،,؛;:!?.)\]}])").unwrap());
static SQUEEZE_PRE_PUNCT_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"([\(«\[{])\s+").unwrap());
static URL_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#"\bhttps?:\/\/[^\s)>"']+|www\.[^\s)>"']+"#).unwrap());
static EMAIL_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"(?i)\b[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}\b").unwrap());
static CODE_BLOCK_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"```[\s\S]*?```").unwrap());
static INLINE_CODE_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"`[^`]*`").unwrap());
static HTML_TAG_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"<\/?[\w\d-]+(\s+[^>]*?)?>").unwrap());
fn to_form(s: &str, form: &UnicodeForm) -> String {
match form {
UnicodeForm::Nfc => s.nfc().collect(),
UnicodeForm::Nfkc => s.nfkc().collect(),
}
}
fn strip_tatweel(s: &str) -> String { s.replace('\u{0640}', "") }
fn strip_diacritics(s: &str) -> String {
s.chars().filter(|&c| !matches!(c, '\u{0610}'..='\u{061A}' | '\u{064B}'..='\u{065F}' | '\u{06D6}'..='\u{06ED}')).collect()
}
fn drop_bidi_controls(s: &str) -> String {
s.chars().filter(|&c| !matches!(c, '\u{200E}'..='\u{200F}' | '\u{202A}'..='\u{202E}' | '\u{2066}'..='\u{2069}')).collect()
}
fn unify_letters_fa(s: &str) -> String {
s.chars().map(arabic_to_persian).collect()
}
fn normalize_digits_impl(s: &str, target_digits: &[char; 10]) -> String {
s.chars().map(|c| {
match c {
'0'..='9' => target_digits[(c as u32 - '0' as u32) as usize],
'۰'..='۹' => target_digits[(c as u32 - '۰' as u32) as usize],
'٠'..='٩' => target_digits[(c as u32 - '٠' as u32) as usize],
_ => c,
}
}).collect()
}
fn normalize_digits(s: &str, policy: &DigitPolicy) -> String {
match policy {
DigitPolicy::Preserve => s.to_string(),
DigitPolicy::Auto => {
if s.chars().any(|c| ('\u{0600}'..='\u{06FF}').contains(&c)) {
normalize_digits_impl(s, &PERSIAN_DIGITS)
} else {
normalize_digits_impl(s, &LATIN_DIGITS)
}
},
DigitPolicy::Fa => normalize_digits_impl(s, &PERSIAN_DIGITS),
DigitPolicy::Ar => normalize_digits_impl(s, &ARABIC_DIGITS),
DigitPolicy::Latin => normalize_digits_impl(s, &LATIN_DIGITS),
}
}
fn normalize_punctuation(s: &str, policy: &PunctPolicy) -> String {
match policy {
PunctPolicy::Keep => s.to_string(),
PunctPolicy::Fa => s.chars().map(punct_fa).collect(),
PunctPolicy::Latin => s.chars().map(punct_latin).collect(),
}
}
fn apply_slang_normalization(s: &str, slang_map: &HashMap<String, String>) -> String {
if slang_map.is_empty() {
return s.to_string();
}
let slang_keys: Vec<String> = slang_map.keys().map(|k| regex::escape(k)).collect();
let pattern = format!(r"\b({})\b", slang_keys.join("|"));
let slang_re = Regex::new(&pattern).unwrap_or_else(|_| Regex::new("a^").unwrap());
slang_re.replace_all(s, |caps: ®ex::Captures| {
slang_map.get(&caps[0]).cloned().unwrap_or_else(|| caps[0].to_string())
}).to_string()
}
fn smart_zwnj(s: &str, zwnj_compound_words: &[String]) -> String {
let mut s_cloned = s.to_string();
s_cloned = MI_NEMI_PREFIX_RE.replace_all(&s_cloned, "$1\u{200c}$2").to_string();
s_cloned = SUFFIXES_RE.replace_all(&s_cloned, "$1\u{200c}$2").to_string();
if !zwnj_compound_words.is_empty() {
for compound in zwnj_compound_words {
if let Some(pos) = compound.find(' ') {
let (part1, part2) = compound.split_at(pos);
let with_zwnj = format!("{}\u{200c}{}", part1, part2.trim_start());
s_cloned = s_cloned.replace(compound, &with_zwnj);
}
}
}
s_cloned.replace(" \u{200c}", "\u{200c}").replace("\u{200c} ", "\u{200c}")
}
fn force_zwnj(s: &str) -> String {
FORCE_ZWNJ_RE.replace_all(s, "$1\u{200c}$2").to_string()
}
fn normalize_whitespace(s: &str, squeeze: bool, trim: bool, normalize_newlines: bool) -> String {
let mut s_cloned = s.to_string();
if normalize_newlines { s_cloned = s_cloned.replace("\r\n", "\n").replace('\r', "\n"); }
if squeeze {
s_cloned = SQUEEZE_WHITESPACE_RE.replace_all(&s_cloned, " ").to_string();
s_cloned = SQUEEZE_POST_PUNCT_RE.replace_all(&s_cloned, "$1").to_string();
s_cloned = SQUEEZE_PRE_PUNCT_RE.replace_all(&s_cloned, "$1").to_string();
}
if trim { s_cloned = s_cloned.trim().to_string(); }
s_cloned
}
#[derive(Debug, Clone, PartialEq)]
enum SpanKind { Text, Url, Email, Code, Html }
#[derive(Debug, Clone)]
struct Span { kind: SpanKind, start: usize, end: usize }
fn segment(text: &str, options: &NormalizeOptions) -> Vec<Span> {
let mut blocks: Vec<Span> = vec![Span { kind: SpanKind::Text, start: 0, end: text.len() }];
let carve = |re: &Regex, kind: SpanKind, blocks: &mut Vec<Span>| {
let mut new_blocks: Vec<Span> = Vec::new();
for b in blocks.iter() {
if b.kind != SpanKind::Text { new_blocks.push(b.clone()); continue; }
let mut current_pos = b.start;
for m in re.find_iter(&text[b.start..b.end]) {
let (s, e) = (b.start + m.start(), b.start + m.end());
if current_pos < s { new_blocks.push(Span { kind: SpanKind::Text, start: current_pos, end: s }); }
new_blocks.push(Span { kind: kind.clone(), start: s, end: e });
current_pos = e;
}
if current_pos < b.end { new_blocks.push(Span { kind: SpanKind::Text, start: current_pos, end: b.end }); }
}
*blocks = new_blocks;
};
if options.protect_html_tags { carve(&HTML_TAG_RE, SpanKind::Html, &mut blocks); }
if options.protect_code {
carve(&CODE_BLOCK_RE, SpanKind::Code, &mut blocks);
carve(&INLINE_CODE_RE, SpanKind::Code, &mut blocks);
}
if options.protect_emails { carve(&EMAIL_RE, SpanKind::Email, &mut blocks); }
if options.protect_urls { carve(&URL_RE, SpanKind::Url, &mut blocks); }
blocks.sort_by_key(|s| s.start);
blocks
}
pub fn normalize_text(input: &str, options: &NormalizeOptions) -> String {
let mut output = String::new();
let spans = segment(input, options);
for span in &spans {
let chunk = &input[span.start..span.end];
if span.kind == SpanKind::Text {
let mut processed_chunk = chunk.to_string();
processed_chunk = to_form(&processed_chunk, &options.unicode_form);
if options.remove_tatweel { processed_chunk = strip_tatweel(&processed_chunk); }
if options.remove_diacritics { processed_chunk = strip_diacritics(&processed_chunk); }
if options.locale == "fa-IR" { processed_chunk = unify_letters_fa(&processed_chunk); }
if !options.custom_rules.is_empty() {
for (from, to) in &options.custom_rules {
processed_chunk = processed_chunk.replace(from, to);
}
}
processed_chunk = apply_slang_normalization(&processed_chunk, &options.slang_map);
processed_chunk = normalize_digits(&processed_chunk, &options.digits);
processed_chunk = normalize_punctuation(&processed_chunk, &options.punctuation);
match &options.zwnj {
ZwnjPolicy::Smart => processed_chunk = smart_zwnj(&processed_chunk, &options.zwnj_compound_words),
ZwnjPolicy::Force => processed_chunk = force_zwnj(&processed_chunk),
ZwnjPolicy::None => (),
}
processed_chunk = normalize_whitespace(&processed_chunk, options.squeeze_whitespace, false, options.normalize_newlines);
if options.drop_bidi_controls { processed_chunk = drop_bidi_controls(&processed_chunk); }
if options.confusable_safe {
processed_chunk = processed_chunk.chars().map(confusable_safe).collect();
}
output.push_str(&processed_chunk);
} else {
output.push_str(chunk);
}
}
if options.trim { output.trim().to_string() } else { output }
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_basic_normalization() {
let text = "مي روم 12,345 ريال، كتاب ها ١٢٣";
let options = NormalizeOptions { digits: DigitPolicy::Fa, punctuation: PunctPolicy::Fa, zwnj: ZwnjPolicy::Smart, ..Default::default() };
let expected = "میروم ۱۲،۳۴۵ ریال، کتابها ۱۲۳";
assert_eq!(normalize_text(text, &options), expected);
}
#[test]
fn test_user_provided_rules() {
let text = "من توی کتاب خانه کار میکنم.";
let mut slang_map = HashMap::new();
slang_map.insert("توی".to_string(), "در".to_string());
slang_map.insert("کار میکنم".to_string(), "کار میکنم".to_string());
let options = NormalizeOptions {
slang_map,
zwnj_compound_words: vec!["کتاب خانه".to_string()],
custom_rules: vec![("من ".to_string(), "بنده ".to_string())],
..Default::default()
};
let expected = "بنده در کتابخانه کار میکنم.";
assert_eq!(normalize_text(text, &options), expected);
}
}