use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use crate::accessor::compile_path;
use crate::error::Result;
use crate::normalize::{normalize_text, normalize_text_cow};
use crate::value::Value;
use super::trigram::{self, Trigram};
use super::{rank, score_item, FuzzyMode, SearchSpec};
pub struct TrigramIndex {
postings: HashMap<Trigram, Vec<u32>>,
len: usize,
field_cache: Mutex<HashMap<String, Arc<Vec<Vec<Trigram>>>>>,
}
impl TrigramIndex {
#[must_use]
pub fn build(items: &[Value]) -> Self {
let mut postings: HashMap<Trigram, Vec<u32>> = HashMap::new();
let mut buf: Vec<Trigram> = Vec::new();
for (index, item) in items.iter().enumerate() {
buf.clear();
collect_string_trigrams(item, &mut buf);
buf.sort_unstable();
buf.dedup();
for tg in &buf {
postings.entry(*tg).or_default().push(index as u32);
}
}
Self {
postings,
len: items.len(),
field_cache: Mutex::new(HashMap::new()),
}
}
fn field_trigrams(&self, items: &[Value], field: &str) -> Arc<Vec<Vec<Trigram>>> {
let mut cache = self.field_cache.lock().expect("field cache poisoned");
if let Some(cached) = cache.get(field) {
return Arc::clone(cached);
}
let built: Vec<Vec<Trigram>> = items
.iter()
.map(|item| match item {
Value::Map(map) => match map.get(field) {
Some(Value::Str(s)) => trigram::trigrams(&normalize_text_cow(s)),
_ => Vec::new(),
},
_ => Vec::new(),
})
.collect();
let arc = Arc::new(built);
cache.insert(field.to_owned(), Arc::clone(&arc));
arc
}
pub(super) fn candidates(&self, query: &[Trigram]) -> Vec<u32> {
let mut seen = vec![false; self.len];
let mut count = 0usize;
for tg in query {
if let Some(list) = self.postings.get(tg) {
for &i in list {
let slot = &mut seen[i as usize];
if !*slot {
*slot = true;
count += 1;
}
}
}
}
let mut out = Vec::with_capacity(count);
for (i, &hit) in seen.iter().enumerate() {
if hit {
out.push(i as u32);
}
}
out
}
}
fn collect_string_trigrams(value: &Value, out: &mut Vec<Trigram>) {
match value {
Value::Str(s) => out.extend(trigram::trigrams(&normalize_text_cow(s))),
Value::Map(map) => map.values().for_each(|v| collect_string_trigrams(v, out)),
Value::List(items) => items.iter().for_each(|v| collect_string_trigrams(v, out)),
_ => {}
}
}
pub fn search_with_index(
items: &[Value],
spec: &SearchSpec,
index: &TrigramIndex,
) -> Result<Vec<usize>> {
if spec.query.trim().chars().count() < spec.min_length {
return Ok((0..items.len()).collect());
}
let paths: Vec<Vec<String>> = spec
.fields
.iter()
.map(|f| compile_path(f))
.collect::<Result<_>>()?;
if spec.fuzzy == FuzzyMode::Exact {
return Ok(super::exact_search(items, &paths, spec));
}
let query = trigram::trigrams(&normalize_text(&spec.query));
if query.is_empty() {
return Ok((0..items.len()).collect());
}
let caches: Vec<Option<Arc<Vec<Vec<Trigram>>>>> = paths
.iter()
.map(|path| (path.len() == 1).then(|| index.field_trigrams(items, &path[0])))
.collect();
let mut scored: Vec<(i64, usize)> = Vec::new();
for candidate in index.candidates(&query) {
let i = candidate as usize;
let score = score_item(items, i, &paths, &caches, &query, spec);
if score > 0 {
scored.push((score, i));
}
}
Ok(rank(scored, spec.max_results))
}