gitcortex_mcp/mcp/

search.rs

//! Fuzzy search over the graph — multi-signal ranking with CamelCase/snake_case
//! tokenisation, token overlap scoring, and edit-distance typo tolerance.
//!
//! Ranking signals (higher score = better match):
//! - exact name match:                   +100
//! - prefix name match:                  +60
//! - all query tokens match name tokens: +50
//! - substring in name:                  +30
//! - partial token overlap:              +10..+25
//! - edit distance ≤1 (typo):            +20
//! - edit distance ≤2:                   +10
//! - substring in qualified_name only:   +10
//! - shorter names break ties
//! - kind boost: Function/Method/Struct/Trait > others

use std::collections::HashSet;

use gitcortex_core::{error::Result, graph::Node, schema::NodeKind, store::GraphStore};
use serde::Serialize;

#[derive(Debug, Clone, Serialize)]
pub struct SearchHit {
    pub name: String,
    pub qualified_name: String,
    pub kind: String,
    pub file: String,
    pub start_line: u32,
    pub score: i32,
}

/// Hits grouped by file — lets a model see "parse is in parser.rs (8 symbols)"
/// at a glance without iterating the full flat hit list.
#[derive(Debug, Clone, Serialize)]
pub struct FileGroup {
    pub file: String,
    pub symbol_count: usize,
    /// Up to 3 highest-scoring symbol names in this file.
    pub top_symbols: Vec<String>,
}

/// Group search hits by file. Groups are sorted by `symbol_count` descending
/// (most hits first), ties broken alphabetically by file path.
///
/// Input must already be sorted by score descending so `top_symbols` picks the
/// highest-scoring symbols per file without extra sorting.
pub fn group_by_file(hits: &[SearchHit]) -> Vec<FileGroup> {
    // Preserve insertion order (first-seen = highest-scored) per file.
    let mut order: Vec<String> = Vec::new();
    let mut map: std::collections::HashMap<String, (usize, Vec<String>)> =
        std::collections::HashMap::new();

    for h in hits {
        let e = map.entry(h.file.clone()).or_insert_with(|| {
            order.push(h.file.clone());
            (0, Vec::new())
        });
        e.0 += 1;
        if e.1.len() < 3 {
            e.1.push(h.name.clone());
        }
    }

    let mut groups: Vec<FileGroup> = order
        .into_iter()
        .map(|file| {
            let (count, top) = map.remove(&file).unwrap();
            FileGroup {
                file,
                symbol_count: count,
                top_symbols: top,
            }
        })
        .collect();

    groups.sort_by(|a, b| {
        b.symbol_count
            .cmp(&a.symbol_count)
            .then_with(|| a.file.cmp(&b.file))
    });
    groups
}

const DEFAULT_LIMIT: usize = 10;
const MAX_LIMIT: usize = 200;
const MIN_TOKEN_LEN: usize = 3;

/// Split a camelCase/snake_case/PascalCase identifier into lowercase tokens.
///
/// "AuthConfig"      → ["auth", "config"]
/// "validate_token"  → ["validate", "token"]
/// "parseJSONResponse" → ["parse", "j", "s", "o", "n", "response"]  (intentional — acronyms split per char)
/// "HTTPClient"      → ["h", "t", "t", "p", "client"]
pub(crate) fn tokenize(s: &str) -> Vec<String> {
    let mut tokens = Vec::new();
    let mut current = String::new();
    let chars: Vec<char> = s.chars().collect();
    for (i, &ch) in chars.iter().enumerate() {
        if ch == '_' || ch == '-' || ch == '.' || ch == ':' || ch == '/' || ch == ' ' {
            if !current.is_empty() {
                tokens.push(current.to_ascii_lowercase());
                current = String::new();
            }
        } else if ch.is_uppercase() {
            // Start new token on uppercase — but keep run of capitals together
            // as one token (e.g. "HTTP" stays "http" not split per char).
            let next_is_lower = chars.get(i + 1).map(|c| c.is_lowercase()).unwrap_or(false);
            let prev_is_upper = i > 0 && chars[i - 1].is_uppercase();
            if !current.is_empty() && (!prev_is_upper || next_is_lower) {
                tokens.push(current.to_ascii_lowercase());
                current = String::new();
            }
            current.push(ch.to_ascii_lowercase());
        } else {
            current.push(ch);
        }
    }
    if !current.is_empty() {
        tokens.push(current.to_ascii_lowercase());
    }
    tokens
}

