hayson_engine/registry/

fulltext_index.rs

use std::collections::{HashMap, HashSet};
use std::sync::{Mutex, MutexGuard, PoisonError};

use super::tokenize::{parse_search_terms, SearchTerm};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SearchMode {
    All,
    Any,
}

#[derive(Debug, Default)]
pub struct CorpusDelta {
    pub added: HashMap<String, Vec<String>>,
    pub removed: HashSet<String>,
}

struct StagedDocument {
    length: usize,
    frequencies: HashMap<String, u32>,
}

struct Corpus {
    documents: i64,
    total_tokens: i64,
    staged: HashMap<String, StagedDocument>,
    hidden: HashSet<String>,
}

struct Posting {
    key: String,
    frequency: u32,
}

const K1: f64 = 1.2;
const B: f64 = 0.75;
const SPLICED_TOKENS: usize = 16;

fn distinct_terms(query: &str) -> Vec<SearchTerm> {
    let mut seen = HashSet::new();
    let mut terms = Vec::new();
    for term in parse_search_terms(query) {
        let identity = if term.prefix {
            format!("{}*", term.token)
        } else {
            term.token.clone()
        };
        if seen.insert(identity) {
            terms.push(term);
        }
    }
    terms
}

/// Whether a document made of `tokens` is one [`FulltextIndex::search`] would answer `query` with —
/// the membership half of a search, decided on one document alone. What a caller holding a record
/// the index has never been told about (a unit of work's staged put) has to ask instead of looking
/// the key up in the postings, which cannot know it. Relevance is not part of the answer: a staged
/// record joins the result set, it does not join the ranking.
#[allow(dead_code)]
pub fn matches_terms(tokens: &[String], query: &str, mode: SearchMode) -> bool {
    let terms = distinct_terms(query);
    if terms.is_empty() {
        return false;
    }
    let mut hits = 0;
    for term in &terms {
        let hit = tokens.iter().any(|token| {
            if term.prefix {
                token.starts_with(&term.token)
            } else {
                token == &term.token
            }
        });
        if hit {
            hits += 1;
        } else if mode == SearchMode::All {
            return false;
        }
    }
    hits > 0
}

/// A poisoned cache lock is recovered rather than propagated: the two vectors behind it are a
/// derived prefix cache, so the worst a panic mid-merge can leave behind is a prefix search that
/// misses a token until it is indexed again — never a wrong record, and never a wrong count.
fn cached<T>(cell: &Mutex<T>) -> MutexGuard<'_, T> {
    cell.lock().unwrap_or_else(PoisonError::into_inner)
}

// NOTE: vocabulary/fresh_tokens are a lazy sorted-prefix cache, mutated by `&self` searches.
#[derive(Default)]
pub struct FulltextIndex {
    postings: HashMap<String, HashMap<String, u32>>,
    lengths: HashMap<String, usize>,
    vocabulary: Mutex<Vec<String>>,
    fresh_tokens: Mutex<Vec<String>>,
    total_tokens: i64,
}

impl FulltextIndex {
    pub fn new() -> Self {
        FulltextIndex::default()
    }

    #[allow(dead_code)]
    pub fn documents(&self) -> usize {
        self.lengths.len()
    }

    pub fn index(&mut self, key: &str, tokens: &[String]) {
        self.remove(key);
        if tokens.is_empty() {
            return;
        }
        for token in tokens {
            match self.postings.get_mut(token) {
                Some(bucket) => {
                    *bucket.entry(key.to_string()).or_insert(0) += 1;
                }
                None => {
                    let mut bucket = HashMap::new();
                    bucket.insert(key.to_string(), 1);
                    self.postings.insert(token.clone(), bucket);
                    cached(&self.fresh_tokens).push(token.clone());
                }
            }
        }
        self.lengths.insert(key.to_string(), tokens.len());
        self.total_tokens += tokens.len() as i64;
    }

