infino 0.5.5

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Infino Authors

//! Textbook BM25 reference implementation. Computes top-k by
//! scoring every doc directly from the BM25 formula; no inverted
//! index, no skip table, no WAND. Used as the correctness oracle
//! for the FTS pipeline's optimized BMW / BMM walks.
//!
//! Formula matches infino's production scoring (standard BM25 IDF +
//! standard BM25 tf factor):
//!
//! ```text
//!   idf(t)      = ln(1 + (N - df(t) + 0.5) / (df(t) + 0.5))
//!   tf_factor   = tf · (K1 + 1) / (tf + K1 · (1 - B + B · dl/avgdl))
//!   score(d, q) = Σ_{t ∈ q} idf(t) · tf_factor(tf(t, d), dl(d), avgdl)
//!
//!   K1 = 1.2,  B = 0.75            (standard BM25 defaults)
//! ```
//!
//! Tokenization runs through whatever `Tokenizer` the caller
//! supplies — pass [`crate::superfile::fts::tokenize::AsciiLowerTokenizer`]
//! to match the production pipeline.
//!
//! Result invariants match the optimized search path:
//! top-k by descending score, ties broken by ascending doc_id.

use std::{cmp::Ordering, collections::HashMap};

use crate::superfile::fts::tokenize::Tokenizer;

/// Standard BM25 default parameters. Match the constants used by
/// the production scoring path.
const K1: f32 = 1.2;
const B: f32 = 0.75;

/// Number of times `phrase` occurs contiguously (in order) in
/// `tokens` — the brute-force phrase tf.
fn phrase_tf(tokens: &[String], phrase: &[String]) -> u32 {
    if phrase.is_empty() || tokens.len() < phrase.len() {
        return 0;
    }
    let mut tf = 0u32;
    for start in 0..=(tokens.len() - phrase.len()) {
        if tokens[start..start + phrase.len()] == *phrase {
            tf += 1;
        }
    }
    tf
}

/// Per-doc statistics derived from the corpus once and reused
/// across queries.
struct DocStats {
    /// Doc id, mirrored from the input.
    doc_id: u64,
    /// Token count (doc length).
    dl: u32,
    /// Term frequencies for this doc: term → count.
    tf: HashMap<String, u32>,
    /// The doc's token sequence, for phrase adjacency scans.
    tokens: Vec<String>,
}

/// Pre-tokenized corpus + per-term df + corpus avgdl. Construct
/// once per fixture, query many times.
pub struct BruteForceBm25 {
    docs: Vec<DocStats>,
    /// Document-frequency: how many docs contain a given term at
    /// least once.
    df: HashMap<String, u32>,
    avgdl: f32,
    n: u32,
}

impl BruteForceBm25 {
    /// Tokenize the corpus once and capture per-doc + per-term
    /// statistics for later scoring. `tokenizer` MUST match what
    /// the production pipeline indexed the same corpus under —
    /// otherwise dl, df, and tf will diverge and recall comparisons
    /// will be meaningless.
    pub fn index(corpus: &[(u64, &str)], tokenizer: &dyn Tokenizer) -> Self {
        let mut docs: Vec<DocStats> = Vec::with_capacity(corpus.len());
        let mut df: HashMap<String, u32> = HashMap::new();
        let mut total_tokens: u64 = 0;

        for (doc_id, text) in corpus {
            let mut tf: HashMap<String, u32> = HashMap::new();
            let mut tokens: Vec<String> = Vec::new();
            let mut dl: u32 = 0;
            tokenizer.tokenize_each(text, &mut |tok| {
                dl += 1;
                *tf.entry(tok.to_owned()).or_insert(0) += 1;
                tokens.push(tok.to_owned());
            });
            for term in tf.keys() {
                *df.entry(term.clone()).or_insert(0) += 1;
            }
            total_tokens += dl as u64;
            docs.push(DocStats {
                doc_id: *doc_id,
                dl,
                tf,
                tokens,
            });
        }

        let n = docs.len() as u32;
        let avgdl = if n == 0 {
            0.0
        } else {
            total_tokens as f32 / n as f32
        };

        Self { docs, df, avgdl, n }
    }

    /// Brute-force top-k for a multi-term OR-mode BM25 query.
    /// `query` is tokenized by the same tokenizer the index was
    /// built with (the caller passes both; we do not capture
    /// the tokenizer to keep the struct `Send + Sync` regardless
    /// of the tokenizer's bounds).
    ///
    /// Returns up to `k` `(doc_id, score)` pairs in descending
    /// score order; tie-breaks ascending by `doc_id`.
    pub fn top_k(&self, query: &str, k: usize, tokenizer: &dyn Tokenizer) -> Vec<(u64, f32)> {
        if k == 0 || self.n == 0 {
            return Vec::new();
        }
        let mut q_terms: Vec<String> = Vec::new();
        tokenizer.tokenize_each(query, &mut |tok| q_terms.push(tok.to_owned()));
        self.top_k_terms(&q_terms, k)
    }

