weirflow 0.1.0

#!/usr/bin/env python3
"""
Performance budget checker for Weir benchmarks.

Reads benches/budgets.json, runs the corresponding benchmarks via cargo bench,
and fails if any benchmark exceeds its budget (baseline median × 1.15).

Usage:
    python3 benches/check_budgets.py
    python3 benches/check_budgets.py --quick
    python3 benches/check_budgets.py --bench-filter csr_normalize
"""

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path


def load_budgets(path: Path) -> dict[str, dict]:
    """Load budgets.json and return a flat dict mapping full_id → budget info."""
    with open(path) as f:
        data = json.load(f)
    budgets = {}
    for group_id, configs in data.items():
        for config_key, info in configs.items():
            full_id = info.get("full_id", f"{group_id}/{config_key}")
            budgets[full_id] = info
    return budgets


def run_benchmarks(cargo_args: list[str]) -> None:
    """Run cargo bench with the given extra arguments."""
    cmd = ["cargo", "bench", "--", "--noplot"] + cargo_args
    print(f"Running: {' '.join(cmd)}")
    result = subprocess.run(cmd)
    if result.returncode != 0:
        print("ERROR: cargo bench failed to run.", file=sys.stderr)
        sys.exit(1)


def gather_results(target_criterion: Path) -> dict[str, float]:
    """Walk target/criterion and gather median point estimates by full_id."""
    results = {}
    for estimates_path in target_criterion.rglob("new/estimates.json"):
        benchmark_path = estimates_path.parent / "benchmark.json"
        if not benchmark_path.exists():
            continue
        with open(benchmark_path) as f:
            benchmark = json.load(f)
        with open(estimates_path) as f:
            estimates = json.load(f)
        full_id = benchmark.get("full_id", benchmark.get("title", ""))
        median_ns = estimates.get("median", {}).get("point_estimate", 0.0)
        results[full_id] = median_ns
    return results


def check_budgets(budgets: dict[str, dict], results: dict[str, float]) -> bool:
    """Compare results against budgets. Return True if all pass."""
    passed = 0
    failed = 0
    missing = 0

    print("\n" + "=" * 70)
    print("PERFORMANCE BUDGET CHECK")
    print("=" * 70)

    for full_id in sorted(budgets.keys()):
        info = budgets[full_id]
        budget_ns = info["budget_ns"]
        baseline_ns = info["baseline_ns"]

        if full_id not in results:
            missing += 1
            print(f"  MISSING   {full_id:60s} (no result found)")
            continue

        actual_ns = results[full_id]
        ratio = actual_ns / baseline_ns if baseline_ns > 0 else 0.0
        status = "PASS" if actual_ns <= budget_ns else "FAIL"

        if status == "PASS":
            passed += 1
        else:
            failed += 1

        print(
            f"  {status:6s}  {full_id:60s}  "
            f"actual={actual_ns:12.2f} ns  "
            f"budget={budget_ns:12d} ns  "
            f"baseline={baseline_ns:12d} ns  "
            f"({ratio:.2f}x)"
        )

    print("=" * 70)
    print(f"Results: {passed} passed, {failed} failed, {missing} missing")
    print("=" * 70)

    if missing > 0:
        print("\nWARNING: Some benchmarks did not produce results.", file=sys.stderr)
        print("This may happen if a benchmark panics or is skipped.", file=sys.stderr)

    return failed == 0 and missing == 0


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Check Weir benchmark performance against budgets."
    )
    parser.add_argument(
        "--budgets",
        type=Path,
        default=Path("benches/budgets.json"),
        help="Path to budgets.json",
    )
    parser.add_argument(
        "--quick",
        action="store_true",
        help="Run a quicker CI-friendly benchmark (reduced sample size)",
    )
    parser.add_argument(
        "--bench-filter",
        dest="bench_filter",
        default=None,
        help="Only run benchmarks matching this filter (passed to cargo bench)",
    )
    parser.add_argument(
        "--target-dir",
        type=Path,
        default=Path("target/criterion"),
        help="Path to Criterion output directory",
    )
    parser.add_argument(
        "--no-run",
        action="store_true",
        help="Do not run benchmarks; only check existing results in target/criterion",
    )
    args = parser.parse_args()

    if not args.budgets.exists():
        print(f"ERROR: Budgets file not found: {args.budgets}", file=sys.stderr)
        return 1

    budgets = load_budgets(args.budgets)
    print(f"Loaded {len(budgets)} budget entries from {args.budgets}")

    if not args.no_run:
        cargo_args = []
        if args.quick:
            cargo_args.extend(["--sample-size", "20", "--measurement-time", "1"])
        if args.bench_filter:
            # Criterion filters come after the -- separator, but cargo bench
            # also accepts a BENCHNAME before --. We pass the filter as cargo's
            # BENCHNAME to select the right benchmark binary, then pass any
            # Criterion args after --.
            cmd = ["cargo", "bench", args.bench_filter, "--", "--noplot"] + cargo_args
        else:
            cmd = ["cargo", "bench", "--", "--noplot"] + cargo_args

        print(f"Running: {' '.join(cmd)}")
        result = subprocess.run(cmd)
        if result.returncode != 0:
            print("ERROR: cargo bench failed.", file=sys.stderr)
            return 1

    results = gather_results(args.target_dir)
    print(f"Gathered {len(results)} benchmark results from {args.target_dir}")

    ok = check_budgets(budgets, results)
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())