    pub fn remove(&mut self, key: &str) {
        let Some(length) = self.lengths.remove(key) else { return };
        self.total_tokens -= length as i64;
        let mut emptied = HashSet::new();
        self.postings.retain(|token, bucket| {
            if bucket.remove(key).is_some() && bucket.is_empty() {
                emptied.insert(token.clone());
                false
            } else {
                true
            }
        });
        if emptied.is_empty() {
            return;
        }
        cached(&self.vocabulary).retain(|token| !emptied.contains(token));
        let mut fresh = cached(&self.fresh_tokens);
        if !fresh.is_empty() {
            fresh.retain(|token| !emptied.contains(token));
        }
    }

    pub fn search(&self, query: &str, mode: SearchMode, delta: Option<&CorpusDelta>) -> HashMap<String, f64> {
        let terms = distinct_terms(query);
        let mut scores = HashMap::new();
        if terms.is_empty() {
            return scores;
        }

        let corpus = self.corpus_with(delta);
        if corpus.documents == 0 {
            return scores;
        }

        let average_length = corpus.total_tokens as f64 / corpus.documents as f64;
        let mut hits: HashMap<String, u32> = HashMap::new();

        for term in &terms {
            self.accumulate_term(term, &corpus, average_length, &mut scores, &mut hits);
        }

        if mode == SearchMode::All {
            let terms_len = terms.len() as u32;
            scores.retain(|key, _| hits.get(key).copied().unwrap_or(0) >= terms_len);
        }
        scores
    }

    fn accumulate_term(
        &self,
        term: &SearchTerm,
        corpus: &Corpus,
        average_length: f64,
        scores: &mut HashMap<String, f64>,
        hits: &mut HashMap<String, u32>,
    ) {
        let tokens = if term.prefix {
            self.expand(&term.token, corpus)
        } else {
            vec![term.token.clone()]
        };
        if tokens.len() == 1 {
            let mut add = |key: String, score: f64| {
                *scores.entry(key.clone()).or_insert(0.0) += score;
                *hits.entry(key).or_insert(0) += 1;
            };
            self.score_token(&tokens[0], corpus, average_length, &mut add);
            return;
        }

        let mut best: HashMap<String, f64> = HashMap::new();
        for token in &tokens {
            let mut take = |key: String, score: f64| {
                let current = best.get(&key).copied();
                if current.is_none_or(|value| score > value) {
                    best.insert(key, score);
                }
            };
            self.score_token(token, corpus, average_length, &mut take);
        }
        for (key, score) in best {
            *scores.entry(key.clone()).or_insert(0.0) += score;
            *hits.entry(key).or_insert(0) += 1;
        }
    }

    fn score_token(&self, token: &str, corpus: &Corpus, average_length: f64, take: &mut dyn FnMut(String, f64)) {
        let postings = self.postings_of(token, corpus);
        if postings.is_empty() {
            return;
        }
        let documents = corpus.documents as f64;
        let df = postings.len() as f64;
        let idf = (1.0 + (documents - df + 0.5) / (df + 0.5)).ln();
        for posting in &postings {
            let length = corpus
                .staged
                .get(&posting.key)
                .map(|doc| doc.length as f64)
                .or_else(|| self.lengths.get(&posting.key).map(|l| *l as f64))
                .unwrap_or(average_length);
            let frequency = posting.frequency as f64;
            let denominator = frequency + K1 * (1.0 - B + (B * length) / average_length);
            take(posting.key.clone(), (idf * (frequency * (K1 + 1.0))) / denominator);
        }
    }

    fn expand(&self, prefix: &str, corpus: &Corpus) -> Vec<String> {
        // NOTE: the range lookup below is valid only once tokens indexed since the last search merge.
        self.merge_fresh_tokens();
        let mut tokens = HashSet::new();
        let mut at = self.lower_bound(prefix);
        let vocabulary = cached(&self.vocabulary);
        while at < vocabulary.len() {
            let token = &vocabulary[at];
            if !token.starts_with(prefix) {
                break;
            }
            tokens.insert(token.clone());
            at += 1;
        }
        drop(vocabulary);
        for document in corpus.staged.values() {
            for token in document.frequencies.keys() {
                if token.starts_with(prefix) {
                    tokens.insert(token.clone());
                }
            }
        }
        tokens.into_iter().collect()
    }

