scryer-engine 0.3.0

//! Index benchmark: how long indexing takes, what it stores, and whether that is the same on
//! every run. Used to measure each phase of Plan 18 (docs/learnings/index-benchmark.md).
//!
//! ```text
//! cargo run --release -p scryer-engine --example index_bench -- \
//!   <root>... [--runs 3] [--batch-cap 500]... [--seed N]... [--probe NAME]... \
//!   [--recall crates/scryer-engine/tests/recall] [--dependencies] \
//!   [--json out.json] [--baseline prev.json] [--dump DIR]
//! ```
//!
//! Every run indexes into a fresh on-disk database. Runs of one root are compared by a
//! digest of the stored references and edges (database ids excluded); the process exits 1
//! when they disagree, so a non-deterministic index fails the benchmark.

#[path = "../tests/support/mod.rs"]
mod support;

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::process::ExitCode;
use std::time::Instant;

use scryer_db::ScryerDb;
use scryer_engine::IndexOptions;
use serde_json::{Value, json};
use support::fingerprint::{fingerprint, reference_count, sample_via};
use support::fixture::index_dir;
use support::recall;
use tempfile::tempdir;

const USAGE: &str = "usage: index_bench <root>... [--runs N] [--batch-cap N]... [--seed N]... \
[--probe NAME]... [--recall DIR] [--sample-via VIA N] [--dependencies] [--json FILE] [--baseline FILE] \
[--dump DIR]";

#[derive(Default)]
struct Args {
    roots: Vec<PathBuf>,
    runs: usize,
    caps: Vec<usize>,
    seeds: Vec<u64>,
    probes: Vec<String>,
    recall: Option<PathBuf>,
    sample_via: Option<(String, usize)>,
    dependencies: bool,
    json: Option<PathBuf>,
    baseline: Option<PathBuf>,
    /// Write the first run's references, edges, symbols and unresolved rows of every corpus
    /// as sorted text files, so two builds can be compared line by line.
    dump: Option<PathBuf>,
}

fn parse_args() -> Result<Args, String> {
    let mut args = Args {
        runs: 3,
        ..Args::default()
    };
    let mut it = std::env::args().skip(1);
    while let Some(arg) = it.next() {
        let mut value = |name: &str| it.next().ok_or(format!("{name} needs a value"));
        match arg.as_str() {
            "--runs" => args.runs = value("--runs")?.parse().map_err(|e| format!("{e}"))?,
            "--batch-cap" => args
                .caps
                .push(value("--batch-cap")?.parse().map_err(|e| format!("{e}"))?),
            "--seed" => args
                .seeds
                .push(value("--seed")?.parse().map_err(|e| format!("{e}"))?),
            "--probe" => args.probes.push(value("--probe")?),
            "--recall" => args.recall = Some(value("--recall")?.into()),
            "--sample-via" => {
                let via = value("--sample-via")?;
                let n = it
                    .next()
                    .ok_or("--sample-via needs a via and a count")?
                    .parse()
                    .map_err(|e| format!("{e}"))?;
                args.sample_via = Some((via, n));
            }
            "--json" => args.json = Some(value("--json")?.into()),
            "--baseline" => args.baseline = Some(value("--baseline")?.into()),
            "--dump" => args.dump = Some(value("--dump")?.into()),
            "--dependencies" => args.dependencies = true,
            "-h" | "--help" => return Err(String::new()),
            other if other.starts_with("--") => return Err(format!("unknown option {other}")),
            root => args.roots.push(root.into()),
        }
    }
    if args.roots.is_empty() && args.recall.is_none() {
        return Err("give at least one <root> or --recall DIR".into());
    }
    if args.caps.is_empty() {
        args.caps.push(scryer_engine::DEFAULT_WRITER_BATCH_CAP);
    }
    Ok(args)
}

fn median(values: &mut [f64]) -> f64 {
    values.sort_by(|a, b| a.total_cmp(b));
    values[values.len() / 2]
}

fn corpus_name(root: &Path) -> String {
    // `work/cargo/source` is named after the suite, not its generic leaf directory.
    let leaf = root.file_name().and_then(|n| n.to_str()).unwrap_or("");
    let name = match (leaf, root.parent().and_then(|p| p.file_name())) {
        ("source", Some(parent)) => parent.to_string_lossy(),
        _ => leaf.into(),
    };
    name.to_string()
}

