switch_case_experiment.py

#!/usr/bin/env python3
"""
Generate switch-heavy D sources and benchmark compile-time scaling vs case count.
"""

from __future__ import annotations

import argparse
import csv
import datetime as dt
import math
import os
import platform
import random
import shutil
import statistics
import subprocess
import sys
import time
from pathlib import Path
from typing import Iterable


def parse_case_counts(raw: str) -> list[int]:
    counts: list[int] = []
    for part in raw.split(","):
        part = part.strip()
        if not part:
            continue
        value = int(part)
        if value <= 0:
            raise ValueError(f"Case count must be positive: {value}")
        counts.append(value)
    if not counts:
        raise ValueError("No case counts provided")
    return sorted(set(counts))


def cpu_brand() -> str:
    brand = ""
    if platform.system() == "Darwin":
        try:
            out = subprocess.check_output(
                ["sysctl", "-n", "machdep.cpu.brand_string"],
                stderr=subprocess.DEVNULL,
                text=True,
            )
            brand = out.strip()
        except Exception:
            brand = ""
    if brand:
        return brand
    return platform.processor() or "unknown"


def median_abs_deviation(values: Iterable[float]) -> float:
    vals = list(values)
    if not vals:
        return float("nan")
    med = statistics.median(vals)
    return statistics.median(abs(v - med) for v in vals)


def bootstrap_ci_median(values: list[float], iterations: int = 2000) -> tuple[float, float]:
    if not values:
        return (float("nan"), float("nan"))
    if len(values) == 1:
        return (values[0], values[0])
    rng = random.Random(42)
    n = len(values)
    meds = []
    for _ in range(iterations):
        sample = [values[rng.randrange(n)] for _ in range(n)]
        meds.append(statistics.median(sample))
    meds.sort()
    low_i = int(0.025 * (iterations - 1))
    high_i = int(0.975 * (iterations - 1))
    return (meds[low_i], meds[high_i])


def write_switch_source(path: Path, cases: int) -> None:
    with path.open("w", encoding="utf-8") as f:
        f.write("// Auto-generated by switch_case_experiment.py\n")
        f.write("import std.stdio;\n\n")
        f.write("int dispatch(int x) {\n")
        f.write("    switch (x) {\n")
        for i in range(cases):
            f.write(f"        case {i}: return {i};\n")
        f.write("        default: return -1;\n")
        f.write("    }\n")
        f.write("}\n\n")
        f.write("void main() {\n")
        f.write("    long acc = 0;\n")
        f.write(f"    enum int N = {cases};\n")
        f.write("    foreach (i; 0 .. 2000) {\n")
        f.write("        acc += dispatch(i % N);\n")
        f.write("    }\n")
        f.write("    if (acc == -1) {\n")
        f.write('        writeln("never");\n')
        f.write("    }\n")
        f.write("}\n")


def run_compile(
    compiler: str,
    source: Path,
    obj: Path,
    timeout_sec: float,
) -> tuple[bool, float | None, str, str, int]:
    cmd = [compiler, str(source), "-c", "-O", f"-of={obj}"]
    start = time.perf_counter_ns()
    try:
        cp = subprocess.run(
            cmd,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
            text=True,
            timeout=timeout_sec,
        )
    except subprocess.TimeoutExpired:
        return (False, None, "timeout", "compile command timed out", 124)

    elapsed_ms = (time.perf_counter_ns() - start) / 1_000_000.0
    stderr = (cp.stderr or "").strip()
    if cp.returncode != 0:
        low = stderr.lower()
        if "segmentation fault" in low or "bus error" in low:
            kind = "runtime_crash"
        elif "killed" in low:
            kind = "killed"
        else:
            kind = "compile_error"
        hint = stderr.splitlines()[-1] if stderr else f"compiler exited with {cp.returncode}"
        return (False, elapsed_ms, kind, hint[:240], cp.returncode)
    return (True, elapsed_ms, "", "", 0)


