codelens_engine/

search.rs

use crate::db::{IndexDb, index_db_path};
use crate::project::ProjectRoot;
use anyhow::Result;
use serde::Serialize;
use strsim::jaro_winkler;

/// Minimum cosine similarity for boosting existing search results with semantic scores.
pub const SEMANTIC_BOOST_THRESHOLD: f64 = 0.10;
/// Minimum cosine similarity for surfacing new semantic-only results.
pub const SEMANTIC_NEW_RESULT_THRESHOLD: f64 = 0.15;
/// Minimum cosine similarity for semantic coupling detection in reports.
pub const SEMANTIC_COUPLING_THRESHOLD: f64 = 0.12;

#[derive(Debug, Clone, Serialize)]
pub struct SearchResult {
    pub name: String,
    pub kind: String,
    pub file: String,
    pub line: usize,
    pub signature: String,
    pub name_path: String,
    pub score: f64,
    pub match_type: String, // "exact", "substring", "fuzzy"
}

/// Maximum PageRank boost added to a result when its file ranks at the top
/// of the import graph. Sized to be smaller than fuzzy/FTS spread so it tilts
/// near-ties toward globally important files without overpowering text relevance.
pub const PAGERANK_MAX_BOOST: f64 = 5.0;

/// Hybrid symbol search: exact → FTS5 → fuzzy → semantic.
///
/// `fuzzy_threshold` — minimum jaro_winkler similarity (0.0–1.0).
/// `semantic_scores` — optional pre-computed semantic similarity scores keyed by
/// "file_path:symbol_name". When provided, semantic-only matches (score > 0.5)
/// are merged as a 4th retrieval path.
///
/// Deduplicated by (name, file, line), sorted by score descending.
pub fn search_symbols_hybrid(
    project: &ProjectRoot,
    query: &str,
    max_results: usize,
    fuzzy_threshold: f64,
) -> Result<Vec<SearchResult>> {
    search_symbols_hybrid_with_semantic(project, query, max_results, fuzzy_threshold, None, None)
}

