bench

#!/usr/bin/env python3

import argparse
from datetime import datetime
import datetime as dt
from operator import itemgetter
from pathlib import Path
import subprocess
import sys

parser = argparse.ArgumentParser(description="Run the Tak benchmark")
parser.add_argument(
    "-n", "--input",
    nargs=3, metavar=("X", "Y", "Z"), default=["48", "20", "12"],
    help="The input to the Tak function to use",
)
parser.add_argument(
    "-r", "--runs",
    type=int, default=4,
    help="The number of times to run each implementation",
)
parser.add_argument(
    "-i", "--impls",
    help="The implementations to run",
)
parser.add_argument(
    "-e", "--exclude",
    help="The implementations to skip running",
)
parser.add_argument(
    "-a", "--all",
    action="store_true",
    help="Run all implementations, including skipped-by-default ones",
)
parser.add_argument(
    "-t", "--timeout",
    type=int, default=180,
    help=(
        "The maximum total time to run each implementation for. "
        "Each implementation will be run at least once in full. "
        "Setting this to 0 will disable timeouts."
    )
)
args = parser.parse_args()

skip_default = [
    "arturo",    # Very slow
    "callisto",  # Requires building compiler and installing stdlib
    "lbforth",   # Very slow
    "oak",       # Very slow
    "uxntal",    # Requires building assembler and emulator
]

WARNING = "\x1B[33m"
ERROR = "\x1B[31m"
DIM = "\x1B[90m"
RESET = "\x1B[0m"

if args.impls and args.exclude:
    sys.exit(
        ERROR
        + "Error: --impls and --exclude cannot be used together"
        + RESET
    )

impls = []
exclude = args.exclude.split(",") if args.exclude else []

if args.impls:
    # Explicitly specified
    for impl in args.impls.split(","):
        path = Path.cwd() / impl
        run = path / "run.sh"
        # Error if the user gave a bad impl name
        if not (path.exists() and run.exists()):
            sys.exit(f"{ERROR}Error: Implementation {impl} not found{RESET}")
        impls.append(path)
else:
    for path in Path.cwd().iterdir():
        if not path.is_dir():
            continue
        run = path / "run.sh"
        if not run.exists():
            continue
        # Skip the tricky ones if -a is not passed
        if not args.all and path.name in skip_default:
            continue
        # Skip the ones the user asked to skip - even when -a is used
        if path.name in exclude:
            continue
        impls.append(path)

impls.sort()

results = []

for path in impls:
    name = path.name
    version = path / "version.sh"
    build = path / "build.sh"
    run = path / "run.sh"

    print(f"{name}: ", end="", flush=True)

    version_output = "Unknown version"
    if version.exists():
        result = subprocess.run(
            version,
            cwd=path,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
        )
        if result.returncode == 0:
            version_output = result.stdout.decode("utf-8").strip()
        else:
            print(f"\n{WARNING}Warning: Failed to get {name} version{RESET}")
    else:
        print(f"\n{WARNING}Warning: Failed to get {name} version{RESET}")

    if build.exists():
        print("building", end="", flush=True)
        result = subprocess.run(
            build,
            cwd=path,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )
        if result.returncode != 0:
            print(f"\n{ERROR}Error: Failed to build {name}{RESET}")
            continue
        print(", ", end="")

    print("running", end="", flush=True)

    total = dt.timedelta()
    runs = 0
    ok = True
    for i in range(args.runs):
        start = datetime.now()
        result = subprocess.run(
            [run] + args.input,
            cwd=path,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        if result.returncode != 0:
            print(f"\n{ERROR}Error: Failed to run {name}{RESET}")
            ok = False
            break
        end = datetime.now()
        total += end - start
        runs += 1

        average = total / runs
        if (
            i < args.runs - 1
            and args.timeout > 0
            and (total + average).total_seconds() > args.timeout
        ):
            print(f", timeout after {runs} run(s)", end="", flush=True)
            break
    if not ok:
        continue
    results.append((path.name, version_output, total / runs))

    print()

if not results:
    print("No results")
    sys.exit()

print("===")

name_width = max(len(result[0]) for result in results)
results.sort(key=itemgetter(2))
for name, version, time in results:
    trimmed = str(time).split(":", 1)[1]
    print(f"{name:{name_width}} | {trimmed} {DIM}| {version}{RESET}")