def scaling_exponent(cases: list[int], medians: list[float]) -> float:
    if len(cases) < 2:
        return float("nan")
    xs = [math.log(float(c)) for c in cases]
    ys = [math.log(float(m)) for m in medians]
    xbar = statistics.mean(xs)
    ybar = statistics.mean(ys)
    denom = sum((x - xbar) ** 2 for x in xs)
    if denom == 0:
        return float("nan")
    numer = sum((x - xbar) * (y - ybar) for x, y in zip(xs, ys))
    return numer / denom


def mermaid_block(diagram_lines: list[str]) -> str:
    return "```mermaid\n" + "\n".join(diagram_lines) + "\n```\n"


def write_plot(out_png: Path, summary_rows: list[dict[str, object]]) -> None:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    x = [int(row["cases"]) for row in summary_rows]
    y = [float(row["median_ms"]) for row in summary_rows]
    yerr_low = [float(row["median_ms"]) - float(row["ci_low_ms"]) for row in summary_rows]
    yerr_high = [float(row["ci_high_ms"]) - float(row["median_ms"]) for row in summary_rows]

    fig, ax = plt.subplots(figsize=(9, 5))
    ax.errorbar(
        x,
        y,
        yerr=[yerr_low, yerr_high],
        fmt="o-",
        linewidth=2,
        markersize=7,
        capsize=4,
        color="#1f77b4",
    )
    ax.set_xscale("log")
    ax.set_xlabel("Switch Cases (log scale)")
    ax.set_ylabel("Compile Time Median (ms)")
    ax.set_title("DMD Compile-Time Scaling for Large switch Statements")
    ax.grid(True, alpha=0.3)

    for row in summary_rows:
        ax.annotate(
            f"{int(row['cases'])}",
            (int(row["cases"]), float(row["median_ms"])),
            textcoords="offset points",
            xytext=(6, 6),
            fontsize=9,
        )

    fig.tight_layout()
    fig.savefig(out_png, dpi=150, bbox_inches="tight")
    plt.close(fig)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--compiler",
        default=str(Path(".locald/dmd-nightly/osx/bin/dmd")),
        help="Path to D compiler binary (default: workspace dmd-nightly)",
    )
    parser.add_argument(
        "--case-counts",
        default="100,1000,10000",
        help="Comma-separated list of switch case counts",
    )
    parser.add_argument("--runs", type=int, default=7, help="Measured runs per case count")
    parser.add_argument("--warmups", type=int, default=2, help="Warmup runs per case count")
    parser.add_argument("--timeout", type=float, default=120.0, help="Compile timeout in seconds")
    parser.add_argument(
        "--out-dir",
        default="artifacts/switch_scaling",
        help="Output directory for CSV/plot/report artifacts",
    )
    args = parser.parse_args()

    counts = parse_case_counts(args.case_counts)
    out_dir = Path(args.out_dir)
    gen_dir = out_dir / "generated"
    cache_dir = out_dir / ".cache" / "matplotlib"
    gen_dir.mkdir(parents=True, exist_ok=True)
    cache_dir.mkdir(parents=True, exist_ok=True)
    os.environ.setdefault("MPLCONFIGDIR", str(cache_dir))
    os.environ.setdefault("XDG_CACHE_HOME", str(out_dir / ".cache"))

    compiler = shutil.which(args.compiler) if os.sep not in args.compiler else args.compiler
    if not compiler or not Path(compiler).exists():
        print(f"Compiler not found: {args.compiler}", file=sys.stderr)
        return 2

    raw_csv = out_dir / "results_raw.csv"
    summary_csv = out_dir / "results_summary.csv"
    out_png = out_dir / "compile_time_vs_cases.png"
    out_md = out_dir / "report.md"

    machine = {
        "hostname": platform.node() or "unknown",
        "cpu_brand": cpu_brand(),
        "os": f"{platform.system()} {platform.release()}",
    }

    raw_rows: list[dict[str, object]] = []
    summary_rows: list[dict[str, object]] = []

    print(f"Compiler: {compiler}")
    print(f"Case counts: {counts}")
    print(f"Measured runs: {args.runs} | Warmups: {args.warmups}")

    for cases in counts:
        src = gen_dir / f"switch_{cases}.d"
        obj = gen_dir / f"switch_{cases}.o"
        write_switch_source(src, cases)

        print(f"[cases={cases}] running warmups...")
        for idx in range(1, args.warmups + 1):
            ok, elapsed_ms, err_kind, err_hint, rc = run_compile(
                compiler=compiler,
                source=src,
                obj=obj,
                timeout_sec=args.timeout,
            )
            raw_rows.append(
                {
                    "cases": cases,
                    "run_idx": idx,
                    "is_warmup": 1,
                    "ok": 1 if ok else 0,
                    "time_ms": "" if elapsed_ms is None else f"{elapsed_ms:.3f}",
                    "object_size_bytes": obj.stat().st_size if ok and obj.exists() else "",
                    "error_kind": err_kind,
                    "error_hint": err_hint,
                    "error_code": rc,
                    "timestamp": dt.datetime.now(dt.UTC).isoformat(),
                    "compiler": compiler,
                    **machine,
                }
            )

        print(f"[cases={cases}] running measured compiles...")
        measured: list[float] = []
        measured_obj_size: int | None = None
        for idx in range(1, args.runs + 1):
            ok, elapsed_ms, err_kind, err_hint, rc = run_compile(
                compiler=compiler,
                source=src,
                obj=obj,
                timeout_sec=args.timeout,
            )
            if ok and elapsed_ms is not None:
                measured.append(elapsed_ms)
                if obj.exists():
                    measured_obj_size = obj.stat().st_size
            raw_rows.append(
                {
                    "cases": cases,
                    "run_idx": idx,
                    "is_warmup": 0,
                    "ok": 1 if ok else 0,
                    "time_ms": "" if elapsed_ms is None else f"{elapsed_ms:.3f}",
                    "object_size_bytes": obj.stat().st_size if ok and obj.exists() else "",
                    "error_kind": err_kind,
                    "error_hint": err_hint,
                    "error_code": rc,
                    "timestamp": dt.datetime.now(dt.UTC).isoformat(),
                    "compiler": compiler,
                    **machine,
                }
            )

        n_ok = len(measured)
        n_fail = args.runs - n_ok
        if n_ok > 0:
            med = statistics.median(measured)
            mad = median_abs_deviation(measured)
            mean = statistics.mean(measured)
            ci_low, ci_high = bootstrap_ci_median(measured)
            print(
                f"[cases={cases}] median={med:.3f} ms "
                f"(ok={n_ok}/{args.runs}, object={measured_obj_size or 'n/a'} bytes)"
            )
        else:
            med = mad = mean = ci_low = ci_high = float("nan")
            print(f"[cases={cases}] no successful measured runs")

        summary_rows.append(
            {
                "cases": cases,
                "n_runs": args.runs,
                "n_ok": n_ok,
                "n_fail": n_fail,
                "median_ms": f"{med:.3f}" if not math.isnan(med) else "",
                "mad_ms": f"{mad:.3f}" if not math.isnan(mad) else "",
                "mean_ms": f"{mean:.3f}" if not math.isnan(mean) else "",
                "ci_low_ms": f"{ci_low:.3f}" if not math.isnan(ci_low) else "",
                "ci_high_ms": f"{ci_high:.3f}" if not math.isnan(ci_high) else "",
                "object_size_bytes": measured_obj_size if measured_obj_size is not None else "",
                "source_file": str(src),
            }
        )

    with raw_csv.open("w", newline="", encoding="utf-8") as f:
        fieldnames = [
            "cases",
            "run_idx",
            "is_warmup",
            "ok",
            "time_ms",
            "object_size_bytes",
            "error_kind",
            "error_hint",
            "error_code",
            "timestamp",
            "compiler",
            "hostname",
            "cpu_brand",
            "os",
        ]
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(raw_rows)

    with summary_csv.open("w", newline="", encoding="utf-8") as f:
        fieldnames = [
            "cases",
            "n_runs",
            "n_ok",
            "n_fail",
            "median_ms",
            "mad_ms",
            "mean_ms",
            "ci_low_ms",
            "ci_high_ms",
            "object_size_bytes",
            "source_file",
        ]
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(summary_rows)

    ok_summary = [r for r in summary_rows if r["n_ok"] > 0]
    if len(ok_summary) >= 2:
        write_plot(out_png, ok_summary)
    else:
        out_png.write_text("Plot skipped: fewer than 2 successful points\n", encoding="utf-8")

    times = [float(r["median_ms"]) for r in ok_summary]
    cases_ok = [int(r["cases"]) for r in ok_summary]
    exp = scaling_exponent(cases_ok, times) if len(times) >= 2 else float("nan")
    ratio_lines = []
    for i in range(1, len(ok_summary)):
        prev = ok_summary[i - 1]
        curr = ok_summary[i]
        prev_t = float(prev["median_ms"])
        curr_t = float(curr["median_ms"])
        ratio = curr_t / prev_t if prev_t > 0 else float("inf")
        ratio_lines.append(
            f"- {prev['cases']} -> {curr['cases']}: x{ratio:.3f} compile-time multiplier"
        )

    with out_md.open("w", encoding="utf-8") as f:
        f.write("# Switch-Case Compile-Time Scaling\n\n")
        f.write(f"Generated: {dt.datetime.now(dt.UTC).strftime('%Y-%m-%d %H:%M:%SZ')}\n\n")
        f.write("Each point in this run comes from a separately generated D source file. ")
        f.write("The metric is compile-only wall time, not program runtime.\n\n")
        f.write(f"- Compiler: `{compiler}`\n")
        f.write(f"- Case counts: `{','.join(str(c) for c in counts)}`\n")
        f.write(f"- Runs per point: `{args.runs}` (warmups: `{args.warmups}`)\n")
        f.write("- Mode: compile-only (`-c -O`)\n")
        f.write(f"- Host: `{machine['hostname']}` / `{machine['cpu_brand']}` / `{machine['os']}`\n\n")
        f.write("## Experiment Pipeline\n\n")
        f.write(
            mermaid_block(
                [
                    "flowchart TD",
                    '    A["case counts\\nparameter sweep"] --> B["write_switch_source()\\ngenerated/switch_N.d"]',
                    '    B --> C["warmup compiles\\nshape the cache, do not score"]',
                    '    B --> D["measured compiles\\nwall-clock compile time"]',
                    '    D --> E["results_raw.csv\\nper-run samples"]',
                    '    E --> F["results_summary.csv\\nmedian, MAD, CI"]',
                    '    F --> G["scaling exponent\\nlog-log slope"]',
                    '    F --> H["compile_time_vs_cases.png\\nerror-bar plot"]',
                    '    G --> I["report.md\\ntrend summary"]',
                    '    H --> I',
                ]
            )
        )
        f.write("\n")
        f.write("## Summary\n\n")
        f.write("| Cases | Median ms | MAD ms | 95% CI ms | Object size bytes |\n")
        f.write("|---:|---:|---:|---:|---:|\n")
        for row in summary_rows:
            ci = ""
            if row["ci_low_ms"] and row["ci_high_ms"]:
                ci = f"{row['ci_low_ms']} - {row['ci_high_ms']}"
            f.write(
                f"| {row['cases']} | {row['median_ms'] or '-'} | {row['mad_ms'] or '-'} | "
                f"{ci or '-'} | {row['object_size_bytes'] or '-'} |\n"
            )
        f.write("\n")
        if not math.isnan(exp):
            f.write(f"- Log-log scaling exponent (median compile time vs cases): `{exp:.3f}`\n")
        if ratio_lines:
            f.write("\n".join(ratio_lines))
            f.write("\n")

    print(f"Wrote raw CSV: {raw_csv}")
    print(f"Wrote summary CSV: {summary_csv}")
    print(f"Wrote plot: {out_png}")
    print(f"Wrote report: {out_md}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())