josie / llm-bench

"""Language runners — compile + execute model output against hidden tests.

Each runner takes a task dict (id, lang, symptom, buggy, fix, test) and
model-generated code, writes it to a fresh temp dir, runs the real toolchain,
and returns (passed: bool, detail: str).

Supported languages: Go, Rust, TypeScript (via bun), Python, Swift (if xcrun
available). The coding-suite tasks reference these by their runner key.

Adding a new language: add a `run_<lang>(task, code)` function here and
register it in RUNNERS.
"""
import os
import re
import subprocess

WORK = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "work")
RUN_TIMEOUT = 120  # seconds per compile+run


def _freshdir(task_id, work=WORK):
    d = os.path.join(work, task_id)
    subprocess.run(["rm", "-rf", d])
    os.makedirs(d, exist_ok=True)
    return d


def _sh(cmd, cwd, env=None):
    p = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True,
                       timeout=RUN_TIMEOUT, env=env)
    return p.returncode, (p.stdout + "\n" + p.stderr).strip()


def run_go(task, code):
    d = _freshdir(task["id"])
    code = re.sub(r'^\s*package\s+\w+', 'package solution', code, count=1, flags=re.M)
    if not re.search(r'^\s*package\s+', code, flags=re.M):
        code = "package solution\n\n" + code
    open(os.path.join(d, "solution.go"), "w").write(code)
    open(os.path.join(d, "solution_test.go"), "w").write(task["test"])
    open(os.path.join(d, "go.mod"), "w").write("module bench\n\ngo 1.26\n")
    try:
        rc, out = _sh(["go", "test", "./..."], d)
    except subprocess.TimeoutExpired:
        return False, "TIMEOUT"
    return rc == 0, out[-1500:]


def run_rust(task, code):
    d = _freshdir(task["id"])
    src = code + "\n\n" + task["test"]
    open(os.path.join(d, "solution.rs"), "w").write(src)
    try:
        rc, out = _sh(["rustc", "--edition", "2021", "--test",
                       "solution.rs", "-o", "tb"], d)
        if rc != 0:
            return False, "COMPILE:\n" + out[-1500:]
        rc, out = _sh([os.path.join(d, "tb")], d)
    except subprocess.TimeoutExpired:
        return False, "TIMEOUT"
    return rc == 0, out[-1500:]


def run_ts(task, code):
    d = _freshdir(task["id"])
    open(os.path.join(d, "solution.ts"), "w").write(code)
    open(os.path.join(d, "runtests.ts"), "w").write(task["test"])
    try:
        rc, out = _sh(["bun", "run", "runtests.ts"], d)
    except subprocess.TimeoutExpired:
        return False, "TIMEOUT"
    return rc == 0, out[-1500:]


def run_python(task, code):
    d = _freshdir(task["id"])
    open(os.path.join(d, "solution.py"), "w").write(code)
    open(os.path.join(d, "test_solution.py"), "w").write(task["test"])
    try:
        rc, out = _sh([os.sys.executable, "test_solution.py"], d)
    except subprocess.TimeoutExpired:
        return False, "TIMEOUT"
    return rc == 0, out[-1500:]


_SWIFT_SDK = None

def _swift_env():
    global _SWIFT_SDK
    env = {k: v for k, v in os.environ.items() if k != "SDKROOT"}
    if _SWIFT_SDK is None:
        try:
            p = subprocess.run(["xcrun", "--sdk", "macosx", "--show-sdk-path"],
                               capture_output=True, text=True, env=env)
            _SWIFT_SDK = p.stdout.strip()
        except FileNotFoundError:
            _SWIFT_SDK = ""
    return env, _SWIFT_SDK


def run_swift(task, code):
    d = _freshdir(task["id"])
    open(os.path.join(d, "solution.swift"), "w").write(code)
    open(os.path.join(d, "main.swift"), "w").write(task["test"])
    env, sdk = _swift_env()
    if not sdk:
        return False, "SKIP: xcrun/Swift not available on this platform"
    try:
        rc, out = _sh(["xcrun", "swiftc", "solution.swift", "main.swift",
                       "-sdk", sdk, "-o", "bin"], d, env=env)
        if rc != 0:
            return False, "COMPILE:\n" + out[-1500:]
        rc, out = _sh([os.path.join(d, "bin")], d, env=env)
    except subprocess.TimeoutExpired:
        return False, "TIMEOUT"
    return rc == 0, out[-1500:]


RUNNERS = {
    "go": run_go,
    "rust": run_rust,
    "ts": run_ts,
    "python": run_python,
    "swift": run_swift,
}


def extract_code(content):
    """Return the most plausible code block from a chat content string."""
    if not content:
        return ""
    blocks = re.findall(r"```[a-zA-Z0-9_+#.-]*\n(.*?)```", content, re.DOTALL)
    if blocks:
        return max(blocks, key=len).strip()
    return content.strip()