scryer-engine 0.2.0

//! Ranked symbol search over the cached BM25 indexes.

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use scryer_db::{DependencyPackage, ProjectDependency, ScryerDb};

use super::cache::{IndexKey, SearchIndexCache, SymbolDoc};
use super::tokenize::{tokenize_identifier, tokenize_query, tokenize_text};

/// Where to search.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub enum SearchScope {
    /// Symbols of the project itself.
    #[default]
    Project,
    /// The project's dependencies in the global (project 0) cache, optionally one crate.
    Dependencies { crate_name: Option<String> },
}

/// A symbol search request. Exactly one of `text` / `like_symbol` must be set.
#[derive(Debug, Clone, Default)]
pub struct SearchQuery {
    pub text: Option<String>,
    /// Find symbols similar to this one (by name or qualified name), excluding itself.
    pub like_symbol: Option<String>,
    /// Case-insensitive kind filter (e.g. `fn`, `struct`); empty means any.
    pub kinds: Vec<String>,
    /// Case-insensitive substring filter on the file path.
    pub file_filter: Option<String>,
    pub scope: SearchScope,
    pub offset: usize,
    pub limit: usize,
}

/// One ranked symbol.
#[derive(Debug, Clone)]
pub struct SymbolHit {
    pub name: String,
    pub qualified_name: String,
    pub kind: String,
    pub file_path: String,
    pub start_line: u32,
    pub end_line: u32,
    pub signature: String,
    pub score: f32,
    pub matched_terms: Vec<String>,
    /// Set for dependency hits.
    pub crate_name: Option<String>,
    pub version: Option<String>,
}

/// A page of ranked symbols.
#[derive(Debug, Clone, Default)]
pub struct SearchResults {
    /// The tokens actually searched for.
    pub query_terms: Vec<String>,
    /// Qualified name of the `like_symbol` target, when used.
    pub like_symbol: Option<String>,
    pub total: usize,
    pub offset: usize,
    pub has_more: bool,
    pub hits: Vec<SymbolHit>,
}

/// `dependency_package_id` → `(crate name, version)`.
type PackageNames = HashMap<u64, (String, String)>;

/// Packages the dependency scope resolved to, with display names.
struct DependencyScope {
    packages: PackageNames,
}

async fn resolve_dependency_scope(
    db: &ScryerDb,
    project_id: u64,
    crate_name: Option<&str>,
) -> anyhow::Result<DependencyScope> {
    let mut guard = db.lock().await;
    let linked: HashSet<u64> =
        ProjectDependency::filter(ProjectDependency::fields().project_id().eq(project_id))
            .exec(&mut *guard)
            .await?
            .into_iter()
            .map(|l| l.dependency_package_id)
            .collect();
    let all = DependencyPackage::all().exec(&mut *guard).await?;

    let selected: Vec<DependencyPackage> = match crate_name {
        Some(name) => {
            let named: Vec<DependencyPackage> = all
                .into_iter()
                .filter(|p| p.name.eq_ignore_ascii_case(name))
                .collect();
            // Prefer the versions this project links; fall back to any cached version.
            if named.iter().any(|p| linked.contains(&p.id)) {
                named
                    .into_iter()
                    .filter(|p| linked.contains(&p.id))
                    .collect()
            } else {
                named
            }
        }
        None => all.into_iter().filter(|p| linked.contains(&p.id)).collect(),
    };

    Ok(DependencyScope {
        packages: selected
            .into_iter()
            .map(|p| (p.id, (p.name, p.version)))
            .collect(),
    })
}

/// Pick the `like_symbol` target: exact name, then qualified name, then qualified suffix.
fn find_like_target<'a>(
    docs: &'a [SymbolDoc],
    target: &str,
    in_scope: &impl Fn(&SymbolDoc) -> bool,
) -> Option<&'a SymbolDoc> {
    let suffix = format!("::{target}");
    docs.iter()
        .filter(|d| in_scope(d))
        .find(|d| d.name == target)
        .or_else(|| {
            docs.iter()
                .filter(|d| in_scope(d))
                .find(|d| d.qualified_name == target)
        })
        .or_else(|| {
            docs.iter()
                .filter(|d| in_scope(d))
                .find(|d| d.qualified_name.ends_with(&suffix))
        })
}

