eidos_kernel/

trigram.rs

//! Trigram inverted index — the hippocampus (fast pattern separation).
//!
//! Inspired by Google Code Search (Russ Cox, 2012): build a posting list of trigrams
//! (3-character substrings) at index time. At query time, extract trigrams from query terms,
//! intersect posting lists to narrow candidates from O(all_nodes) to O(matching_nodes).
//!
//! ## How it works
//!
//! Each node's haystack (id + title + summary + aliases + query_examples, lowercased) is split
//! into overlapping 3-char trigrams. A `HashMap<Trigram, RoaringBitSet<NodeIndex>>` maps each
//! trigram to the set of nodes that contain it.
//!
//! At query time: each stemmed query term produces a set of trigrams. The intersection of
//! posting lists for those trigrams gives the set of nodes whose haystack contains that term
//! (as a substring, matching the existing `haystack.contains(term)` semantics exactly).
//!
//! ## Eval-byte-identical guarantee
//!
//! The trigram index only NARROWS the candidate set — it doesn't change which nodes get scored.
//! A node is a candidate if and only if its haystack contains at least one query term as a
//! substring. The trigram intersection is a necessary-but-not-sufficient condition: if a node's
//! haystack contains a term, it necessarily contains all the term's trigrams. So the index never
//! misses a real candidate (no false negatives), and the final `contains()` check eliminates
//! false positives (nodes that have the trigrams but not the actual term).
//!
//! Result: identical scoring, dramatically less work on large graphs.

use std::collections::HashMap;

/// A trigram: a 3-byte substring, stored as a packed u32 for cheap hashing.
type Trigram = u32;

/// The trigram inverted index: maps each trigram to the set of node indices that contain it.
///
/// Built once at `GroundIndex::build` time. Query-side: for each query term, extract its
/// trigrams, intersect posting lists → candidate set. The candidate set is a superset of
/// the actual matches (trigrams are necessary but not sufficient); the final `contains()`
/// check in the scoring loop filters false positives.
///
/// As of D0.6 the indexing side is dormant: `GroundIndex::build` no longer constructs one
/// (scoring.rs:48-50 commented out the call site — see Sprint D0.6). `Default::default()`
/// gives an empty index for callers that need one; `build()` is preserved for the D3
/// re-enable and unit tests that round-trip the original behavior. `candidates()` is
/// therefore safe to call only against a populated index; today the scoring loop in
/// `ground_scoped` (retrieval.rs:835-841) bypasses it and falls through to the full scan.
#[derive(Default)]
pub struct TrigramIndex {
    /// trigram → sorted list of node indices containing it
    postings: HashMap<Trigram, Vec<usize>>,
    /// total node count (for "match all" fallback when a term is < 3 chars)
    node_count: usize,
}

impl Clone for TrigramIndex {
    fn clone(&self) -> Self {
        Self {
            postings: self.postings.clone(),
            node_count: self.node_count,
        }
    }
}

impl TrigramIndex {
    /// Build the index from the per-node haystacks.
    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();

            // Extract all trigrams from this haystack
            if bytes.len() < 3 {
                // Short haystacks (< 3 chars) get the "match all" treatment
                continue;
            }

            for window in bytes.windows(3) {
                let tri = pack_trigram(window);
                postings.entry(tri).or_default().push(node_idx);
            }
        }

        // Dedup and sort each posting list (a node can contain the same trigram multiple times)
        for indices in postings.values_mut() {
            indices.sort_unstable();
            indices.dedup();
        }

        TrigramIndex {
            postings,
            node_count: haystacks.len(),
        }
    }

    /// Find candidate node indices for a query term.
    ///
    /// Returns the set of nodes whose haystack CONTAINS all trigrams of `term`.
    /// This is a superset of nodes whose haystack `contains(term)` — the scoring loop
    /// does the final `contains()` check.
    ///
    /// For terms shorter than 3 characters (no trigrams), returns None = "scan all"
    /// (the fallback — short terms are rare and cheap).
    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; // Too short for trigrams — scan all
        }

        // Extract trigrams from the term
        let trigrams: Vec<Trigram> = bytes.windows(3).map(pack_trigram).collect();

        if trigrams.is_empty() {
            return None;
        }

        // Start with the first trigram's posting list, then intersect with subsequent ones
        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()), // This trigram doesn't exist anywhere → no candidates
            };

            candidates = Some(match candidates.take() {
                None => posting.clone(),
                Some(existing) => intersect_sorted(&existing, posting),
            });

            // Early exit: if the intersection is empty, no node contains all trigrams
            if candidates.as_ref().is_none_or(std::vec::Vec::is_empty) {
                return Some(Vec::new());
            }
        }

        candidates
    }

    /// Find candidate node indices for ALL query terms: the UNION of per-term candidates.
    ///
    /// A node is a candidate if ANY query term's trigrams match (because the scoring loop
    /// checks each term independently — a node that matches one term gets a non-zero score).
    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, // A short term → scan all (can't narrow)
                Some(term_candidates) => {
                    union = Some(match union.take() {
                        None => term_candidates,
                        Some(existing) => union_sorted(&existing, &term_candidates),
                    });
                }
            }
        }

        union
    }
}

/// Pack 3 bytes into a u32 for cheap hashing.
fn pack_trigram(bytes: &[u8]) -> Trigram {
    ((bytes[0] as u32) << 16) | ((bytes[1] as u32) << 8) | (bytes[2] as u32)
}

/// Intersect two sorted vectors, preserving order.
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
}

/// Union two sorted vectors, preserving order and deduplicating.
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(),       // 0
            "task scheduler queue".to_string(),       // 1
            "dead letter queue handling".to_string(), // 2
            "jitter exponential".to_string(),         // 3
        ];

        let index = TrigramIndex::build(&haystacks);

        // "retry" → trigrams: ret, etr, try
        let cands = index.candidates_for_term("retry").unwrap();
        assert!(cands.contains(&0)); // haystack 0 contains "retry"
        assert!(!cands.contains(&1)); // haystack 1 does not

        // "scheduler" → trigrams: sch, che, edu, dul, ula, lat, ato, tor
        let cands = index.candidates_for_term("scheduler").unwrap();
        assert!(cands.contains(&1));
        assert!(!cands.contains(&0));

        // Short term (< 3 chars) → None = scan all
        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);

        // Query "retry queue" → nodes matching EITHER term
        let cands = index.candidates(&["retry".to_string(), "queue".to_string()]);
        if let Some(cands) = cands {
            assert!(cands.contains(&0)); // has "retry"
            assert!(cands.contains(&1)); // has "queue"
            assert!(cands.contains(&2)); // has "queue"
        } 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());
    }
}