use std::collections::{BTreeSet, HashMap};
use super::data::SpecialPiece;
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct RawToken {
pub text: String,
pub norm: Option<String>,
pub spacy: bool,
}
impl RawToken {
pub(crate) fn lex(text: &str) -> Self {
Self {
text: text.to_string(),
norm: None,
spacy: false,
}
}
}
#[derive(Debug, Default)]
pub(crate) struct Specials {
map: HashMap<String, Vec<SpecialPiece>>,
}
impl Specials {
pub(crate) fn new(rules: &[(String, Vec<SpecialPiece>)]) -> Self {
Self {
map: rules.iter().cloned().collect(),
}
}
pub(crate) fn contains(&self, chunk: &str) -> bool {
self.map.contains_key(chunk)
}
pub(crate) fn push(&self, chunk: &str, doc: &mut Vec<RawToken>) -> bool {
let Some(pieces) = self.map.get(chunk) else {
return false;
};
doc.extend(pieces.iter().map(|(orth, norm)| RawToken {
text: orth.clone(),
norm: norm.clone(),
spacy: false,
}));
true
}
}
#[derive(Debug, Default)]
pub(crate) struct SpecialMatcher {
by_first: HashMap<String, Vec<Vec<String>>>,
}
impl SpecialMatcher {
pub(crate) fn add(&mut self, pattern: Vec<String>) {
if let Some(first) = pattern.first() {
self.by_first
.entry(first.clone())
.or_default()
.push(pattern);
}
}
fn find_matches(&self, doc: &[RawToken]) -> Vec<(usize, usize)> {
let mut matches = Vec::new();
for start in 0..doc.len() {
let Some(patterns) = self.by_first.get(&doc[start].text) else {
continue;
};
for pattern in patterns {
let end = start + pattern.len();
if end <= doc.len()
&& doc[start..end]
.iter()
.zip(pattern)
.all(|(t, p)| t.text == *p)
{
matches.push((start, end));
}
}
}
matches
}
pub(crate) fn apply(&self, doc: Vec<RawToken>, specials: &Specials) -> Vec<RawToken> {
let matches = self.find_matches(&doc);
if matches.is_empty() {
return doc;
}
let spans = filter_spans(matches);
let mut out = Vec::with_capacity(doc.len());
let mut span_iter = spans.into_iter().peekable();
let mut i = 0;
while i < doc.len() {
let Some(&(start, end)) = span_iter.peek().filter(|(s, _)| *s == i) else {
out.push(doc[i].clone());
i += 1;
continue;
};
span_iter.next();
let text = span_text(&doc[start..end]);
let before = out.len();
if specials.push(&text, &mut out) {
if let Some(last) = out[before..].last_mut() {
last.spacy = doc[end - 1].spacy;
}
} else {
out.extend_from_slice(&doc[start..end]);
}
i = end;
}
out
}
}
fn span_text(tokens: &[RawToken]) -> String {
let mut text = String::new();
for (k, t) in tokens.iter().enumerate() {
text.push_str(&t.text);
if t.spacy && k + 1 < tokens.len() {
text.push(' ');
}
}
text
}
fn filter_spans(mut spans: Vec<(usize, usize)>) -> Vec<(usize, usize)> {
spans.sort_by(|a, b| {
let (la, lb) = (a.1 - a.0, b.1 - b.0);
la.cmp(&lb).then(b.0.cmp(&a.0))
});
let mut seen = BTreeSet::new();
let mut kept = Vec::new();
for &(start, end) in spans.iter().rev() {
if !seen.contains(&start) && !seen.contains(&(end - 1)) {
kept.push((start, end));
}
seen.extend(start..end);
}
kept.sort_by_key(|&(start, _)| start);
kept
}
#[cfg(test)]
mod tests {
use super::*;
fn toks(texts: &[(&str, bool)]) -> Vec<RawToken> {
texts
.iter()
.map(|&(t, spacy)| RawToken {
text: t.to_string(),
norm: None,
spacy,
})
.collect()
}
#[test]
fn filter_prefers_longest_then_leftmost() {
assert_eq!(filter_spans(vec![(0, 2), (1, 3), (0, 1)]), vec![(0, 2)]);
assert_eq!(
filter_spans(vec![(2, 3), (0, 1), (1, 2)]),
vec![(0, 1), (1, 2), (2, 3)]
);
assert_eq!(filter_spans(vec![(0, 3), (1, 4), (3, 4)]), vec![(0, 3)]);
}
#[test]
fn replaces_only_exact_text() {
let specials = Specials::new(&[(":)".to_string(), vec![(":)".to_string(), None)])]);
let mut m = SpecialMatcher::default();
m.add(vec![":".to_string(), ")".to_string()]);
let joined = m.apply(toks(&[("a", true), (":", false), (")", true)]), &specials);
assert_eq!(joined, toks(&[("a", true), (":)", true)]));
let spaced = toks(&[(":", true), (")", false)]);
assert_eq!(m.apply(spaced.clone(), &specials), spaced);
}
}