use std::collections::HashSet;
use tantivy::tokenizer::{Token, TokenStream, Tokenizer};
use unicode_normalization::UnicodeNormalization;
pub const DEFAULT_MIN_GRAM: usize = 2;
pub const DEFAULT_MAX_GRAM: usize = 5;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct NgramConfig {
pub min_gram: usize,
pub max_gram: usize,
pub ascii_folding: bool,
}
impl Default for NgramConfig {
fn default() -> Self {
NgramConfig {
min_gram: DEFAULT_MIN_GRAM,
max_gram: DEFAULT_MAX_GRAM,
ascii_folding: true,
}
}
}
pub fn ascii_fold_lower(input: &str, ascii_folding: bool) -> String {
let lowered = input.to_lowercase();
if !ascii_folding {
return lowered;
}
let mut out = String::with_capacity(lowered.len());
for ch in lowered.nfkd() {
if is_combining_mark(ch) {
continue;
}
match ch {
'ł' => out.push('l'),
'đ' => out.push('d'),
'ø' => out.push('o'),
'æ' => out.push_str("ae"),
'œ' => out.push_str("oe"),
'ß' => out.push_str("ss"),
'þ' => out.push_str("th"),
'ð' => out.push('d'),
_ => out.push(ch),
}
}
out
}
fn is_combining_mark(ch: char) -> bool {
matches!(ch as u32,
0x0300..=0x036F | 0x1AB0..=0x1AFF | 0x1DC0..=0x1DFF | 0x20D0..=0x20FF | 0xFE20..=0xFE2F)
}
fn is_word_char(ch: char) -> bool {
ch.is_alphanumeric()
}
fn split_words(folded: &str) -> Vec<String> {
folded
.split(|c: char| !is_word_char(c))
.filter(|w| !w.is_empty())
.map(|w| w.to_string())
.collect()
}
pub fn ngrams_of_word(word: &str, cfg: &NgramConfig, out: &mut Vec<String>) {
let chars: Vec<char> = word.chars().collect();
let len = chars.len();
if len == 0 {
return;
}
if len < cfg.min_gram {
out.push(chars.iter().collect());
return;
}
for n in cfg.min_gram..=cfg.max_gram {
if n > len {
break;
}
for start in 0..=(len - n) {
let gram: String = chars[start..start + n].iter().collect();
out.push(gram);
}
}
}
pub fn ngram_tokens(text: &str, cfg: &NgramConfig) -> Vec<String> {
let folded = ascii_fold_lower(text, cfg.ascii_folding);
let mut out = Vec::new();
for word in split_words(&folded) {
ngrams_of_word(&word, cfg, &mut out);
}
out
}
pub fn ngram_set(text: &str, cfg: &NgramConfig) -> HashSet<String> {
ngram_tokens(text, cfg).into_iter().collect()
}
pub fn ngram_match(haystack: &str, needle: &str, cfg: &NgramConfig) -> bool {
let needle_grams = ngram_set(needle, cfg);
if needle_grams.is_empty() {
return false;
}
let hay_grams = ngram_set(haystack, cfg);
needle_grams.iter().all(|g| hay_grams.contains(g))
}
pub fn pack_typmod(parts: &[String]) -> Result<i32, String> {
let min: usize = parts
.first()
.map(|s| s.trim().parse().map_err(|_| "min_gram must be an integer".to_string()))
.unwrap_or(Ok(DEFAULT_MIN_GRAM))?;
let max: usize = parts
.get(1)
.map(|s| s.trim().parse().map_err(|_| "max_gram must be an integer".to_string()))
.unwrap_or(Ok(DEFAULT_MAX_GRAM))?;
let mut ascii_folding = true;
if let Some(opts) = parts.get(2) {
for kv in opts.split([';', ',']) {
let mut it = kv.splitn(2, '=');
let k = it.next().unwrap_or("").trim().to_lowercase();
let v = it.next().unwrap_or("").trim().to_lowercase();
if k == "ascii_folding" {
ascii_folding = v != "false" && v != "0" && v != "off";
}
}
}
if min < 1 || min > 255 {
return Err(format!("min_gram out of range: {min}"));
}
if max < min || max > 255 {
return Err(format!("max_gram must be between min_gram and 255: {max}"));
}
Ok(((min as i32) << 16) | ((max as i32) << 8) | (ascii_folding as i32))
}
pub fn unpack_typmod(tm: i32) -> NgramConfig {
if tm < 0 {
return NgramConfig::default();
}
NgramConfig {
min_gram: ((tm >> 16) & 0xFF) as usize,
max_gram: ((tm >> 8) & 0xFF) as usize,
ascii_folding: (tm & 1) != 0,
}
}
#[derive(Clone)]
pub struct AsciiNgramTokenizer {
pub cfg: NgramConfig,
}
impl AsciiNgramTokenizer {
pub fn new(cfg: NgramConfig) -> Self {
AsciiNgramTokenizer { cfg }
}
}
pub struct VecTokenStream {
tokens: Vec<Token>,
idx: usize,
}
impl TokenStream for VecTokenStream {
fn advance(&mut self) -> bool {
if self.idx < self.tokens.len() {
self.idx += 1;
true
} else {
false
}
}
fn token(&self) -> &Token {
&self.tokens[self.idx - 1]
}
fn token_mut(&mut self) -> &mut Token {
&mut self.tokens[self.idx - 1]
}
}
impl Tokenizer for AsciiNgramTokenizer {
type TokenStream<'a> = VecTokenStream;
fn token_stream<'a>(&'a mut self, text: &'a str) -> Self::TokenStream<'a> {
let grams = ngram_tokens(text, &self.cfg);
let mut tokens = Vec::with_capacity(grams.len());
for (position, text) in grams.into_iter().enumerate() {
let len = text.len();
tokens.push(Token {
offset_from: 0,
offset_to: len,
position,
text,
position_length: 1,
});
}
VecTokenStream { tokens, idx: 0 }
}
}
#[cfg(test)]
mod tests {
use super::*;
fn cfg() -> NgramConfig {
NgramConfig::default()
}
#[test]
fn folds_lithuanian_diacritics() {
assert_eq!(ascii_fold_lower("Ąžuolas", true), "azuolas");
assert_eq!(ascii_fold_lower("ČĘĖĮŠŲŪŽ", true), "ceeisuuz");
assert_eq!(ascii_fold_lower("Kavinė", true), "kavine");
}
#[test]
fn folds_common_latin_extras() {
assert_eq!(ascii_fold_lower("Łódź", true), "lodz");
assert_eq!(ascii_fold_lower("Straße", true), "strasse");
assert_eq!(ascii_fold_lower("Œuvre", true), "oeuvre");
}
#[test]
fn lowercases_when_folding_disabled() {
assert_eq!(ascii_fold_lower("Ąžuolas", false), "ąžuolas");
}
#[test]
fn ngrams_basic_word() {
let mut out = Vec::new();
ngrams_of_word("kava", &cfg(), &mut out);
assert!(out.contains(&"ka".to_string()));
assert!(out.contains(&"av".to_string()));
assert!(out.contains(&"va".to_string()));
assert!(out.contains(&"kava".to_string()));
assert!(out.contains(&"kav".to_string()));
assert!(out.contains(&"ava".to_string()));
assert!(!out.contains(&"k".to_string()));
}
#[test]
fn short_word_below_min_kept_whole() {
let mut out = Vec::new();
ngrams_of_word("a", &cfg(), &mut out);
assert_eq!(out, vec!["a".to_string()]);
}
#[test]
fn max_gram_capped_at_five() {
let mut out = Vec::new();
ngrams_of_word("aparatas", &cfg(), &mut out);
assert!(out.iter().all(|g| g.chars().count() <= 5));
assert!(out.contains(&"apara".to_string()));
assert!(!out.contains(&"aparat".to_string()));
}
#[test]
fn tokens_split_on_whitespace_and_punct() {
let toks = ngram_set("Kavos aparatas", &cfg());
assert!(toks.contains("ka"));
assert!(toks.contains("ap"));
assert!(!toks.contains("s "));
assert!(!toks.contains("s a"));
}
#[test]
fn match_accent_insensitive() {
assert!(ngram_match("Ąžuolų baldai", "azuol", &cfg()));
assert!(ngram_match("Kavos aparatas", "kavos", &cfg()));
assert!(ngram_match("Kavos aparatas", "apar", &cfg()));
}
#[test]
fn match_is_conjunctive_substring() {
assert!(!ngram_match("Telefonas Samsung", "canon", &cfg()));
assert!(!ngram_match("Kavos aparatas", "kava", &cfg()));
assert!(ngram_match("Canon EOS fotoaparatas", "canon", &cfg()));
assert!(ngram_match("Sony PlayStation 5", "playstation", &cfg()));
}
#[test]
fn match_case_insensitive() {
assert!(ngram_match("HELLO World", "hello", &cfg()));
}
#[test]
fn no_match_disjoint() {
assert!(!ngram_match("kavos aparatas", "zzzz", &cfg()));
}
#[test]
fn two_char_query_matches_substring() {
assert!(ngram_match("kaina", "ka", &cfg()));
assert!(ngram_match("sukamasis", "ka", &cfg()));
}
#[test]
fn typmod_roundtrip() {
let tm = pack_typmod(&[
"2".to_string(),
"5".to_string(),
"ascii_folding=true".to_string(),
])
.unwrap();
let cfg = unpack_typmod(tm);
assert_eq!(cfg.min_gram, 2);
assert_eq!(cfg.max_gram, 5);
assert!(cfg.ascii_folding);
assert!(tm >= 0);
}
#[test]
fn typmod_ascii_folding_false() {
let tm = pack_typmod(&[
"3".to_string(),
"4".to_string(),
"ascii_folding=false".to_string(),
])
.unwrap();
let cfg = unpack_typmod(tm);
assert_eq!(cfg.min_gram, 3);
assert_eq!(cfg.max_gram, 4);
assert!(!cfg.ascii_folding);
}
#[test]
fn typmod_rejects_bad_range() {
assert!(pack_typmod(&["5".to_string(), "2".to_string(), "".to_string()]).is_err());
assert!(pack_typmod(&["0".to_string(), "5".to_string(), "".to_string()]).is_err());
}
#[test]
fn tantivy_tokenizer_emits_same_tokens() {
let mut tk = AsciiNgramTokenizer::new(cfg());
let mut stream = tk.token_stream("Kava");
let mut emitted = Vec::new();
while stream.advance() {
emitted.push(stream.token().text.clone());
}
let mut expected = Vec::new();
ngrams_of_word("kava", &cfg(), &mut expected);
assert_eq!(emitted, expected);
}
}