    /// As [`Self::top_k`] but skips the query-side tokenization
    /// step — the caller has already tokenized.
    pub fn top_k_terms(&self, terms: &[String], k: usize) -> Vec<(u64, f32)> {
        if k == 0 || terms.is_empty() || self.n == 0 {
            return Vec::new();
        }

        let n = self.n as f32;
        let avgdl = self.avgdl;

        let mut scored: Vec<(u64, f32)> = Vec::with_capacity(self.docs.len());
        for doc in &self.docs {
            let mut score: f32 = 0.0;
            let dl = doc.dl as f32;
            let dl_norm = K1 * (1.0 - B + B * dl / avgdl.max(f32::MIN_POSITIVE));
            for term in terms {
                let Some(&tf) = doc.tf.get(term) else {
                    continue;
                };
                let df = *self.df.get(term).unwrap_or(&0) as f32;
                if df == 0.0 {
                    continue;
                }
                let idf = (1.0 + (n - df + 0.5) / (df + 0.5)).ln();
                let tf_f = tf as f32;
                score += idf * tf_f * (K1 + 1.0) / (tf_f + dl_norm);
            }
            if score > 0.0 {
                scored.push((doc.doc_id, score));
            }
        }

        scored.sort_by(|a, b| {
            b.1.partial_cmp(&a.1)
                .unwrap_or(Ordering::Equal)
                .then(a.0.cmp(&b.0))
        });
        scored.truncate(k);
        scored
    }

    /// Clause-model top-k: the match set is the docs containing
    /// **every** must (or, with no musts, **any** should), minus docs
    /// containing any negative. A matching doc's score sums its must
    /// contributions plus each should that lands on it — shoulds are
    /// scoring-only and never add or remove a match. Mirrors the
    /// production must/should walk; tie-break is ascending doc_id,
    /// identical to the other helpers.
    pub fn top_k_clauses(
        &self,
        musts: &[String],
        shoulds: &[String],
        negatives: &[String],
        k: usize,
    ) -> Vec<(u64, f32)> {
        if k == 0 || (musts.is_empty() && shoulds.is_empty()) || self.n == 0 {
            return Vec::new();
        }

        let n = self.n as f32;
        let avgdl = self.avgdl;

        let mut scored: Vec<(u64, f32)> = Vec::with_capacity(self.docs.len());
        'docs: for doc in &self.docs {
            for neg in negatives {
                if doc.tf.contains_key(neg) {
                    continue 'docs;
                }
            }
            for must in musts {
                if !doc.tf.contains_key(must) {
                    continue 'docs;
                }
            }
            let dl = doc.dl as f32;
            let dl_norm = K1 * (1.0 - B + B * dl / avgdl.max(f32::MIN_POSITIVE));
            let mut score: f32 = 0.0;
            let mut any_should = false;
            for term in musts.iter().chain(shoulds) {
                let Some(&tf) = doc.tf.get(term) else {
                    continue;
                };
                any_should |= !musts.contains(term);
                let df = *self.df.get(term).unwrap_or(&0) as f32;
                if df == 0.0 {
                    continue;
                }
                let idf = (1.0 + (n - df + 0.5) / (df + 0.5)).ln();
                let tf_f = tf as f32;
                score += idf * tf_f * (K1 + 1.0) / (tf_f + dl_norm);
            }
            // With musts, every surviving doc matches (score > 0 since
            // idf is always positive); with none, only docs hit by at
            // least one should are in the union.
            if !musts.is_empty() || any_should {
                scored.push((doc.doc_id, score));
            }
        }

