use crate::display;
use crate::message::ConversationMessage;
use crate::text;
use std::collections::{HashMap, HashSet};
struct Hit {
entry_id: String,
score: f64,
}
const BM25_K1: f64 = 1.5;
const BM25_B: f64 = 0.75;
const PROXIMITY_FACTOR: f64 = 2.0;
#[allow(clippy::cast_precision_loss)]
fn usize_as_f64(value: usize) -> f64 {
value as f64
}
fn min_span_all(positions: &[&[usize]]) -> usize {
let mut all: Vec<(usize, usize)> = Vec::new();
for (term_idx, pos_list) in positions.iter().enumerate() {
for &pos in *pos_list {
all.push((pos, term_idx));
}
}
all.sort_by_key(|(pos, _)| *pos);
let num_terms = positions.len();
let mut counts = vec![0usize; num_terms];
let mut matched = 0;
let mut left = 0;
let mut min_span = usize::MAX;
for right in 0..all.len() {
let term_idx = all[right].1;
if counts[term_idx] == 0 {
matched += 1;
}
counts[term_idx] += 1;
while matched == num_terms {
let span = all[right].0 - all[left].0 + 1;
min_span = min_span.min(span);
let left_term = all[left].1;
counts[left_term] -= 1;
if counts[left_term] == 0 {
matched -= 1;
}
left += 1;
}
}
min_span
}
fn adjacent_fallback(
messages: &[ConversationMessage],
tokenized: &[Vec<String>],
terms: &[String],
page: usize,
page_size: usize,
) -> Vec<String> {
let hits: Vec<Hit> = tokenized
.iter()
.enumerate()
.filter_map(|(i, msg_words)| {
let word_set: HashSet<&str> =
msg_words.iter().map(std::string::String::as_str).collect();
let matching_pairs = terms
.windows(2)
.filter(|pair| {
word_set.contains(pair[0].as_str()) && word_set.contains(pair[1].as_str())
})
.count();
if matching_pairs == 0 {
return None;
}
Some(Hit {
entry_id: messages[i].entry_id.clone(),
score: usize_as_f64(matching_pairs),
})
})
.collect();
sort_and_page(hits, page, page_size)
}
fn sort_and_page(mut hits: Vec<Hit>, page: usize, page_size: usize) -> Vec<String> {
sort_by_score_desc(&mut hits);
let start = (page.saturating_sub(1)) * page_size;
hits[start..std::cmp::min(start + page_size, hits.len())]
.iter()
.map(|h| h.entry_id.clone())
.collect()
}
fn sort_by_score_desc(hits: &mut [Hit]) {
hits.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
}
pub fn grep(messages: &[ConversationMessage], pattern: &str) -> Vec<String> {
let pattern_lc = pattern.to_ascii_lowercase();
messages
.iter()
.filter_map(|msg| {
let searchable = display::searchable_text(msg);
text::glob_search(&pattern_lc, &searchable.to_ascii_lowercase())
.then(|| msg.entry_id.clone())
})
.collect()
}
struct TermSearchIndex {
tokenized: Vec<Vec<String>>,
doc_lengths: Vec<usize>,
df: HashMap<String, usize>,
avgdl: f64,
n: f64,
}
fn normalized_terms(query_str: &str) -> Vec<String> {
text::split_words(query_str)
.into_iter()
.filter(|w| w.len() > 1 && !text::is_stop_word(w))
.collect()
}
fn build_term_search_index(messages: &[ConversationMessage], terms: &[String]) -> TermSearchIndex {
let n = usize_as_f64(messages.len());
let mut df: HashMap<String, usize> = HashMap::new();
let mut tokenized: Vec<Vec<String>> = Vec::with_capacity(messages.len());
let mut doc_lengths: Vec<usize> = Vec::with_capacity(messages.len());
let mut total_words: usize = 0;
for msg in messages {
let searchable = display::searchable_text(msg);
let msg_words: Vec<String> = text::split_words(&searchable)
.into_iter()
.map(|w| w.to_ascii_lowercase())
.collect();
let dl = msg_words.len();
doc_lengths.push(dl);
total_words += dl;
let word_set: HashSet<&str> = msg_words.iter().map(String::as_str).collect();
for term in terms {
if word_set.contains(term.as_str()) {
*df.entry(term.clone()).or_insert(0) += 1;
}
}
tokenized.push(msg_words);
}
let avgdl = if n > 0.0 {
usize_as_f64(total_words) / n
} else {
0.0
};
TermSearchIndex {
tokenized,
doc_lengths,
df,
avgdl,
n,
}
}
fn positions_by_term<'a>(
msg_words: &'a [String],
term_set: &HashSet<&str>,
) -> HashMap<&'a str, Vec<usize>> {
let mut positions_by_term: HashMap<&str, Vec<usize>> = HashMap::new();
for (pos, word) in msg_words.iter().enumerate() {
let word = word.as_str();
if term_set.contains(word) {
positions_by_term.entry(word).or_default().push(pos);
}
}
positions_by_term
}
fn score_message_terms(
index: &TermSearchIndex,
terms: &[String],
positions_by_term: &HashMap<&str, Vec<usize>>,
doc_index: usize,
) -> f64 {
let dl = usize_as_f64(index.doc_lengths[doc_index]);
let len_norm = 1.0 - BM25_B + BM25_B * (dl / index.avgdl.max(1.0));
let mut score = 0.0;
let mut term_positions: Vec<&[usize]> = Vec::with_capacity(terms.len());
for term in terms {
let Some(positions) = positions_by_term.get(term.as_str()) else {
continue;
};
let tf = positions.len();
term_positions.push(positions.as_slice());
let df_val = usize_as_f64(*index.df.get(term).unwrap_or(&1));
let idf = (1.0 + (index.n - df_val + 0.5) / (df_val + 0.5)).ln();
let tf = usize_as_f64(tf);
let tf_score = (tf * (BM25_K1 + 1.0)) / (tf + BM25_K1 * len_norm);
score += idf * tf_score;
}
if score > 0.0 && term_positions.len() > 1 {
let min_span = min_span_all(&term_positions).max(1);
score *= 1.0 + PROXIMITY_FACTOR / (1.0 + usize_as_f64(min_span));
}
score
}
fn query_terms(
messages: &[ConversationMessage],
terms: &[String],
page: usize,
page_size: usize,
) -> Vec<String> {
let index = build_term_search_index(messages, terms);
let term_set: HashSet<&str> = terms.iter().map(String::as_str).collect();
let hits: Vec<Hit> = messages
.iter()
.enumerate()
.filter_map(|(i, msg)| {
let positions_by_term = positions_by_term(&index.tokenized[i], &term_set);
let score = score_message_terms(&index, terms, &positions_by_term, i);
(score > 0.0).then(|| Hit {
entry_id: msg.entry_id.clone(),
score,
})
})
.collect();
let total = hits.len();
let result = sort_and_page(hits, page, page_size);
if result.is_empty() && total == 0 && terms.len() >= 2 {
return adjacent_fallback(messages, &index.tokenized, terms, page, page_size);
}
result
}
pub fn query(
messages: &[ConversationMessage],
query_str: &str,
page: usize,
page_size: usize,
) -> Vec<String> {
let terms = normalized_terms(query_str);
if terms.is_empty() {
return Vec::new();
}
query_terms(messages, &terms, page, page_size)
}