    fn merge_fresh_tokens(&self) {
        // NOTE: taken under its own guard — `lower_bound` below locks `vocabulary`, never this one.
        let mut fresh = std::mem::take(&mut *cached(&self.fresh_tokens));
        if fresh.is_empty() {
            return;
        }

        if fresh.len() <= SPLICED_TOKENS {
            for token in fresh {
                let at = self.lower_bound(&token);
                cached(&self.vocabulary).insert(at, token);
            }
            return;
        }

        fresh.sort();
        let sorted = std::mem::take(&mut *cached(&self.vocabulary));
        let mut merged = Vec::with_capacity(sorted.len() + fresh.len());
        let mut left = 0;
        let mut right = 0;
        while left < sorted.len() || right < fresh.len() {
            if right == fresh.len() || (left < sorted.len() && sorted[left] <= fresh[right]) {
                merged.push(sorted[left].clone());
                left += 1;
            } else {
                merged.push(fresh[right].clone());
                right += 1;
            }
        }
        *cached(&self.vocabulary) = merged;
    }

    fn lower_bound(&self, token: &str) -> usize {
        let vocabulary = cached(&self.vocabulary);
        let mut low = 0;
        let mut high = vocabulary.len();
        while low < high {
            let middle = (low + high) / 2;
            if vocabulary[middle].as_str() < token {
                low = middle + 1;
            } else {
                high = middle;
            }
        }
        low
    }

    fn corpus_with(&self, delta: Option<&CorpusDelta>) -> Corpus {
        let mut staged = HashMap::new();
        let mut hidden = HashSet::new();
        let mut documents = self.lengths.len() as i64;
        let mut total_tokens = self.total_tokens;
        let Some(delta) = delta else {
            return Corpus {
                documents,
                total_tokens,
                staged,
                hidden,
            };
        };

        for key in &delta.removed {
            hidden.insert(key.clone());
        }
        for (key, tokens) in &delta.added {
            hidden.insert(key.clone());
            if tokens.is_empty() {
                continue;
            }
            let mut frequencies = HashMap::new();
            for token in tokens {
                *frequencies.entry(token.clone()).or_insert(0) += 1;
            }
            staged.insert(
                key.clone(),
                StagedDocument {
                    length: tokens.len(),
                    frequencies,
                },
            );
        }

        for key in &hidden {
            if let Some(length) = self.lengths.get(key) {
                documents -= 1;
                total_tokens -= *length as i64;
            }
        }
        for document in staged.values() {
            documents += 1;
            total_tokens += document.length as i64;
        }
        Corpus {
            documents,
            total_tokens,
            staged,
            hidden,
        }
    }