        scored.sort_by(|a, b| {
            b.1.partial_cmp(&a.1)
                .unwrap_or(Ordering::Equal)
                .then(a.0.cmp(&b.0))
        });
        scored.truncate(k);
        scored
    }

    /// Phrase-aware clause-model top-k: like
    /// [`Self::top_k_clauses`], with exact phrases in every polarity.
    /// A phrase matches a doc when its token sequence occurs
    /// contiguously in order; its per-doc tf is the number of
    /// occurrence starts, and it scores as one BM25 atom with
    /// `idf = Σ member idf` — mirroring the production phrase atom.
    #[allow(clippy::too_many_arguments)]
    pub fn top_k_atoms(
        &self,
        musts: &[String],
        must_phrases: &[Vec<String>],
        shoulds: &[String],
        should_phrases: &[Vec<String>],
        negatives: &[String],
        negative_phrases: &[Vec<String>],
        k: usize,
    ) -> Vec<(u64, f32)> {
        let no_positive = musts.is_empty()
            && must_phrases.is_empty()
            && shoulds.is_empty()
            && should_phrases.is_empty();
        if k == 0 || no_positive || self.n == 0 {
            return Vec::new();
        }

        let n = self.n as f32;
        let idf_of = |term: &String| -> f32 {
            let df = *self.df.get(term).unwrap_or(&0) as f32;
            match df > 0.0 {
                true => (1.0 + (n - df + 0.5) / (df + 0.5)).ln(),
                false => 0.0,
            }
        };
        let phrase_idf = |p: &Vec<String>| -> f32 { p.iter().map(idf_of).sum() };

        let avgdl = self.avgdl;
        let mut scored: Vec<(u64, f32)> = Vec::with_capacity(self.docs.len());
        'docs: for doc in &self.docs {
            for neg in negatives {
                if doc.tf.contains_key(neg) {
                    continue 'docs;
                }
            }
            for p in negative_phrases {
                if phrase_tf(&doc.tokens, p) > 0 {
                    continue 'docs;
                }
            }
            for must in musts {
                if !doc.tf.contains_key(must) {
                    continue 'docs;
                }
            }
            let mut must_phrase_tfs: Vec<u32> = Vec::with_capacity(must_phrases.len());
            for p in must_phrases {
                let tf = phrase_tf(&doc.tokens, p);
                if tf == 0 {
                    continue 'docs;
                }
                must_phrase_tfs.push(tf);
            }

            let dl = doc.dl as f32;
            let dl_norm = K1 * (1.0 - B + B * dl / avgdl.max(f32::MIN_POSITIVE));
            let tf_factor = |tf: u32| -> f32 { tf as f32 * (K1 + 1.0) / (tf as f32 + dl_norm) };

            let mut score: f32 = 0.0;
            let mut matched_any_should = false;
            for term in musts {
                score += idf_of(term) * tf_factor(doc.tf[term]);
            }
            for (p, tf) in must_phrases.iter().zip(&must_phrase_tfs) {
                score += phrase_idf(p) * tf_factor(*tf);
            }
            for term in shoulds {
                if let Some(&tf) = doc.tf.get(term) {
                    matched_any_should = true;
                    score += idf_of(term) * tf_factor(tf);
                }
            }
            for p in should_phrases {
                let tf = phrase_tf(&doc.tokens, p);
                if tf > 0 {
                    matched_any_should = true;
                    score += phrase_idf(p) * tf_factor(tf);
                }
            }

            let has_musts = !musts.is_empty() || !must_phrases.is_empty();
            if has_musts || matched_any_should {
                scored.push((doc.doc_id, score));
            }
        }

        scored.sort_by(|a, b| {
            b.1.partial_cmp(&a.1)
                .unwrap_or(Ordering::Equal)
                .then(a.0.cmp(&b.0))
        });
        scored.truncate(k);
        scored
    }

    /// Same as [`Self::top_k_terms`] but with AND semantics: only docs
    /// that contain *every* query term contribute a score. Used by the
    /// AND-mode oracle tests against the reader's leapfrog
    /// intersection. Each term still contributes its own BM25 score to
    /// the per-doc sum; tie-break is ascending doc_id, identical to
    /// the OR helper.
    pub fn top_k_terms_and(&self, terms: &[String], k: usize) -> Vec<(u64, f32)> {
        if k == 0 || terms.is_empty() || self.n == 0 {
            return Vec::new();
        }

        let n = self.n as f32;
        let avgdl = self.avgdl;

        let mut scored: Vec<(u64, f32)> = Vec::with_capacity(self.docs.len());
        'docs: for doc in &self.docs {
            let dl = doc.dl as f32;
            let dl_norm = K1 * (1.0 - B + B * dl / avgdl.max(f32::MIN_POSITIVE));
            let mut score: f32 = 0.0;
            for term in terms {
                let Some(&tf) = doc.tf.get(term) else {
                    continue 'docs;
                };
                let df = *self.df.get(term).unwrap_or(&0) as f32;
                if df == 0.0 {
                    continue 'docs;
                }
                let idf = (1.0 + (n - df + 0.5) / (df + 0.5)).ln();
                let tf_f = tf as f32;
                score += idf * tf_f * (K1 + 1.0) / (tf_f + dl_norm);
            }
            scored.push((doc.doc_id, score));
        }

        scored.sort_by(|a, b| {
            b.1.partial_cmp(&a.1)
                .unwrap_or(Ordering::Equal)
                .then(a.0.cmp(&b.0))
        });
        scored.truncate(k);
        scored
    }

    /// Number of indexed docs.
    pub fn n_docs(&self) -> u32 {
        self.n
    }
}