codesearch 0.1.14

//! BM25 (Best Match 25) Relevance Scoring
//!
//! BM25 is a probabilistic ranking function used by search engines.
//! It scores documents based on term frequency, inverse document frequency,
//! and document length normalization.
//!
//! Formula for a single term:
//! ```text
//! score(D, q) = IDF(q) * (tf(q,D) * (k1 + 1)) / (tf(q,D) + k1 * (1 - b + b * |D| / avgdl))
//! ```
//!
//! Where:
//! - `tf(q,D)` = term frequency of query q in document D
//! - `|D|` = length of document D (in lines)
//! - `avgdl` = average document length
//! - `k1`, `b` = tuning parameters (defaults: k1=1.2, b=0.75)
//! - `IDF(q) = log((N - n(q) + 0.5) / (n(q) + 0.5) + 1)`
//!   - `N` = total number of documents
//!   - `n(q)` = number of documents containing the query term

/// Collection-level statistics needed for BM25 scoring.
/// Computed once per search, then shared across all result scoring.
#[derive(Debug, Clone, Copy)]
pub struct Bm25Stats {
    /// Total number of documents (files) in the collection
    pub total_docs: usize,
    /// Number of documents containing the query term at least once
    pub docs_with_term: usize,
    /// Average document length (in lines)
    pub avg_doc_length: f64,
    /// BM25 k1 parameter (controls term frequency saturation)
    pub k1: f64,
    /// BM25 b parameter (controls length normalization)
    pub b: f64,
}

impl Bm25Stats {
    /// Create stats with default BM25 parameters (k1=1.2, b=0.75).
    pub fn new(total_docs: usize, docs_with_term: usize, avg_doc_length: f64) -> Self {
        Self {
            total_docs,
            docs_with_term,
            avg_doc_length,
            k1: 1.2,
            b: 0.75,
        }
    }

    /// Compute the inverse document frequency (IDF) component.
    ///
    /// Uses the standard BM25 IDF variant that avoids negative values:
    /// `log((N - n(q) + 0.5) / (n(q) + 0.5) + 1)`
    fn idf(&self) -> f64 {
        let n = self.docs_with_term as f64;
        let n_total = self.total_docs as f64;
        let numerator = n_total - n + 0.5;
        let denominator = n + 0.5;
        (numerator / denominator).ln() + 1.0
    }

    /// Compute the BM25 score for a single document.
    ///
    /// # Arguments
    ///
    /// * `term_freq` — Number of times the query appears in the document
    /// * `doc_length` — Length of the document in lines (or other unit)
    ///
    /// # Returns
    ///
    /// A raw BM25 score (unbounded; higher is better).
    pub fn score(&self, term_freq: usize, doc_length: usize) -> f64 {
        if self.total_docs == 0 || self.avg_doc_length == 0.0 {
            return 0.0;
        }

        let tf = term_freq as f64;
        let dl = doc_length as f64;
        let idf = self.idf();

        let tf_component = tf * (self.k1 + 1.0);
        let normalization = tf + self.k1 * (1.0 - self.b + self.b * (dl / self.avg_doc_length));

        idf * (tf_component / normalization)
    }

    /// Normalize a raw BM25 score to a 0–100 scale for display.
    ///
    /// Uses a sigmoid-like clamp to avoid extreme outliers dominating.
    pub fn normalized_score(&self, term_freq: usize, doc_length: usize) -> f64 {
        let raw = self.score(term_freq, doc_length);
        // Scale: typical BM25 scores range 0–10 for a single term.
        // Multiply by 10 and clamp to 100.
        (raw * 10.0).clamp(0.0, 100.0)
    }
}

/// Compute collection statistics for BM25 from a set of files.
///
/// This scans each file to count lines and detect term presence.
/// It is a one-time cost per search.
pub fn compute_collection_stats(
    files: &[std::path::PathBuf],
    query: &str,
    case_sensitive: bool,
) -> Bm25Stats {
    let total_docs = files.len();
    if total_docs == 0 {
        return Bm25Stats::new(0, 0, 1.0);
    }

    let mut total_lines = 0usize;
    let mut docs_with_term = 0usize;

    for file_path in files {
        if let Ok(content) = std::fs::read_to_string(file_path) {
            let lines = content.lines().count();
            total_lines += lines;

            let contains = if case_sensitive {
                content.contains(query)
            } else {
                content.to_lowercase().contains(&query.to_lowercase())
            };

            if contains {
                docs_with_term += 1;
            }
        }
    }

    let avg_doc_length = if total_docs > 0 {
        total_lines as f64 / total_docs as f64
    } else {
        1.0
    };

    Bm25Stats::new(total_docs, docs_with_term, avg_doc_length)
}

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

    #[test]
    fn test_bm25_score_basic() {
        let stats = Bm25Stats::new(100, 10, 50.0);
        // A doc with 5 occurrences in 50 lines (avg length)
        let score = stats.score(5, 50);
        assert!(score > 0.0);
    }

    #[test]
    fn test_bm25_longer_doc_penalty() {
        let stats = Bm25Stats::new(100, 10, 50.0);
        // Same term frequency, but doc is 2x average length
        let score_short = stats.score(5, 50);
        let score_long = stats.score(5, 100);
        assert!(score_short > score_long, "Shorter doc should score higher");
    }

    #[test]
    fn test_bm25_rare_term_boosts_score() {
        // Rare term (only in 1 doc) vs common term (in 50 docs)
        let rare = Bm25Stats::new(100, 1, 50.0);
        let common = Bm25Stats::new(100, 50, 50.0);

        let score_rare = rare.score(3, 50);
        let score_common = common.score(3, 50);

        assert!(
            score_rare > score_common,
            "Rare term should score higher than common term"
        );
    }

    #[test]
    fn test_bm25_normalized_score_range() {
        let stats = Bm25Stats::new(100, 10, 50.0);
        let score = stats.normalized_score(10, 50);
        assert!(score >= 0.0 && score <= 100.0);
    }

    #[test]
    fn test_bm25_empty_collection() {
        let stats = Bm25Stats::new(0, 0, 1.0);
        assert_eq!(stats.score(5, 50), 0.0);
    }

    #[test]
    fn test_idf_computation() {
        let stats = Bm25Stats::new(100, 10, 50.0);
        let idf = stats.idf();
        assert!(idf > 0.0);
    }
}