    fn postings_of(&self, token: &str, corpus: &Corpus) -> Vec<Posting> {
        let mut postings = Vec::new();
        if let Some(bucket) = self.postings.get(token) {
            for (key, frequency) in bucket {
                if corpus.hidden.contains(key) {
                    continue;
                }
                postings.push(Posting {
                    key: key.clone(),
                    frequency: *frequency,
                });
            }
        }
        for (key, document) in &corpus.staged {
            if let Some(frequency) = document.frequencies.get(token) {
                postings.push(Posting {
                    key: key.clone(),
                    frequency: *frequency,
                });
            }
        }
        postings
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::registry::tokenize::tokenize;

    fn strings(tokens: &[&str]) -> Vec<String> {
        tokens.iter().map(|t| t.to_string()).collect()
    }

    #[test]
    fn requires_every_query_token_by_default_and_ranks_the_closer_document_higher() {
        let mut index = FulltextIndex::new();
        index.index("a", &tokenize("backend developer node"));
        index.index("b", &tokenize("frontend developer"));

        let hits: Vec<String> = index
            .search("developer backend", SearchMode::All, None)
            .into_keys()
            .collect();
        assert_eq!(hits, vec!["a".to_string()]);

        let any = index.search("developer frontend", SearchMode::Any, None);
        let mut keys: Vec<&String> = any.keys().collect();
        keys.sort();
        assert_eq!(keys, vec!["a", "b"]);
        assert!(any["b"] > any["a"]);
    }

    #[test]
    fn forgets_a_removed_document() {
        let mut index = FulltextIndex::new();
        index.index("a", &strings(&["backend"]));
        index.remove("a");
        assert_eq!(index.documents(), 0);
        assert!(index.search("backend", SearchMode::All, None).is_empty());
    }

    #[test]
    fn scores_a_staged_document_against_the_corpus_it_is_added_to() {
        let mut index = FulltextIndex::new();
        index.index("a", &strings(&["backend", "developer"]));
        let alone = index.search("backend", SearchMode::All, None)["a"];

        let delta = CorpusDelta {
            added: HashMap::from([
                ("b".to_string(), strings(&["backend", "developer"])),
                ("c".to_string(), strings(&["backend", "architect"])),
            ]),
            removed: HashSet::new(),
        };
        let scores = index.search("backend", SearchMode::All, Some(&delta));

        let mut keys: Vec<&String> = scores.keys().collect();
        keys.sort();
        assert_eq!(keys, vec!["a", "b", "c"]);
        assert!(scores["a"] < alone);
    }

    #[test]
    fn leaves_a_removed_document_out_of_the_scores_and_statistics() {
        let mut index = FulltextIndex::new();
        index.index("a", &strings(&["backend", "developer"]));
        index.index("b", &strings(&["frontend", "developer"]));
        let with_both = index.search("developer", SearchMode::All, None)["a"];

        let delta = CorpusDelta {
            added: HashMap::new(),
            removed: HashSet::from(["b".to_string()]),
        };
        let scores = index.search("developer", SearchMode::All, Some(&delta));

        let keys: Vec<&String> = scores.keys().collect();
        assert_eq!(keys, vec!["a"]);
        assert_ne!(scores["a"], with_both);
    }

    #[test]
    fn matches_a_token_that_only_begins_with_the_query_term_when_marked() {
        let mut index = FulltextIndex::new();
        index.index("a", &strings(&["wybor", "silnika", "skladowania"]));
        index.index("b", &strings(&["wybor", "formatu"]));

        assert!(index.search("silnik", SearchMode::All, None).is_empty());
        let hits: Vec<String> = index.search("silnik*", SearchMode::All, None).into_keys().collect();
        assert_eq!(hits, vec!["a".to_string()]);
    }

    #[test]
    fn scores_each_expanded_token_by_its_own_document_frequency() {
        let mut index = FulltextIndex::new();
        index.index("rare", &strings(&["silnikowego"]));
        index.index("common", &strings(&["silnik"]));
        for at in 0..20 {
            index.index(&format!("filler_{at}"), &strings(&["silnik"]));
        }

        let scores = index.search("silnik*", SearchMode::All, None);
        assert!(scores["rare"] > scores["common"]);
        assert_eq!(scores["common"], scores["filler_0"]);
    }

    #[test]
    fn judges_an_unindexed_document_by_the_same_membership_rule_as_the_postings() {
        let mut index = FulltextIndex::new();
        let tokens = tokenize("Backend architecture");
        index.index("indexed", &tokens);

        for (query, mode) in [
            ("backend", SearchMode::All),
            ("architec*", SearchMode::All),
            ("architek* backend", SearchMode::All),
            ("architek* backend", SearchMode::Any),
            ("frontend", SearchMode::All),
            ("", SearchMode::All),
        ] {
            let indexed = index.search(query, mode, None).contains_key("indexed");
            assert_eq!(
                matches_terms(&tokens, query, mode),
                indexed,
                "query \"{query}\" disagreed"
            );
        }
    }

    #[test]
    fn forgets_a_token_from_the_vocabulary_when_its_last_document_goes() {
        let mut index = FulltextIndex::new();
        index.index("a", &strings(&["silnika"]));
        index.index("b", &strings(&["silnikowego"]));
        index.remove("a");

        let hits: Vec<String> = index.search("silnik*", SearchMode::All, None).into_keys().collect();
        assert_eq!(hits, vec!["b".to_string()]);
        index.remove("b");
        assert!(index.search("silnik*", SearchMode::All, None).is_empty());
    }
}