rustytokens/src/index/pipeline.rs

//! Index pipeline: hash, run ctags, extract imports, write to SQLite in one transaction.

use std::collections::HashMap;
use std::path::Path;

use anyhow::{Context, Result};
use rayon::prelude::*;
use rusqlite::params;

use crate::ctags;
use crate::discovery::{self, SourceFile};
use crate::index::{hashing, imports, store};

#[derive(Debug, serde::Serialize)]
pub struct IndexStats {
    pub files: i64,
    pub symbols: i64,
    pub imports: i64,
}

/// Read the `[package] name` from `Cargo.toml`, normalized to the crate
/// identifier (`-` → `_`) used in `use <crate>::…`.
fn read_crate_name(project: &Path) -> Option<String> {
    let toml = std::fs::read_to_string(project.join("Cargo.toml")).ok()?;
    let mut in_package = false;
    for line in toml.lines() {
        let l = line.trim();
        if l.starts_with('[') {
            in_package = l == "[package]";
            continue;
        }
        if in_package
            && let Some(v) = l.strip_prefix("name").map(str::trim)
            && let Some(v) = v.strip_prefix('=').map(str::trim)
        {
            let name = v.trim_matches('"');
            if !name.is_empty() {
                return Some(name.replace('-', "_"));
            }
        }
    }
    None
}