/// Levenshtein edit distance between two strings (capped early at `max`).
fn edit_distance(a: &str, b: &str) -> usize {
    let a: Vec<char> = a.chars().collect();
    let b: Vec<char> = b.chars().collect();
    let m = a.len();
    let n = b.len();
    // Quick bounds: length difference alone is a lower bound.
    if m.abs_diff(n) > 3 {
        return usize::MAX;
    }
    let mut prev: Vec<usize> = (0..=n).collect();
    let mut curr = vec![0usize; n + 1];
    for i in 1..=m {
        curr[0] = i;
        for j in 1..=n {
            curr[j] = if a[i - 1] == b[j - 1] {
                prev[j - 1]
            } else {
                1 + prev[j - 1].min(prev[j]).min(curr[j - 1])
            };
        }
        std::mem::swap(&mut prev, &mut curr);
    }
    prev[n]
}

/// Score a node against a query. Returns `None` when the node is not a match.
fn score(n: &Node, q_lower: &str, q_tokens: &[String]) -> Option<i32> {
    let name_lower = n.name.to_ascii_lowercase();
    let qname_lower = n.qualified_name.to_ascii_lowercase();
    let name_tokens = tokenize(&n.name);

    let base = if name_lower == q_lower {
        // Exact name match — highest confidence.
        100
    } else if name_lower.starts_with(q_lower) {
        60
    } else if !q_tokens.is_empty() && q_tokens.iter().all(|t| name_tokens.contains(t)) {
        // All query tokens present in name tokens.
        // "auth config" fully matches "AuthConfig" or "auth_config".
        50
    } else if name_lower.contains(q_lower) {
        30
    } else {
        // Partial token overlap.
        let overlap = q_tokens
            .iter()
            .filter(|qt| qt.len() >= MIN_TOKEN_LEN && name_tokens.contains(*qt))
            .count();
        if overlap > 0 {
            10 + (overlap as i32 * 5).min(15)
        } else if qname_lower.contains(q_lower) {
            // Match only in qualified path (e.g. module prefix).
            10
        } else if q_lower.len() >= 4 && q_lower.len() <= 15 && name_lower.len() <= 25 {
            // Typo tolerance: edit distance on short-ish queries.
            let dist = edit_distance(q_lower, &name_lower);
            if dist <= 1 {
                20
            } else if dist <= 2 {
                10
            } else {
                return None;
            }
        } else {
            return None;
        }
    };

    Some(base + kind_boost(&n.kind))
}

fn kind_boost(k: &NodeKind) -> i32 {
    match k {
        NodeKind::Function | NodeKind::Method => 5,
        NodeKind::Struct | NodeKind::Trait | NodeKind::Interface => 4,
        NodeKind::Enum | NodeKind::TypeAlias => 3,
        NodeKind::Constant | NodeKind::Macro | NodeKind::Annotation => 2,
        _ => 0,
    }
}

fn to_hit(n: Node, score: i32) -> SearchHit {
    SearchHit {
        name: n.name,
        qualified_name: n.qualified_name,
        kind: n.kind.to_string(),
        file: n.file.display().to_string(),
        start_line: n.span.start_line,
        score,
    }
}

