lora-store 0.32.1

//! Trigram inverted index for `TEXT` Cypher indexes.
//!
//! Each indexed (label-or-type, property) pair owns a
//! [`TrigramScope`]. The scope maps every overlapping 3-byte window of
//! the string's UTF-8 encoding to the set of entity ids whose value
//! contains that window. This lets `CONTAINS`, `STARTS WITH`, and
//! `ENDS WITH` predicates skip the full property scan: probe by the
//! query's own trigrams, intersect, and refilter the candidate set.
//!
//! ## Why byte-level trigrams
//!
//! Some references describe *Unicode code-point* trigrams.
//! Byte-level trigrams over UTF-8 yield the same set of valid
//! 3-byte aligned windows for pure-ASCII text and a strict
//! superset for multibyte text — the index *over*-collects, the
//! refilter step (mandatory anyway, since any trigram set can match
//! a substring not actually present) trims it back. In return we
//! avoid the cost of decoding every char to look up a single key.
//!
//! ## Maintenance contract
//!
//! All mutations go through the existing `MutationEvent` stream;
//! `InMemoryGraph::on_node_property_set` etc. delegate updates here
//! when the registry has at least one matching scope. When the last
//! scope referencing a (label, property) pair is dropped, the bucket
//! itself is removed so an empty TEXT catalog pays zero memory.

use std::collections::{BTreeSet, HashMap};

use super::cow::CowMap;
use super::entity_index_store::ScopedPropertyKey;
use super::id_map::CowIdMap;

/// Leaf size of the trigram trie. A ~20-character value touches ~20
/// trigrams spread over most leaves, and a staged write copies each leaf
/// it touches: a quarter of the default leaf size cut that copying (and
/// the old leaves' drop) and made indexed MERGE/CREATE ~10% faster.
const GRAM_LEAF: usize = 112;

/// Registry of trigram scopes for either nodes or relationships.
#[derive(Debug, Default, Clone)]
pub(super) struct TrigramRegistry {
    pub(super) by_scope: HashMap<ScopedPropertyKey, TrigramScope>,
}

#[derive(Debug, Default, Clone)]
pub(super) struct TrigramScope {
    /// Trigram → entity ids whose property value contains it.
    pub(super) grams: CowMap<[u8; 3], CowIdMap<()>, GRAM_LEAF>,
    /// Reference count: how many catalog entries point at this scope.
    /// We allow multiple TEXT indexes on the same `(label, property)`
    /// (different names, redundant); the scope is freed only when the
    /// last reference is dropped.
    refcount: u32,
}

impl TrigramRegistry {
    /// Whether any of `scopes` has an index on `property`. Checked through
    /// a read guard before maintenance takes a (copy-on-write) write guard.
    pub(super) fn covers_any<'a>(
        &self,
        scopes: impl IntoIterator<Item = &'a str>,
        property: &str,
    ) -> bool {
        !self.by_scope.is_empty()
            && scopes.into_iter().any(|scope| {
                self.by_scope
                    .contains_key(&ScopedPropertyKey::new(scope, property))
            })
    }

    /// Mark a scope as in-use, allocating it if missing. Returns
    /// `true` if the scope was freshly created (caller should
    /// backfill it from existing data).
    pub(super) fn add_scope(&mut self, label: &str, property: &str) -> bool {
        let key = ScopedPropertyKey::new(label, property);
        let entry = self.by_scope.entry(key).or_default();
        let was_empty = entry.refcount == 0;
        entry.refcount = entry.refcount.saturating_add(1);
        was_empty
    }

    /// Decrement the refcount on a scope. The scope is removed once
    /// the last reference is gone.
    pub(super) fn remove_scope(&mut self, label: &str, property: &str) {
        let key = ScopedPropertyKey::new(label, property);
        if let Some(scope) = self.by_scope.get_mut(&key) {
            scope.refcount = scope.refcount.saturating_sub(1);
            if scope.refcount == 0 {
                self.by_scope.remove(&key);
            }
        }
    }

    #[cfg(test)]
    pub(super) fn has_scope(&self, label: &str, property: &str) -> bool {
        self.by_scope
            .contains_key(&ScopedPropertyKey::new(label, property))
    }

    pub(super) fn insert(&mut self, label: &str, property: &str, id: u64, value: &str) {
        if let Some(scope) = self
            .by_scope
            .get_mut(&ScopedPropertyKey::new(label, property))
        {
            scope.insert(id, value);
        }
    }

    #[cfg(test)]
    pub(super) fn remove(&mut self, label: &str, property: &str, id: u64, value: &str) {
        if let Some(scope) = self
            .by_scope
            .get_mut(&ScopedPropertyKey::new(label, property))
        {
            scope.remove(id, value);
        }
    }

    pub(super) fn update(
        &mut self,
        label: &str,
        property: &str,
        id: u64,
        old: Option<&str>,
        new: Option<&str>,
    ) {
        let Some(scope) = self
            .by_scope
            .get_mut(&ScopedPropertyKey::new(label, property))
        else {
            return;
        };
        if let Some(old) = old {
            scope.remove(id, old);
        }
        if let Some(new) = new {
            scope.insert(id, new);
        }
    }

    /// Candidate ids whose property value *might* contain `query`.
    /// `None` when the query is too short for trigram matching (callers
    /// should fall back to scan).
    pub(super) fn candidates(
        &self,
        label: &str,
        property: &str,
        query: &str,
    ) -> Option<BTreeSet<u64>> {
        let scope = self
            .by_scope
            .get(&ScopedPropertyKey::new(label, property))?;
        scope.candidates(query)
    }
}

impl TrigramScope {
    fn insert(&mut self, id: u64, value: &str) {
        for tri in trigrams(value) {
            self.grams
                .get_or_insert_with(tri, CowIdMap::default)
                .insert(id, ());
        }
    }