/// Query tokens describing a symbol: its name, qualified path, signature and doc snippet.
fn like_terms(doc: &SymbolDoc) -> Vec<String> {
    let mut terms = Vec::new();
    for t in tokenize_identifier(&doc.name)
        .into_iter()
        .chain(tokenize_text(doc.parent_path()))
        .chain(tokenize_text(&doc.signature))
        .chain(tokenize_text(&doc.doc_snippet))
    {
        if !terms.contains(&t) {
            terms.push(t);
        }
    }
    terms
}

pub(crate) async fn search_symbols(
    db: &ScryerDb,
    rayon_pool: &Arc<rayon::ThreadPool>,
    cache: &SearchIndexCache,
    project_id: u64,
    query: SearchQuery,
) -> anyhow::Result<SearchResults> {
    anyhow::ensure!(
        query.text.is_some() != query.like_symbol.is_some(),
        "Exactly one of `query` or `like_symbol` must be provided"
    );

    let (key, packages): (IndexKey, Option<PackageNames>) = match &query.scope {
        SearchScope::Project => ((project_id, None), None),
        SearchScope::Dependencies { crate_name } => {
            let scope = resolve_dependency_scope(db, project_id, crate_name.as_deref()).await?;
            if scope.packages.is_empty() {
                return Ok(SearchResults {
                    offset: query.offset,
                    ..SearchResults::default()
                });
            }
            let key = if scope.packages.len() == 1 {
                (0, scope.packages.keys().next().copied())
            } else {
                (0, None)
            };
            (key, Some(scope.packages))
        }
    };

    let index = cache.get_or_build(db, rayon_pool, key).await?;

    let kinds: Vec<String> = query.kinds.iter().map(|k| k.to_lowercase()).collect();
    let file_filter = query.file_filter.as_deref().map(str::to_lowercase);
    let in_scope = |d: &SymbolDoc| {
        packages.as_ref().is_none_or(|pkgs| {
            d.dependency_package_id
                .is_some_and(|id| pkgs.contains_key(&id))
        })
    };

    let (terms, exclude_id, like_name) = match (&query.text, &query.like_symbol) {
        (Some(text), _) => (tokenize_query(text), None, None),
        (None, Some(target)) => {
            let doc = find_like_target(index.index.docs(), target, &in_scope).ok_or_else(|| {
                anyhow::anyhow!(
                    "Symbol `{target}` not found in this scope; try a `query` search instead"
                )
            })?;
            (
                like_terms(doc),
                Some(doc.id),
                Some(doc.qualified_name.clone()),
            )
        }
        (None, None) => unreachable!("validated above"),
    };

    let hits = index.index.search(&terms, |d| {
        in_scope(d)
            && Some(d.id) != exclude_id
            && (kinds.is_empty() || kinds.contains(&d.kind.to_lowercase()))
            && file_filter
                .as_ref()
                .is_none_or(|f| d.file_path.to_lowercase().contains(f))
    });

    let total = hits.len();
    let page: Vec<SymbolHit> = hits
        .into_iter()
        .skip(query.offset)
        .take(query.limit)
        .map(|hit| {
            let d = index.index.doc(hit.doc_idx);
            let pkg = packages
                .as_ref()
                .zip(d.dependency_package_id)
                .and_then(|(pkgs, id)| pkgs.get(&id));
            SymbolHit {
                name: d.name.clone(),
                qualified_name: d.qualified_name.clone(),
                kind: d.kind.clone(),
                file_path: d.file_path.clone(),
                start_line: d.start_line,
                end_line: d.end_line,
                signature: d.signature.clone(),
                score: hit.score,
                matched_terms: hit.matched_terms,
                crate_name: pkg.map(|(n, _)| n.clone()),
                version: pkg.map(|(_, v)| v.clone()),
            }
        })
        .collect();

    Ok(SearchResults {
        query_terms: terms,
        like_symbol: like_name,
        total,
        offset: query.offset,
        has_more: query.offset + page.len() < total,
        hits: page,
    })
}