async fn bench_root(root: &Path, args: &Args) -> anyhow::Result<Value> {
    let seeds: Vec<Option<u64>> = if args.seeds.is_empty() {
        vec![None]
    } else {
        args.seeds.iter().copied().map(Some).collect()
    };
    let mut runs = Vec::new();
    println!("\n== {} ({})", corpus_name(root), root.display());
    println!(
        "{:>5} {:>6} {:>3} {:>8} {:>7} {:>8} {:>8}  digest            probes",
        "cap", "seed", "run", "seconds", "files", "refs", "edges"
    );
    for &cap in &args.caps {
        for &seed in &seeds {
            for run in 1..=args.runs {
                let dir = tempdir()?;
                let db = ScryerDb::connect_file(&dir.path().join("bench.db")).await?;
                let options = IndexOptions {
                    writer_batch_cap: cap,
                    file_order_seed: seed,
                    dependencies: args.dependencies,
                };
                let started = Instant::now();
                let (project, report) = index_dir(&db, root, options).await?;
                let seconds = started.elapsed().as_secs_f64();
                let fp = fingerprint(&db, project).await?;
                let health = scryer_engine::index_health(&db, project).await?;
                let syntax_errors = health
                    .stats
                    .as_ref()
                    .map(|s| s.syntax_error_files.clone())
                    .unwrap_or_default();
                let mut probes = BTreeMap::new();
                for name in &args.probes {
                    probes.insert(name.clone(), reference_count(&db, project, name).await?);
                }
                let probe_text: Vec<String> =
                    probes.iter().map(|(k, v)| format!("{k}={v}")).collect();
                println!(
                    "{cap:>5} {:>6} {run:>3} {seconds:>8.2} {:>7} {:>8} {:>8}  {}  {}",
                    seed.map_or("-".into(), |s| s.to_string()),
                    report.scanned_files,
                    fp.references.len(),
                    fp.edges.len(),
                    fp.digest(),
                    probe_text.join(" ")
                );
                if let (Some(dir), true) = (&args.dump, runs.is_empty()) {
                    dump_fingerprint(dir, &corpus_name(root), &fp)?;
                    let mut text = syntax_errors.join("\n");
                    text.push('\n');
                    std::fs::write(
                        dir.join(format!("{}.syntax-errors.txt", corpus_name(root))),
                        text,
                    )?;
                }
                if let (Some((via, n)), true) = (&args.sample_via, run == 1) {
                    println!("  sample of {n} references with via={via}:");
                    for line in sample_via(&db, project, via, *n).await? {
                        println!("    {line}");
                    }
                }
                runs.push(json!({
                    "cap": cap, "seed": seed, "run": run, "seconds": seconds,
                    "files": report.scanned_files, "symbols": fp.symbols,
                    "references": fp.references.len(), "edges": fp.edges.len(),
                    "digest": fp.digest(), "probes": probes, "roles": fp.refs_by_role,
                    "vias": fp.refs_by_via, "syntax_error_files": syntax_errors.len(),
                }));
            }
        }
    }
    let mut digests: Vec<&str> = runs.iter().map(|r| r["digest"].as_str().unwrap()).collect();
    digests.sort_unstable();
    digests.dedup();
    let mut seconds: Vec<f64> = runs
        .iter()
        .map(|r| r["seconds"].as_f64().unwrap())
        .collect();
    let refs: Vec<u64> = runs
        .iter()
        .map(|r| r["references"].as_u64().unwrap())
        .collect();
    let deterministic = digests.len() == 1;
    if let Some(first) = runs.first() {
        println!(
            "references by via (first run): {}   files with syntax errors: {}",
            first["vias"], first["syntax_error_files"]
        );
    }
    println!(
        "deterministic: {}   median {:.2}s   references {}..{}",
        if deterministic {
            "yes".to_string()
        } else {
            format!("NO ({} distinct digests)", digests.len())
        },
        median(&mut seconds),
        refs.iter().min().unwrap(),
        refs.iter().max().unwrap()
    );
    Ok(json!({
        "root": root.to_string_lossy(), "runs": runs, "deterministic": deterministic,
        "distinct_digests": digests.len(), "median_seconds": median(&mut seconds),
        "references_min": refs.iter().min(), "references_max": refs.iter().max(),
    }))
}