/// (Re)index a project. `recreate` drops any existing index first.
pub fn run(project: &Path, db_path: &Path, recreate: bool) -> Result<IndexStats> {
    if recreate && db_path.exists() {
        std::fs::remove_file(db_path).context("remove old index")?;
    }
    let conn = store::open(db_path)?;

    let files = discovery::discover(project)?;

    let existing: HashMap<String, (i64, String)> = store::all_files(&conn)?
        .into_iter()
        .map(|f| (f.path.clone(), (f.id, f.content_hash.clone())))
        .collect();

    // Hash all files in parallel.
    let hashed: Vec<(String, String, i64, &'static str)> = files
        .par_iter()
        .map(|f| {
            let h = hashing::hash_file(&f.path).unwrap_or_else(|_| "0000000000000000".to_string());
            let m = hashing::mtime_secs(&f.path);
            (f.rel.clone(), h, m, f.language)
        })
        .collect();

    let by_rel: HashMap<&str, &SourceFile> = files.iter().map(|f| (f.rel.as_str(), f)).collect();
    let abs_to_rel: HashMap<String, String> = files
        .iter()
        .map(|f| (f.path.to_string_lossy().into_owned(), f.rel.clone()))
        .collect();

    let changed: Vec<&SourceFile> = hashed
        .iter()
        .filter(|(rel, h, _, _)| existing.get(rel).map(|(_, oh)| oh != h).unwrap_or(true))
        .filter_map(|(rel, ..)| by_rel.get(rel.as_str()).copied())
        .collect();

    let deleted_ids: Vec<i64> = existing
        .iter()
        .filter(|(rel, _)| !by_rel.contains_key(rel.as_str()))
        .map(|(_, (id, _))| *id)
        .collect();

    conn.execute_batch("BEGIN;")?;

    let mut file_ids: HashMap<String, i64> = HashMap::new();
    for (rel, h, m, lang) in &hashed {
        let id = store::upsert_file(&conn, rel, h, lang, *m)?;
        file_ids.insert(rel.clone(), id);
    }

    for id in &deleted_ids {
        store::delete_file(&conn, *id)?;
        conn.execute("DELETE FROM content_fts WHERE rowid = ?1", [id])?;
    }

    // Run ctags once over changed files.
    let changed_paths: Vec<&Path> = changed.iter().map(|f| f.path.as_path()).collect();
    let tags = ctags::run(&changed_paths)?;
    let mut tags_by_rel: HashMap<String, Vec<ctags::Tag>> = HashMap::new();
    for t in tags {
        if let Some(rel) = abs_to_rel.get(&t.path) {
            tags_by_rel.entry(rel.clone()).or_default().push(t);
        }
    }

    let resolver = imports::Resolver::build(
        file_ids.iter().map(|(k, v)| (k.clone(), *v)).collect(),
        read_crate_name(project).as_deref(),
    );

    for f in &changed {
        let fid = file_ids[&f.rel];
        let content = std::fs::read_to_string(&f.path).unwrap_or_default();
        let tags = tags_by_rel.get(&f.rel).cloned().unwrap_or_default();
        store::replace_symbols(&conn, fid, &tags_to_symbols(&tags, &content))?;
        store::clear_imports_for_file(&conn, fid)?;
        let dir = f
            .rel
            .rsplit_once('/')
            .map(|(d, _)| d)
            .unwrap_or("")
            .to_string();
        for imp in imports::extract(f) {
            if let Some(to) = resolver.resolve_with_context(&imp.target, &dir, imp.levels)
                && to != fid
            {
                store::insert_import(&conn, fid, to, &imp.import_type)?;
            }
        }
        conn.execute("DELETE FROM content_fts WHERE rowid = ?1", [fid])?;
        conn.execute(
            "INSERT INTO content_fts(rowid, path, content) VALUES (?1, ?2, ?3)",
            params![fid, f.rel, content],
        )?;
    }

    conn.execute_batch("COMMIT;")?;

    let (files_n, symbols_n, imports_n) = store::stats(&conn)?;
    Ok(IndexStats {
        files: files_n,
        symbols: symbols_n,
        imports: imports_n,
    })
}

fn tags_to_symbols(tags: &[ctags::Tag], content: &str) -> Vec<store::NewSymbol> {
    let lines: Vec<&str> = content.lines().collect();
    tags.iter()
        .filter_map(|t| {
            let line = t.line? as i64;
            let end = t.end.map(|e| e as i64).unwrap_or(line);
            let kind = t.kind.clone().unwrap_or_default();
            let end = extend_block_end(&lines, line, end, &kind);
            let name = t.name.clone();
            let qualified_name = match &t.scope {
                Some(s) if !name.starts_with(&format!("{s}.")) => format!("{s}.{name}"),
                _ => name.clone(),
            };
            Some(store::NewSymbol {
                name,
                qualified_name,
                kind,
                signature: t.signature.clone(),
                line_start: line,
                line_end: end,
            })
        })
        .collect()
}

/// ctags omits a body span for some kinds (e.g. Rust `fn`), leaving `end ==
/// start` even for a multi-line block. Extend brace-based blocks to their
/// closing brace so stored spans are correct for every reader.
fn extend_block_end(lines: &[&str], start: i64, end: i64, kind: &str) -> i64 {
    if !is_block_kind(kind) || end > start {
        return end;
    }
    // Semicolon-terminated declarations (`mod foo;`, `struct Foo;`) have no body.
    if let Some(line) = lines.get((start - 1) as usize)
        && line.trim_end().ends_with(';')
    {
        return end;
    }
    let mut depth = 0i64;
    let mut opened = false;
    let mut out = end;
    // ponytail: naive char brace count; braces in string/comment literals skew
    // depth. Swap for a tokenizer only if real spans come back wrong.
    for (i, line) in lines.iter().enumerate() {
        let n = (i as i64) + 1;
        if n < start {
            continue;
        }
        for c in line.chars() {
            match c {
                '{' => {
                    depth += 1;
                    opened = true;
                }
                '}' => depth -= 1,
                _ => {}
            }
        }
        if opened {
            out = n;
        }
        if opened && depth <= 0 {
            break;
        }
    }
    out
}

fn is_block_kind(kind: &str) -> bool {
    matches!(
        kind,
        "function"
            | "method"
            | "struct"
            | "class"
            | "interface"
            | "enum"
            | "trait"
            | "impl"
            | "implementation"
            | "module"
            | "namespace"
            | "union"
    )
}