use async_trait::async_trait;
use std::collections::HashMap;
use crate::processor::{Correction, CorrectionKind, ProcessError, ProcessResult, TextProcessor};
use crate::types::{ContextSnapshot, Language, Span};
const MIN_ALIAS_CHARS: usize = 2;
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize)]
pub struct TermEntry {
pub term: String,
pub aliases: Vec<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum MatchScope {
WordBoundary,
Substring,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MatchPolicy {
scope: MatchScope,
fold_case: bool,
fold_fullwidth: bool,
fold_kana: bool,
fold_whitespace: bool,
}
impl MatchPolicy {
pub fn for_language(lang: Language) -> Self {
let scope = match lang {
Language::Japanese | Language::Chinese => MatchScope::Substring,
_ => MatchScope::WordBoundary,
};
Self {
scope,
fold_case: true,
fold_fullwidth: true,
fold_kana: matches!(lang, Language::Japanese),
fold_whitespace: true,
}
}
pub fn none() -> Self {
Self {
scope: MatchScope::Substring,
fold_case: false,
fold_fullwidth: false,
fold_kana: false,
fold_whitespace: false,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum Problem {
EmptyTerm { entry: usize },
EmptyAlias { entry: usize, alias: usize },
AliasEqualsTerm {
entry: usize,
alias: usize,
term: String,
},
AliasTooShort {
entry: usize,
alias: usize,
text: String,
folded_chars: usize,
},
ConflictingAlias {
alias: String,
first_term: String,
second_term: String,
},
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("dictionary has {} problem(s): {problems:?}", problems.len())]
pub struct DictionaryError {
pub problems: Vec<Problem>,
}
struct Folded {
chars: Vec<char>,
origin: Vec<usize>,
}
fn fold(text: &str, policy: &MatchPolicy) -> Folded {
let mut chars = Vec::new();
let mut origin = Vec::new();
for (index, c) in text.chars().enumerate() {
let c = if policy.fold_fullwidth {
match c {
'\u{FF01}'..='\u{FF5E}' => char::from_u32(c as u32 - 0xFEE0).unwrap_or(c),
'\u{3000}' => ' ',
other => other,
}
} else {
c
};
let c = if policy.fold_kana && ('\u{3041}'..='\u{3096}').contains(&c) {
char::from_u32(c as u32 + 0x60).unwrap_or(c)
} else {
c
};
if policy.fold_whitespace && c.is_whitespace() {
if chars.last() == Some(&' ') {
continue;
}
chars.push(' ');
origin.push(index);
continue;
}
if policy.fold_case {
for lowered in c.to_lowercase() {
chars.push(lowered);
origin.push(index);
}
} else {
chars.push(c);
origin.push(index);
}
}
Folded { chars, origin }
}
#[derive(Debug)]
struct CompiledAlias {
folded: Vec<char>,
term: String,
}
#[derive(Debug)]
pub struct TermDictionary {
aliases: Vec<CompiledAlias>,
policy: MatchPolicy,
}
impl TermDictionary {
pub fn new(
entries: impl IntoIterator<Item = TermEntry>,
policy: MatchPolicy,
) -> Result<Self, DictionaryError> {
let entries: Vec<TermEntry> = entries.into_iter().collect();
let mut problems = Vec::new();
let mut claimed: HashMap<String, (String, String)> = HashMap::new();
let mut aliases: Vec<CompiledAlias> = Vec::new();
for (entry_index, entry) in entries.iter().enumerate() {
if entry.term.trim().is_empty() {
problems.push(Problem::EmptyTerm { entry: entry_index });
continue;
}
for (alias_index, alias) in entry.aliases.iter().enumerate() {
if alias.trim().is_empty() {
problems.push(Problem::EmptyAlias {
entry: entry_index,
alias: alias_index,
});
continue;
}
if *alias == entry.term {
problems.push(Problem::AliasEqualsTerm {
entry: entry_index,
alias: alias_index,
term: entry.term.clone(),
});
continue;
}
let folded = fold(alias, &policy).chars;
if folded.len() < MIN_ALIAS_CHARS {
problems.push(Problem::AliasTooShort {
entry: entry_index,
alias: alias_index,
text: alias.clone(),
folded_chars: folded.len(),
});
continue;
}
let key: String = folded.iter().collect();
match claimed.get(&key) {
Some((term, _)) if *term == entry.term => continue,
Some((term, written)) => {
problems.push(Problem::ConflictingAlias {
alias: written.clone(),
first_term: term.clone(),
second_term: entry.term.clone(),
});
continue;
}
None => {}
}
claimed.insert(key, (entry.term.clone(), alias.clone()));
aliases.push(CompiledAlias {
folded,
term: entry.term.clone(),
});
}
}
if !problems.is_empty() {
return Err(DictionaryError { problems });
}
aliases.sort_by_key(|alias| std::cmp::Reverse(alias.folded.len()));
Ok(Self { aliases, policy })
}
pub fn len(&self) -> usize {
self.aliases.len()
}
pub fn is_empty(&self) -> bool {
self.aliases.is_empty()
}
pub fn apply(&self, text: &str) -> (String, Vec<Correction>) {
let original: Vec<char> = text.chars().collect();
let folded = fold(text, &self.policy);
let mut corrections = Vec::new();
let mut plan: Vec<(usize, usize, &str)> = Vec::new();
let mut i = 0usize;
while i < folded.chars.len() {
let hit = self
.aliases
.iter()
.find(|alias| self.matches_at(&folded.chars, i, alias));
match hit {
Some(alias) => {
let end = i + alias.folded.len();
let from = folded.origin[i];
let to = folded.origin[end - 1] + 1;
corrections.push(Correction {
kind: CorrectionKind::DictionaryMatch,
original: original[from..to].iter().collect(),
replacement: alias.term.clone(),
span: Some(Span {
start: from,
end: to,
}),
});
plan.push((from, to, &alias.term));
i = end;
}
None => i += 1,
}
}
if plan.is_empty() {
return (text.to_string(), corrections);
}
let mut out = String::with_capacity(text.len());
let mut cursor = 0usize;
for (from, to, term) in plan {
out.extend(original[cursor..from].iter());
out.push_str(term);
cursor = to;
}
out.extend(original[cursor..].iter());
(out, corrections)
}
fn matches_at(&self, haystack: &[char], at: usize, alias: &CompiledAlias) -> bool {
let end = at + alias.folded.len();
if end > haystack.len() || haystack[at..end] != alias.folded[..] {
return false;
}
match self.policy.scope {
MatchScope::Substring => true,
MatchScope::WordBoundary => {
let before_ok = at == 0 || !haystack[at - 1].is_alphanumeric();
let after_ok = end == haystack.len() || !haystack[end].is_alphanumeric();
before_ok && after_ok
}
}
}
}
#[async_trait]
impl TextProcessor for TermDictionary {
async fn process(
&self,
text: &str,
_context: &ContextSnapshot,
) -> Result<ProcessResult, ProcessError> {
let (text, corrections) = self.apply(text);
Ok(ProcessResult { text, corrections })
}
}
#[cfg(test)]
mod tests {
use super::*;
fn entry(term: &str, aliases: &[&str]) -> TermEntry {
TermEntry {
term: term.into(),
aliases: aliases.iter().map(|a| (*a).into()).collect(),
}
}
fn ja(entries: Vec<TermEntry>) -> TermDictionary {
TermDictionary::new(entries, MatchPolicy::for_language(Language::Japanese))
.expect("valid dictionary")
}
fn en(entries: Vec<TermEntry>) -> TermDictionary {
TermDictionary::new(entries, MatchPolicy::for_language(Language::English))
.expect("valid dictionary")
}
#[test]
fn a_registered_term_replaces_its_alias() {
let dict = ja(vec![entry("typwrtr", &["タイプライター"])]);
let (text, corrections) = dict.apply("タイプライターで書いている");
assert_eq!(text, "typwrtrで書いている");
assert_eq!(corrections.len(), 1);
assert_eq!(corrections[0].kind, CorrectionKind::DictionaryMatch);
assert_eq!(corrections[0].original, "タイプライター");
assert_eq!(corrections[0].replacement, "typwrtr");
}
#[test]
fn text_with_no_match_is_returned_unchanged() {
let dict = ja(vec![entry("typwrtr", &["タイプライター"])]);
let (text, corrections) = dict.apply("今日はいい天気だ");
assert_eq!(text, "今日はいい天気だ");
assert!(corrections.is_empty());
}
#[test]
fn every_occurrence_is_replaced() {
let dict = ja(vec![entry("typwrtr", &["タイプライター"])]);
let (text, corrections) = dict.apply("タイプライターとタイプライター");
assert_eq!(text, "typwrtrとtypwrtr");
assert_eq!(corrections.len(), 2);
}
#[test]
fn the_span_locates_the_match_in_codepoints() {
let dict = ja(vec![entry("typwrtr", &["タイプライター"])]);
let (_, corrections) = dict.apply("私はタイプライターを使う");
let span = corrections[0].span.expect("a substitution knows where it was");
assert_eq!(span, Span { start: 2, end: 9 });
let chars: String = "私はタイプライターを使う"
.chars()
.skip(span.start)
.take(span.len())
.collect();
assert_eq!(chars, "タイプライター");
}
#[test]
fn spans_of_several_matches_are_in_text_order() {
let dict = ja(vec![entry("typwrtr", &["タイプライター"])]);
let (_, corrections) = dict.apply("タイプライターとタイプライター");
let spans: Vec<Span> = corrections.iter().filter_map(|c| c.span).collect();
assert_eq!(spans[0], Span { start: 0, end: 7 });
assert_eq!(spans[1], Span { start: 8, end: 15 });
}
#[test]
fn an_english_alias_does_not_fire_inside_a_longer_word() {
let dict = en(vec![entry("Category", &["cat"])]);
let (text, corrections) = dict.apply("concatenate the cat");
assert_eq!(text, "concatenate the Category");
assert_eq!(corrections.len(), 1);
}
#[test]
fn a_japanese_alias_fires_mid_string() {
let dict = ja(vec![entry("typwrtr", &["タイプライター"])]);
let (text, _) = dict.apply("これはタイプライターだ");
assert_eq!(text, "これはtypwrtrだ");
}
#[test]
fn punctuation_still_bounds_an_english_match() {
let dict = en(vec![entry("typwrtr", &["typewriter"])]);
let (text, _) = dict.apply("a typewriter, and a typewriter.");
assert_eq!(text, "a typwrtr, and a typwrtr.");
}
#[test]
fn japanese_folds_hiragana_and_katakana_together() {
let dict = ja(vec![entry("typwrtr", &["タイプライター"])]);
let (text, _) = dict.apply("たいぷらいたーを使う");
assert_eq!(text, "typwrtrを使う");
}
#[test]
fn japanese_does_not_fold_the_long_vowel_mark() {
let dict = ja(vec![entry("typwrtr", &["タイプライター"])]);
let (text, corrections) = dict.apply("タイプライタを使う");
assert_eq!(text, "タイプライタを使う");
assert!(corrections.is_empty(), "got {corrections:?}");
}
#[test]
fn full_width_ascii_folds_to_half_width() {
let dict = ja(vec![entry("euhadra", &["ユーハドラ"])]);
let (text, _) = dict.apply("GitHub のユーハドラ");
assert_eq!(text, "GitHub のeuhadra");
let dict = ja(vec![entry("OK", &["OKです"])]);
let (text, _) = dict.apply("OKですね");
assert_eq!(text, "OKね");
}
#[test]
fn case_is_folded_for_matching_but_never_for_the_replacement() {
let dict = en(vec![entry("TensorFlow", &["tensor flow"])]);
let (text, _) = dict.apply("Import Tensor Flow now");
assert_eq!(
text, "Import TensorFlow now",
"the term is emitted verbatim; only the match side folds"
);
}
#[test]
fn a_sentence_initial_capital_still_matches() {
let dict = en(vec![entry("typwrtr", &["typewriter"])]);
let (text, _) = dict.apply("Typewriter is the tool.");
assert_eq!(text, "typwrtr is the tool.");
}
#[test]
fn repeated_whitespace_collapses_for_matching() {
let dict = en(vec![entry("TensorFlow", &["tensor flow"])]);
let (text, corrections) = dict.apply("use tensor flow here");
assert_eq!(text, "use TensorFlow here");
assert_eq!(
corrections[0].original, "tensor flow",
"the correction reports what was actually in the text"
);
}
#[test]
fn a_policy_of_none_folds_nothing() {
let dict = TermDictionary::new(
vec![entry("typwrtr", &["typewriter"])],
MatchPolicy::none(),
)
.unwrap();
let (text, _) = dict.apply("Typewriter and typewriter");
assert_eq!(
text, "Typewriter and typwrtr",
"exact matching only; the capitalised one is a different string"
);
}
#[test]
fn the_longest_alias_wins_at_a_position() {
let dict = ja(vec![
entry("Type", &["タイプ"]),
entry("typwrtr", &["タイプライター"]),
]);
let (text, _) = dict.apply("タイプライターとタイプ");
assert_eq!(text, "typwrtrとType");
}
#[test]
fn a_replacement_is_not_rescanned() {
let dict = ja(vec![
entry("beta", &["アルファ"]),
entry("gamma", &["beta"]),
]);
let (text, corrections) = dict.apply("アルファ");
assert_eq!(text, "beta", "one pass: `beta` is output, not re-matched");
assert_eq!(corrections.len(), 1);
}
#[test]
fn every_problem_is_reported_at_once() {
let err = TermDictionary::new(
vec![
entry("", &["something"]),
entry("ok", &[""]),
entry("dup", &["e"]),
],
MatchPolicy::for_language(Language::English),
)
.unwrap_err();
assert_eq!(
err.problems.len(),
3,
"a caller fixing a dictionary should see all of it, got {:?}",
err.problems
);
}
#[test]
fn an_alias_shorter_than_the_minimum_is_refused() {
let err = TermDictionary::new(
vec![entry("euhadra", &["e"])],
MatchPolicy::for_language(Language::English),
)
.unwrap_err();
assert!(
matches!(
&err.problems[0],
Problem::AliasTooShort { folded_chars: 1, .. }
),
"got {:?}",
err.problems
);
}
#[test]
fn a_two_character_alias_is_allowed() {
let dict = en(vec![entry("Information Technology", &["IT"])]);
let (text, _) = dict.apply("the IT department");
assert_eq!(text, "the Information Technology department");
}
#[test]
fn two_terms_cannot_claim_the_same_alias() {
let err = TermDictionary::new(
vec![
entry("typwrtr", &["タイプライター"]),
entry("Typewriter Co.", &["タイプライター"]),
],
MatchPolicy::for_language(Language::Japanese),
)
.unwrap_err();
assert!(
matches!(&err.problems[0], Problem::ConflictingAlias { .. }),
"got {:?}",
err.problems
);
}
#[test]
fn aliases_that_fold_together_conflict() {
let err = TermDictionary::new(
vec![
entry("typwrtr", &["タイプライター"]),
entry("Typewriter Co.", &["たいぷらいたー"]),
],
MatchPolicy::for_language(Language::Japanese),
)
.unwrap_err();
assert!(
matches!(&err.problems[0], Problem::ConflictingAlias { .. }),
"got {:?}",
err.problems
);
}
#[test]
fn entries_sharing_a_term_are_merged() {
let dict = ja(vec![
entry("typwrtr", &["タイプライター"]),
entry("typwrtr", &["タイプライター", "typewriter"]),
]);
assert_eq!(dict.len(), 2, "the duplicate alias is folded away");
let (text, _) = dict.apply("タイプライターと typewriter");
assert_eq!(text, "typwrtrと typwrtr");
}
#[test]
fn an_alias_identical_to_its_term_is_refused() {
let err = TermDictionary::new(
vec![entry("typwrtr", &["typwrtr"])],
MatchPolicy::for_language(Language::English),
)
.unwrap_err();
assert!(
matches!(&err.problems[0], Problem::AliasEqualsTerm { .. }),
"got {:?}",
err.problems
);
}
#[test]
fn an_alias_differing_only_in_case_is_allowed() {
let dict = en(vec![entry("typwrtr", &["Typwrtr"])]);
let (text, _) = dict.apply("Typwrtr here");
assert_eq!(text, "typwrtr here");
}
#[test]
fn a_substring_relation_between_aliases_is_allowed() {
let dict = ja(vec![
entry("Type", &["タイプ"]),
entry("typwrtr", &["タイプライター"]),
]);
assert_eq!(dict.len(), 2);
}
#[test]
fn an_empty_dictionary_is_valid_and_changes_nothing() {
let dict = TermDictionary::new(
Vec::<TermEntry>::new(),
MatchPolicy::for_language(Language::English),
)
.unwrap();
assert!(dict.is_empty());
let (text, corrections) = dict.apply("nothing happens here");
assert_eq!(text, "nothing happens here");
assert!(corrections.is_empty());
}
#[tokio::test]
async fn it_runs_as_a_text_processor() {
let dict = ja(vec![entry("typwrtr", &["タイプライター"])]);
let result = dict
.process("タイプライターを使う", &ContextSnapshot::default())
.await
.unwrap();
assert_eq!(result.text, "typwrtrを使う");
assert_eq!(result.corrections.len(), 1);
}
}