use std::collections::HashMap;
type Trigram = u32;
#[derive(Default)]
pub struct TrigramIndex {
postings: HashMap<Trigram, Vec<usize>>,
node_count: usize,
}
impl Clone for TrigramIndex {
fn clone(&self) -> Self {
Self {
postings: self.postings.clone(),
node_count: self.node_count,
}
}
}
impl TrigramIndex {
pub fn build(haystacks: &[String]) -> Self {
let mut postings: HashMap<Trigram, Vec<usize>> = HashMap::new();
for (node_idx, haystack) in haystacks.iter().enumerate() {
let lower = haystack.to_lowercase();
let bytes = lower.as_bytes();
if bytes.len() < 3 {
continue;
}
for window in bytes.windows(3) {
let tri = pack_trigram(window);
postings.entry(tri).or_default().push(node_idx);
}
}
for indices in postings.values_mut() {
indices.sort_unstable();
indices.dedup();
}
TrigramIndex {
postings,
node_count: haystacks.len(),
}
}
pub fn candidates_for_term(&self, term: &str) -> Option<Vec<usize>> {
let lower = term.to_lowercase();
let bytes = lower.as_bytes();
if bytes.len() < 3 {
return None; }
let trigrams: Vec<Trigram> = bytes.windows(3).map(pack_trigram).collect();
if trigrams.is_empty() {
return None;
}
let mut candidates: Option<Vec<usize>> = None;
for tri in &trigrams {
let posting = match self.postings.get(tri) {
Some(p) => p,
None => return Some(Vec::new()), };
candidates = Some(match candidates.take() {
None => posting.clone(),
Some(existing) => intersect_sorted(&existing, posting),
});
if candidates.as_ref().is_none_or(std::vec::Vec::is_empty) {
return Some(Vec::new());
}
}
candidates
}
pub fn candidates(&self, terms: &[String]) -> Option<Vec<usize>> {
let mut union: Option<Vec<usize>> = None;
for term in terms {
match self.candidates_for_term(term) {
None => return None, Some(term_candidates) => {
union = Some(match union.take() {
None => term_candidates,
Some(existing) => union_sorted(&existing, &term_candidates),
});
}
}
}
union
}
}
fn pack_trigram(bytes: &[u8]) -> Trigram {
((bytes[0] as u32) << 16) | ((bytes[1] as u32) << 8) | (bytes[2] as u32)
}
fn intersect_sorted(a: &[usize], b: &[usize]) -> Vec<usize> {
let mut result = Vec::with_capacity(a.len().min(b.len()));
let (mut i, mut j) = (0, 0);
while i < a.len() && j < b.len() {
match a[i].cmp(&b[j]) {
std::cmp::Ordering::Equal => {
result.push(a[i]);
i += 1;
j += 1;
}
std::cmp::Ordering::Less => i += 1,
std::cmp::Ordering::Greater => j += 1,
}
}
result
}
fn union_sorted(a: &[usize], b: &[usize]) -> Vec<usize> {
let mut result = Vec::with_capacity(a.len() + b.len());
let (mut i, mut j) = (0, 0);
while i < a.len() && j < b.len() {
match a[i].cmp(&b[j]) {
std::cmp::Ordering::Equal => {
result.push(a[i]);
i += 1;
j += 1;
}
std::cmp::Ordering::Less => {
result.push(a[i]);
i += 1;
}
std::cmp::Ordering::Greater => {
result.push(b[j]);
j += 1;
}
}
}
while i < a.len() {
result.push(a[i]);
i += 1;
}
while j < b.len() {
result.push(b[j]);
j += 1;
}
result
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn trigram_index_narrows_candidates() {
let haystacks = vec![
"retry policy backoff".to_string(), "task scheduler queue".to_string(), "dead letter queue handling".to_string(), "jitter exponential".to_string(), ];
let index = TrigramIndex::build(&haystacks);
let cands = index.candidates_for_term("retry").unwrap();
assert!(cands.contains(&0)); assert!(!cands.contains(&1));
let cands = index.candidates_for_term("scheduler").unwrap();
assert!(cands.contains(&1));
assert!(!cands.contains(&0));
assert_eq!(index.candidates_for_term("go"), None);
}
#[test]
fn trigram_index_candidates_for_multiple_terms() {
let haystacks = vec![
"retry policy backoff".to_string(),
"task scheduler queue".to_string(),
"dead letter queue".to_string(),
];
let index = TrigramIndex::build(&haystacks);
let cands = index.candidates(&["retry".to_string(), "queue".to_string()]);
if let Some(cands) = cands {
assert!(cands.contains(&0)); assert!(cands.contains(&1)); assert!(cands.contains(&2)); } else {
panic!("expected candidates");
}
}
#[test]
fn nonexistent_term_returns_empty() {
let haystacks = vec!["hello world".to_string()];
let index = TrigramIndex::build(&haystacks);
let cands = index.candidates_for_term("zzzzz").unwrap();
assert!(cands.is_empty());
}
}