    fn remove(&mut self, id: u64, value: &str) {
        for tri in trigrams(value) {
            let emptied = self.grams.get_mut(&tri).is_some_and(|set| {
                set.remove(&id);
                set.is_empty()
            });
            if emptied {
                self.grams.remove(&tri);
            }
        }
    }

    fn candidates(&self, query: &str) -> Option<BTreeSet<u64>> {
        let mut grams: Vec<[u8; 3]> = trigrams(query).collect();
        if grams.is_empty() {
            // Query shorter than 3 bytes — every entity is a potential
            // match. Returning `None` signals the caller to scan.
            return None;
        }
        // Probe lowest-cardinality trigram first — once that bucket is
        // narrowed, intersection on subsequent probes runs over a
        // smaller working set.
        grams.sort_by_key(|tri| self.grams.get(tri).map(CowIdMap::len).unwrap_or(usize::MAX));
        // The scope indexes every value, so a trigram it has never seen
        // is in no value: nothing matches (`None` would mean a full scan).
        let Some(first) = self.grams.get(&grams[0]) else {
            return Some(BTreeSet::new());
        };
        let mut out: BTreeSet<u64> = first.keys().collect();
        for tri in &grams[1..] {
            let Some(next) = self.grams.get(tri) else {
                return Some(BTreeSet::new());
            };
            out.retain(|id| next.contains_key(id));
            if out.is_empty() {
                return Some(out);
            }
        }
        Some(out)
    }
}

fn trigrams(value: &str) -> impl Iterator<Item = [u8; 3]> + '_ {
    let bytes = value.as_bytes();
    (0..bytes.len().saturating_sub(2)).map(move |i| {
        let window = &bytes[i..i + 3];
        [window[0], window[1], window[2]]
    })
}

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

    #[test]
    fn empty_query_returns_none() {
        let mut reg = TrigramRegistry::default();
        reg.add_scope("Person", "name");
        reg.insert("Person", "name", 1, "Alice");
        assert!(reg.candidates("Person", "name", "ab").is_none());
    }

    #[test]
    fn substring_match_intersects_trigrams() {
        let mut reg = TrigramRegistry::default();
        reg.add_scope("Person", "name");
        reg.insert("Person", "name", 1, "Alexandra");
        reg.insert("Person", "name", 2, "Alexander");
        reg.insert("Person", "name", 3, "Bob");

        let candidates = reg.candidates("Person", "name", "Alex").unwrap();
        assert!(candidates.contains(&1));
        assert!(candidates.contains(&2));
        assert!(!candidates.contains(&3));
    }

    #[test]
    fn unseen_trigram_matches_nothing_without_a_scan() {
        let mut reg = TrigramRegistry::default();
        reg.add_scope("Person", "name");
        reg.insert("Person", "name", 1, "Alice");
        assert_eq!(
            reg.candidates("Person", "name", "zzz"),
            Some(BTreeSet::new())
        );
        assert_eq!(
            reg.candidates("Person", "name", "Alizz"),
            Some(BTreeSet::new())
        );
        // An empty scope has seen no trigram either.
        reg.add_scope("Person", "nick");
        assert_eq!(
            reg.candidates("Person", "nick", "Ali"),
            Some(BTreeSet::new())
        );
        assert!(reg.candidates("Other", "name", "Ali").is_none());
    }

    #[test]
    fn remove_after_insert_clears_buckets() {
        let mut reg = TrigramRegistry::default();
        reg.add_scope("Person", "name");
        reg.insert("Person", "name", 1, "Alice");
        reg.remove("Person", "name", 1, "Alice");
        assert_eq!(
            reg.candidates("Person", "name", "Alic"),
            Some(BTreeSet::new())
        );
    }

    #[test]
    fn update_moves_an_id_between_values_sharing_trigrams() {
        let mut reg = TrigramRegistry::default();
        reg.add_scope("Person", "name");
        reg.insert("Person", "name", 1, "Festival number 42");
        reg.insert("Person", "name", 2, "Festival number 7");
        reg.update(
            "Person",
            "name",
            1,
            Some("Festival number 42"),
            Some("Festival nummer 42 aaaa"),
        );
        fn hits(reg: &TrigramRegistry, q: &str) -> BTreeSet<u64> {
            reg.candidates("Person", "name", q).unwrap_or_default()
        }
        assert_eq!(hits(&reg, "Festival"), BTreeSet::from([1, 2]));
        assert_eq!(hits(&reg, "nummer 42"), BTreeSet::from([1]));
        assert_eq!(hits(&reg, "aaaa"), BTreeSet::from([1]));
        assert!(
            hits(&reg, "number 4").is_empty(),
            "'r 4' was only in the old value"
        );
        assert_eq!(hits(&reg, "number"), BTreeSet::from([2]));
        // Dropping the value clears every trigram it alone held.
        reg.update("Person", "name", 1, Some("Festival nummer 42 aaaa"), None);
        assert!(hits(&reg, "aaa").is_empty());
        assert!(hits(&reg, "mme").is_empty());
        reg.update("Person", "name", 2, Some("Festival number 7"), None);
        assert!(reg.by_scope.values().all(|scope| scope.grams.is_empty()));
    }

    #[test]
    fn refcount_keeps_scope_until_last_remove() {
        let mut reg = TrigramRegistry::default();
        assert!(reg.add_scope("Person", "name"));
        assert!(!reg.add_scope("Person", "name")); // already exists
        reg.insert("Person", "name", 1, "Alice");
        reg.remove_scope("Person", "name");
        assert!(reg.has_scope("Person", "name"));
        reg.remove_scope("Person", "name");
        assert!(!reg.has_scope("Person", "name"));
    }
}