/// Full hybrid search with optional semantic scores and PageRank file importance.
///
/// `pagerank_scores` — optional per-file PageRank values from `GraphCache::file_pagerank_scores`.
/// When supplied, the top-ranked file in the import graph gains `PAGERANK_MAX_BOOST`
/// points; lower-ranked files scale linearly. The boost is applied before the
/// final sort so it can break near-ties between text-relevance matches.
pub fn search_symbols_hybrid_with_semantic(
    project: &ProjectRoot,
    query: &str,
    max_results: usize,
    fuzzy_threshold: f64,
    semantic_scores: Option<&std::collections::HashMap<String, f64>>,
    pagerank_scores: Option<&std::collections::HashMap<String, f64>>,
) -> Result<Vec<SearchResult>> {
    let db_path = index_db_path(project.as_path());
    let db = IndexDb::open(&db_path)?;

    let mut seen: std::collections::HashSet<(String, String, i64)> =
        std::collections::HashSet::new();
    let mut results: Vec<SearchResult> = Vec::new();

    // ── 1. Exact match (score 100) ──────────────────────────────────────────
    for (row, file) in db.find_symbols_with_path(query, true, max_results)? {
        let key = (row.name.clone(), file.clone(), row.line);
        if seen.insert(key) {
            results.push(SearchResult {
                name: row.name,
                kind: row.kind,
                file,
                line: row.line as usize,
                signature: row.signature,
                name_path: row.name_path,
                score: 100.0,
                match_type: "exact".to_owned(),
            });
        }
    }

    // ── 2. FTS5 full-text search (score 40-80 from rank) ────────────────────
    // Falls back to LIKE search on pre-v4 databases automatically.
    for (row, file, rank) in db.search_symbols_fts(query, max_results)? {
        let key = (row.name.clone(), file.clone(), row.line);
        if seen.insert(key) {
            // FTS5 rank is negative (lower = better), normalize to 40-80 range
            let fts_score = (80.0 + rank.clamp(-40.0, 0.0)).max(40.0);
            results.push(SearchResult {
                name: row.name,
                kind: row.kind,
                file,
                line: row.line as usize,
                signature: row.signature,
                name_path: row.name_path,
                score: fts_score,
                match_type: "fts".to_owned(),
            });
        }
    }

    // ── 3. Fuzzy match (score = similarity * 100) ───────────────────────────
    let query_lower = query.to_ascii_lowercase();
    let prefix: String = query_lower.chars().take(2).collect();
    let fuzzy_candidates = if prefix.len() >= 2 {
        db.find_symbols_with_path(&prefix, false, 500)?
    } else {
        db.find_symbols_with_path(&query_lower, false, 500)?
    };
    for (row, file) in fuzzy_candidates {
        let key = (row.name.clone(), file.clone(), row.line);
        if seen.contains(&key) {
            continue;
        }
        let sim = jaro_winkler(&query_lower, &row.name.to_ascii_lowercase());
        if sim >= fuzzy_threshold {
            seen.insert(key);
            results.push(SearchResult {
                name: row.name,
                kind: row.kind,
                file,
                line: row.line as usize,
                signature: row.signature,
                name_path: row.name_path,
                score: sim * 100.0,
                match_type: "fuzzy".to_owned(),
            });
        }
    }

    // ── 4. Semantic matches (score = cosine_similarity * 90, capped below exact) ─
    if let Some(scores) = semantic_scores {
        // Collect semantic-only discoveries not found by text/fts/fuzzy paths.
        // Only include high-confidence matches (> 0.5 cosine similarity).
        // Use lightweight all_symbol_names() instead of all_symbols_with_bytes()
        // to avoid loading byte offsets into memory.
        let all_symbols = db.all_symbol_names()?;
        for (name, kind, file_path, line, signature, name_path) in all_symbols {
            let key = (name.clone(), file_path.clone(), line);
            if seen.contains(&key) {
                let sem_key = format!("{file_path}:{name}");
                if let Some(&sem_score) = scores.get(&sem_key)
                    && sem_score > SEMANTIC_BOOST_THRESHOLD
                    && let Some(existing) = results
                        .iter_mut()
                        .find(|r| r.name == name && r.file == file_path && r.line == line as usize)
                {
                    existing.score += (sem_score * 15.0).min(10.0);
                }
                continue;
            }
            let sem_key = format!("{file_path}:{name}");
            if let Some(&sem_score) = scores
                .get(&sem_key)
                .filter(|&&s| s > SEMANTIC_NEW_RESULT_THRESHOLD)
            {
                seen.insert(key);
                results.push(SearchResult {
                    name,
                    kind,
                    file: file_path,
                    line: line as usize,
                    signature,
                    name_path,
                    score: sem_score * 90.0,
                    match_type: "semantic".to_owned(),
                });
            }
        }
    }

    if let Some(pr) = pagerank_scores {
        apply_pagerank_boost(&mut results, pr, PAGERANK_MAX_BOOST);
    }

    results.sort_by(|a, b| {
        b.score
            .partial_cmp(&a.score)
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    // Diversity pass: limit results per file to avoid clustering from one module.
    // Promotes cross-file discovery while keeping the highest-scored results.
    const MAX_PER_FILE: usize = 3;
    if results.len() > max_results {
        let mut file_counts: std::collections::HashMap<String, usize> =
            std::collections::HashMap::new();
        let mut promoted = Vec::with_capacity(max_results);
        let mut demoted = Vec::new();
        for r in results {
            let count = file_counts.entry(r.file.clone()).or_insert(0);
            if *count < MAX_PER_FILE {
                *count += 1;
                promoted.push(r);
            } else {
                demoted.push(r);
            }
        }
        promoted.extend(demoted);
        results = promoted;
    }

    results.truncate(max_results);
    Ok(results)
}

/// Adds a PageRank-derived boost (0 .. `max_boost`) to each result based on
/// where its file ranks in the import graph. Uses max-normalization so the
/// boost magnitude is independent of graph size or absolute PageRank values.
/// Applied in-place before the final sort.
fn apply_pagerank_boost(
    results: &mut [SearchResult],
    pagerank_scores: &std::collections::HashMap<String, f64>,
    max_boost: f64,
) {
    let max_pr = pagerank_scores.values().copied().fold(0.0_f64, f64::max);
    if max_pr <= 0.0 {
        return;
    }
    for result in results.iter_mut() {
        if let Some(&pr) = pagerank_scores.get(&result.file) {
            result.score += (pr / max_pr) * max_boost;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::db::{IndexDb, NewSymbol, index_db_path};

    /// Create a temp directory seeded with test symbols.
    ///
    /// Returns the owned [`tempfile::TempDir`] (keep it alive for the
    /// test duration so the directory isn't reaped early) plus a
    /// [`ProjectRoot`] pointed at it.
    ///
    /// Why `tempfile::TempDir` instead of `std::env::temp_dir().join(nanos)`:
    /// the previous version generated paths from `subsec_nanos()` alone,
    /// which collided in macOS CI when nextest scheduled multiple
    /// fixture-using tests on overlapping nanoseconds (9 tests in this
    /// file share this helper). Two tests racing into the same SQLite
    /// file hit `journal_mode = WAL`'s schema-level write lock before
    /// our own `busy_timeout = 5000` PRAGMA could be applied (it's the
    /// 6th PRAGMA in the open batch, so the first PRAGMA still uses
    /// the default 0 ms timeout) — surfacing as `Error code 5:
    /// The database file is locked` at ~0.14 s into the test. macOS
    /// reproduced this on every CI run since 2026-04; Ubuntu happened
    /// to schedule those nine tests in an order that avoided the
    /// collision. `tempfile::tempdir()` guarantees a process-unique
    /// path, removing the collision class entirely.
    fn make_project_with_symbols() -> (tempfile::TempDir, ProjectRoot) {
        let temp = tempfile::tempdir().expect("create temp dir for search test fixture");
        let root = temp.path();

        // Write a dummy source file so ProjectRoot is happy
        std::fs::write(root.join("hello.txt"), "hello").unwrap();

        // Seed the SQLite index
        let db_path = index_db_path(root);
        let db = IndexDb::open(&db_path).unwrap();
        let fid = db
            .upsert_file("main.py", 100, "h1", 10, Some("py"))
            .unwrap();
        db.insert_symbols(
            fid,
            &[
                NewSymbol {
                    name: "ServiceManager",
                    kind: "class",
                    line: 1,
                    column_num: 0,
                    start_byte: 0,
                    end_byte: 100,
                    signature: "class ServiceManager:",
                    name_path: "ServiceManager",
                    parent_id: None,
                },
                NewSymbol {
                    name: "run_service",
                    kind: "function",
                    line: 10,
                    column_num: 0,
                    start_byte: 101,
                    end_byte: 200,
                    signature: "def run_service():",
                    name_path: "run_service",
                    parent_id: None,
                },
                NewSymbol {
                    name: "helper",
                    kind: "function",
                    line: 20,
                    column_num: 0,
                    start_byte: 201,
                    end_byte: 300,
                    signature: "def helper():",
                    name_path: "helper",
                    parent_id: None,
                },
            ],
        )
        .unwrap();

        let project = ProjectRoot::new(root.to_str().unwrap()).unwrap();
        (temp, project)
    }

    #[test]
    fn exact_match_gets_highest_score() {
        let (_root, project) = make_project_with_symbols();
        let results = search_symbols_hybrid(&project, "ServiceManager", 10, 0.6).unwrap();
        assert!(!results.is_empty());
        assert_eq!(results[0].name, "ServiceManager");
        assert_eq!(results[0].match_type, "exact");
        assert_eq!(results[0].score, 100.0);
    }

    #[test]
    fn substring_match_returns_bm25_type() {
        let (_root, project) = make_project_with_symbols();
        // "service" is a substring of "ServiceManager" and "run_service"
        // threshold 0.99 ensures fuzzy won't fire, so only exact/fts/substring contribute
        let results = search_symbols_hybrid(&project, "service", 10, 0.99).unwrap();
        let text_matches: Vec<_> = results
            .iter()
            .filter(|r| r.match_type == "substring" || r.match_type == "fts")
            .collect();
        assert!(!text_matches.is_empty());
    }

    #[test]
    fn fuzzy_match_finds_approximate_name() {
        let (_root, project) = make_project_with_symbols();
        // "helpr" is close to "helper" via jaro_winkler
        let results = search_symbols_hybrid(&project, "helpr", 10, 0.7).unwrap();
        let fuzzy: Vec<_> = results.iter().filter(|r| r.match_type == "fuzzy").collect();
        assert!(!fuzzy.is_empty(), "expected a fuzzy match for 'helpr'");
        assert_eq!(fuzzy[0].name, "helper");
    }

    #[test]
    fn results_sorted_by_score_descending() {
        let (_root, project) = make_project_with_symbols();
        let results = search_symbols_hybrid(&project, "run_service", 20, 0.5).unwrap();
        for window in results.windows(2) {
            assert!(window[0].score >= window[1].score);
        }
    }

    #[test]
    fn no_duplicates_in_results() {
        let (_root, project) = make_project_with_symbols();
        let results = search_symbols_hybrid(&project, "ServiceManager", 20, 0.5).unwrap();
        let mut keys = std::collections::HashSet::new();
        for r in &results {
            let key = (r.name.clone(), r.file.clone(), r.line);
            assert!(keys.insert(key), "duplicate entry found");
        }
    }

    #[test]
    fn semantic_scores_add_new_results() {
        let (_root, project) = make_project_with_symbols();
        let mut scores = std::collections::HashMap::new();
        // "helper" wouldn't match "authentication" textually, but semantic says it's relevant
        scores.insert("main.py:helper".to_owned(), 0.8);

        let results = search_symbols_hybrid_with_semantic(
            &project,
            "authentication",
            10,
            0.99, // high fuzzy threshold to disable fuzzy path
            Some(&scores),
            None,
        )
        .unwrap();

        let semantic_matches: Vec<_> = results
            .iter()
            .filter(|r| r.match_type == "semantic")
            .collect();
        assert!(
            !semantic_matches.is_empty(),
            "semantic path should surface 'helper' for 'authentication' query"
        );
        assert_eq!(semantic_matches[0].name, "helper");
        assert!(semantic_matches[0].score > 0.0);
    }

    #[test]
    fn semantic_scores_boost_existing_results() {
        let (_root, project) = make_project_with_symbols();
        // Get baseline score for exact match
        let baseline = search_symbols_hybrid(&project, "ServiceManager", 10, 0.5).unwrap();
        let baseline_score = baseline[0].score;

        // Now add semantic boost
        let mut scores = std::collections::HashMap::new();
        scores.insert("main.py:ServiceManager".to_owned(), 0.9);

        let boosted = search_symbols_hybrid_with_semantic(
            &project,
            "ServiceManager",
            10,
            0.5,
            Some(&scores),
            None,
        )
        .unwrap();

        assert!(
            boosted[0].score > baseline_score,
            "semantic boost should increase score: {} > {}",
            boosted[0].score,
            baseline_score
        );
    }

    #[test]
    fn semantic_low_scores_filtered_out() {
        let (_root, project) = make_project_with_symbols();
        let mut scores = std::collections::HashMap::new();
        // Score below 0.15 threshold should not produce semantic match
        scores.insert("main.py:helper".to_owned(), 0.1);

        let results = search_symbols_hybrid_with_semantic(
            &project,
            "unrelated_query_xyz",
            10,
            0.99,
            Some(&scores),
            None,
        )
        .unwrap();

        let semantic_matches: Vec<_> = results
            .iter()
            .filter(|r| r.match_type == "semantic")
            .collect();
        assert!(
            semantic_matches.is_empty(),
            "low semantic scores should not surface results"
        );
    }

    #[test]
    fn no_duplicates_with_semantic() {
        let (_root, project) = make_project_with_symbols();
        let mut scores = std::collections::HashMap::new();
        scores.insert("main.py:ServiceManager".to_owned(), 0.9);
        scores.insert("main.py:helper".to_owned(), 0.7);

        let results = search_symbols_hybrid_with_semantic(
            &project,
            "ServiceManager",
            20,
            0.5,
            Some(&scores),
            None,
        )
        .unwrap();

        let mut keys = std::collections::HashSet::new();
        for r in &results {
            let key = (r.name.clone(), r.file.clone(), r.line);
            assert!(keys.insert(key.clone()), "duplicate entry found: {:?}", key);
        }
    }

    #[test]
    fn pagerank_boost_max_normalized_by_top_file() {
        let mut results = vec![
            SearchResult {
                name: "a".into(),
                kind: "function".into(),
                file: "popular.py".into(),
                line: 1,
                signature: "".into(),
                name_path: "a".into(),
                score: 50.0,
                match_type: "fts".into(),
            },
            SearchResult {
                name: "b".into(),
                kind: "function".into(),
                file: "obscure.py".into(),
                line: 2,
                signature: "".into(),
                name_path: "b".into(),
                score: 50.0,
                match_type: "fts".into(),
            },
        ];
        let mut pr = std::collections::HashMap::new();
        pr.insert("popular.py".into(), 0.4);
        pr.insert("obscure.py".into(), 0.05);
        apply_pagerank_boost(&mut results, &pr, 5.0);
        // popular.py is the max → +5.0; obscure.py is at 0.05/0.4 → +0.625
        assert!((results[0].score - 55.0).abs() < 1e-6);
        assert!((results[1].score - 50.625).abs() < 1e-6);
    }

    #[test]
    fn pagerank_boost_skips_unranked_files() {
        let mut results = vec![SearchResult {
            name: "a".into(),
            kind: "function".into(),
            file: "unmapped.py".into(),
            line: 1,
            signature: "".into(),
            name_path: "a".into(),
            score: 50.0,
            match_type: "fts".into(),
        }];
        let mut pr = std::collections::HashMap::new();
        pr.insert("other.py".into(), 0.3);
        apply_pagerank_boost(&mut results, &pr, 5.0);
        assert_eq!(results[0].score, 50.0);
    }

    #[test]
    fn pagerank_boost_zero_max_is_noop() {
        let mut results = vec![SearchResult {
            name: "a".into(),
            kind: "function".into(),
            file: "x.py".into(),
            line: 1,
            signature: "".into(),
            name_path: "a".into(),
            score: 50.0,
            match_type: "fts".into(),
        }];
        let pr = std::collections::HashMap::new();
        apply_pagerank_boost(&mut results, &pr, 5.0);
        assert_eq!(results[0].score, 50.0);
    }
}