/// Run a fuzzy search across all nodes on `branch`.
///
/// Candidate set is built by querying the store for the whole query string AND
/// for each individual token (for multi-word / camelCase queries). Candidates
/// are deduplicated, scored with the multi-signal scorer, sorted by score
/// descending, and truncated to `limit`.
pub fn search<S: GraphStore + ?Sized>(
    store: &S,
    branch: &str,
    query: &str,
    limit: Option<usize>,
) -> Result<Vec<SearchHit>> {
    let limit = limit.unwrap_or(DEFAULT_LIMIT).min(MAX_LIMIT);
    let q = query.trim();
    if q.is_empty() {
        return Ok(Vec::new());
    }

    let q_lower = q.to_ascii_lowercase();
    let q_tokens = tokenize(q);
    let candidate_limit = (limit * 50).max(500);

    // Fetch candidates: whole query first, then per token.
    let mut seen: HashSet<String> = HashSet::new();
    let mut nodes: Vec<Node> = Vec::new();

    let push = |nodes: &mut Vec<Node>, seen: &mut HashSet<String>, batch: Vec<Node>| {
        for n in batch {
            let id = n.id.as_str();
            if seen.insert(id) {
                nodes.push(n);
            }
        }
    };

    push(
        &mut nodes,
        &mut seen,
        store.search_nodes(branch, q, candidate_limit)?,
    );

    // Per-token expansion: lets "validate token" find "validate_token" even
    // when the store's CONTAINS filter requires the full substring.
    for token in &q_tokens {
        if token.len() < MIN_TOKEN_LEN {
            continue;
        }
        // Skip if token equals the whole query (already fetched above).
        if token.as_str() == q_lower {
            continue;
        }
        push(
            &mut nodes,
            &mut seen,
            store.search_nodes(branch, token, candidate_limit)?,
        );
    }

    // Typo-fallback: CONTAINS can't find misspelled queries ("Greetter" won't
    // match "Greeter"). When no candidates found and query is short enough for
    // edit-distance to be meaningful, scan all nodes so the scorer can apply
    // typo tolerance.
    if nodes.is_empty() && q_lower.len() >= 4 && q_lower.len() <= 20 {
        push(&mut nodes, &mut seen, store.list_all_nodes(branch)?);
    }

    let mut hits: Vec<SearchHit> = nodes
        .into_iter()
        .filter_map(|n| score(&n, &q_lower, &q_tokens).map(|s| to_hit(n, s)))
        .collect();

    hits.sort_by(|a, b| {
        b.score
            .cmp(&a.score)
            .then_with(|| a.name.len().cmp(&b.name.len()))
            .then_with(|| a.qualified_name.cmp(&b.qualified_name))
    });
    hits.truncate(limit);
    Ok(hits)
}

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

    #[test]
    fn tokenize_camel_case() {
        assert_eq!(tokenize("AuthConfig"), vec!["auth", "config"]);
        assert_eq!(tokenize("validateToken"), vec!["validate", "token"]);
        assert_eq!(tokenize("HTTPClient"), vec!["http", "client"]);
    }

    #[test]
    fn tokenize_snake_case() {
        assert_eq!(tokenize("validate_token"), vec!["validate", "token"]);
        assert_eq!(tokenize("auth_middleware"), vec!["auth", "middleware"]);
    }

    #[test]
    fn tokenize_pascal_case() {
        assert_eq!(tokenize("KuzuGraphStore"), vec!["kuzu", "graph", "store"]);
    }

    #[test]
    fn edit_distance_exact() {
        assert_eq!(edit_distance("validate", "validate"), 0);
    }

    #[test]
    fn edit_distance_typo() {
        assert_eq!(edit_distance("vlidate", "validate"), 1);
        assert_eq!(edit_distance("authnticate", "authenticate"), 1);
    }

    #[test]
    fn edit_distance_length_short_circuit() {
        // length difference > 3 → MAX
        assert_eq!(edit_distance("a", "abcde"), usize::MAX);
    }

    // ── group_by_file ─────────────────────────────────────────────────────────

    fn hit(name: &str, file: &str, score: i32) -> SearchHit {
        SearchHit {
            name: name.to_owned(),
            qualified_name: name.to_owned(),
            kind: "Function".to_owned(),
            file: file.to_owned(),
            start_line: 1,
            score,
        }
    }

    #[test]
    fn group_by_file_empty_returns_empty() {
        assert!(group_by_file(&[]).is_empty());
    }

    #[test]
    fn group_by_file_single_file() {
        let hits = vec![hit("parse_args", "src/parser.rs", 100)];
        let groups = group_by_file(&hits);
        assert_eq!(groups.len(), 1);
        assert_eq!(groups[0].file, "src/parser.rs");
        assert_eq!(groups[0].symbol_count, 1);
        assert_eq!(groups[0].top_symbols, vec!["parse_args"]);
    }

    #[test]
    fn group_by_file_sorted_by_count_desc() {
        let hits = vec![
            hit("a", "src/big.rs", 80),
            hit("b", "src/big.rs", 70),
            hit("c", "src/big.rs", 60),
            hit("x", "src/small.rs", 50),
        ];
        let groups = group_by_file(&hits);
        assert_eq!(groups[0].file, "src/big.rs");
        assert_eq!(groups[0].symbol_count, 3);
        assert_eq!(groups[1].file, "src/small.rs");
        assert_eq!(groups[1].symbol_count, 1);
    }

    #[test]
    fn group_by_file_top_symbols_capped_at_three() {
        let hits: Vec<SearchHit> = (0..8)
            .map(|i| hit(&format!("fn{i}"), "src/big.rs", 100 - i))
            .collect();
        let groups = group_by_file(&hits);
        assert_eq!(groups[0].symbol_count, 8);
        assert_eq!(groups[0].top_symbols.len(), 3);
        // First three are highest-scored (fn0, fn1, fn2).
        assert_eq!(groups[0].top_symbols[0], "fn0");
    }

    #[test]
    fn group_by_file_alphabetical_tie_break() {
        let hits = vec![hit("a", "src/z.rs", 50), hit("b", "src/a.rs", 50)];
        let groups = group_by_file(&hits);
        // Both have count=1; alphabetical tie-break: a.rs before z.rs.
        assert_eq!(groups[0].file, "src/a.rs");
    }
}