/// `<dir>/<corpus>.{refs,edges,symbols,unresolved}.txt`, one sorted fact per line.
fn dump_fingerprint(
    dir: &Path,
    corpus: &str,
    fp: &support::fingerprint::Fingerprint,
) -> anyhow::Result<()> {
    std::fs::create_dir_all(dir)?;
    for (suffix, lines) in [
        ("refs", &fp.references),
        ("edges", &fp.edges),
        ("symbols", &fp.symbol_rows),
        ("unresolved", &fp.unresolved),
    ] {
        let mut text = lines.join("\n");
        text.push('\n');
        std::fs::write(dir.join(format!("{corpus}.{suffix}.txt")), text)?;
    }
    Ok(())
}

async fn bench_recall(dir: &Path) -> anyhow::Result<Value> {
    let mut out = BTreeMap::new();
    let mut line = Vec::new();
    for language in ["rust", "python", "typescript"] {
        let root = dir.join(language);
        if !root.is_dir() {
            continue;
        }
        let r = recall::run(language, &root, IndexOptions::default()).await?;
        line.push(format!("{language} {}/{}", r.found(), r.expected.len()));
        out.insert(
            language,
            json!({"found": r.found(), "expected": r.expected.len()}),
        );
    }
    println!("\n== recall fixtures: {}", line.join("   "));
    Ok(json!(out))
}

/// Print how this run differs from an earlier `--json` file.
fn compare(current: &Value, baseline: &Value) {
    println!("\n== versus baseline");
    let Some(corpora) = current["corpora"].as_object() else {
        return;
    };
    for (name, now) in corpora {
        let Some(then) = baseline["corpora"].get(name) else {
            println!("{name}: not in baseline");
            continue;
        };
        let pct = |a: f64, b: f64| if b == 0.0 { 0.0 } else { (a - b) / b * 100.0 };
        let (t_now, t_then) = (
            now["median_seconds"].as_f64().unwrap_or(0.0),
            then["median_seconds"].as_f64().unwrap_or(0.0),
        );
        let (r_now, r_then) = (
            now["references_max"].as_f64().unwrap_or(0.0),
            then["references_max"].as_f64().unwrap_or(0.0),
        );
        println!(
            "{name}: time {t_then:.2}s -> {t_now:.2}s ({:+.1}%)   refs(max) {r_then} -> {r_now} ({:+.1}%)   deterministic {} -> {}",
            pct(t_now, t_then),
            pct(r_now, r_then),
            then["deterministic"],
            now["deterministic"]
        );
    }
    if let (Some(now), Some(then)) = (
        current["recall"].as_object(),
        baseline["recall"].as_object(),
    ) {
        for (lang, n) in now {
            println!(
                "recall {lang}: {} -> {} of {}",
                then.get(lang).map_or(Value::Null, |t| t["found"].clone()),
                n["found"],
                n["expected"]
            );
        }
    }
}

#[tokio::main]
async fn main() -> ExitCode {
    let args = match parse_args() {
        Ok(a) => a,
        Err(msg) => {
            if !msg.is_empty() {
                eprintln!("error: {msg}\n");
            }
            eprintln!("{USAGE}");
            return ExitCode::from(2);
        }
    };
    let mut corpora = serde_json::Map::new();
    let mut ok = true;
    for root in &args.roots {
        match bench_root(root, &args).await {
            Ok(v) => {
                ok &= v["deterministic"].as_bool().unwrap_or(false);
                corpora.insert(corpus_name(root), v);
            }
            Err(err) => {
                eprintln!("{}: {err:#}", root.display());
                return ExitCode::from(2);
            }
        }
    }
    let recall = match &args.recall {
        Some(dir) => match bench_recall(dir).await {
            Ok(v) => v,
            Err(err) => {
                eprintln!("recall: {err:#}");
                return ExitCode::from(2);
            }
        },
        None => Value::Null,
    };
    let result = json!({ "corpora": corpora, "recall": recall });
    if let Some(path) = &args.baseline {
        match std::fs::read_to_string(path)
            .map_err(anyhow::Error::from)
            .and_then(|t| Ok(serde_json::from_str::<Value>(&t)?))
        {
            Ok(baseline) => compare(&result, &baseline),
            Err(err) => eprintln!("baseline {}: {err}", path.display()),
        }
    }
    if let Some(path) = &args.json {
        if let Some(parent) = path.parent() {
            let _ = std::fs::create_dir_all(parent);
        }
        let text = serde_json::to_string_pretty(&result).expect("serializable");
        if let Err(err) = std::fs::write(path, text) {
            eprintln!("{}: {err}", path.display());
            return ExitCode::from(2);
        }
    }
    if ok {
        ExitCode::SUCCESS
    } else {
        ExitCode::from(1)
    }
}