tests/bench_lean.py

"""
Run inside Docker container:
  docker run --rm leanmigrate python tests/bench_lean.py

Decision rule:
  stdin p50 < 500ms  -> ship with stdin backend
  stdin p50 > 800ms  -> switch to kimina as default
"""

from __future__ import annotations

import os
import statistics
import subprocess
import textwrap
import time
from pathlib import Path
import shutil

REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_LEAN_BIN = str(Path(shutil.which("lean") or Path.home() / ".elan/bin/lean"))
DEFAULT_LEAN_PATH = str(REPO_ROOT / "lean" / ".lake" / "build" / "lib")
DEFAULT_CWD = REPO_ROOT / "lean"
DEFAULT_LAKE_BIN = str(Path(shutil.which("lake") or Path.home() / ".elan/bin/lake"))

LEAN_BIN = os.environ.get("LEAN_BIN", DEFAULT_LEAN_BIN)
LEAN_PATH = os.environ.get("LEAN_PATH", DEFAULT_LEAN_PATH)
CWD = Path(os.environ.get("LEAN_CWD", str(DEFAULT_CWD)))
LAKE_BIN = os.environ.get("LAKE_BIN", DEFAULT_LAKE_BIN)

TEST_CODE = textwrap.dedent(
    """
    import lean.AuthSpec
    open AuthSpec

    #check AuthSpec.canAccess
    """
)


def ensure_build() -> None:
    marker = Path(LEAN_PATH) / "AuthSpec.olean"
    if marker.exists():
        return

    process = subprocess.run(
        [LAKE_BIN, "build", "AuthSpec", "PricingSpec", "SagaSpec"],
        capture_output=True,
        text=True,
        timeout=120,
        cwd=str(CWD),
        env={**os.environ, "LEAN_PATH": LEAN_PATH},
    )
    if process.returncode != 0:
        raise RuntimeError(
            "Lean build failed before benchmarking.\n"
            f"stdout:\n{process.stdout}\n"
            f"stderr:\n{process.stderr}"
        )


def bench_stdin(n: int = 20) -> dict[str, float]:
    times: list[float] = []
    for _ in range(n):
        started = time.monotonic()
        result = subprocess.run(
            [LEAN_BIN, "--stdin"],
            input=TEST_CODE,
            capture_output=True,
            text=True,
            timeout=30,
            cwd=str(CWD),
            env={**os.environ, "LEAN_PATH": LEAN_PATH},
        )
        elapsed_ms = (time.monotonic() - started) * 1000
        times.append(elapsed_ms)
        if result.returncode != 0:
            print(f"  LEAN error: {result.stderr[:200]}")

    sorted_times = sorted(times)
    return {
        "p50": sorted_times[n // 2],
        "p95": sorted_times[int(n * 0.95) - 1],
        "min": min(times),
        "max": max(times),
        "mean": statistics.mean(times),
    }


if __name__ == "__main__":
    ensure_build()
    print(f"LEAN binary: {LEAN_BIN}")
    print(f"LEAN_PATH:   {LEAN_PATH}")
    print(f"LEAN_CWD:    {CWD}")
    print("\nRunning 20 calls to bench stdin backend...")

    stats = bench_stdin(20)
    print("\nResults:")
    print(f"  p50  = {stats['p50']:.0f}ms")
    print(f"  p95  = {stats['p95']:.0f}ms")
    print(f"  min  = {stats['min']:.0f}ms")
    print(f"  max  = {stats['max']:.0f}ms")
    print(f"  mean = {stats['mean']:.0f}ms")

    if stats["p50"] < 500:
        print("\n  DECISION: stdin backend is fast enough. Ship it.")
    elif stats["p50"] < 800:
        print("\n  DECISION: Borderline. Consider Kimina if you have time.")
    else:
        print("\n  DECISION: Too slow. Switch to Kimina backend.")