vsmith

#!/usr/bin/env python3
"""Differential fuzzer for VeIR optimization passes.

Repeatedly generates random MLIR programs, interprets them before and after
running an optimizer pipeline, and reports a reproducer if the observable
outputs differ.
"""

from __future__ import annotations

import argparse
import functools
import itertools
import multiprocessing
import os
import random
import re
import secrets
import shutil
import signal
import subprocess
import sys
import tempfile
import time
import traceback
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

sys.path.insert(0, str(Path(__file__).parent))
from mir_llc import emit_asm_cmd, regalloc_cmd
from vsmith_generate import generate


DEFAULT_PASSES = "O"
RISCV_PASSES = "riscv"

# --- RISC-V MIR execution cross-check -------------------------------------
# In --riscv mode we further lower the optimized riscv-dialect MLIR to LLVM
# pre-RA MIR (Tools/veir2mir + llc), assemble/link it into a static Linux ELF
# whose _start calls the lowered `main`, compares its i64 return value against
# veir-interpret's result, and exits with 0 (agree) or MIR_MISMATCH_EXIT
# (disagree). The ELF runs under qemu-riscv64 user mode: directly if it is in
# PATH (Linux), otherwise inside the Lima VM RISCV_VM (macOS). It is piped in
# over stdin, so the VM needn't see the workdir.
MIR_TOOLS = ("llc", "clang")
MIR_MARCH = "rv64im_zba_zbb_zbs_zbc_zbkb_zicond"
MIR_ABI = "lp64"
# Disable C so compressed instructions trap; A, F and D are left on.
QEMU_CPU = "rv64,c=false,zba=true,zbb=true,zbs=true,zbc=true,zbkb=true,zicond=true"
RISCV_VM = "rv64"
QEMU_TIMEOUT = 30
MIR_MISMATCH_EXIT = 99
# Runs the ELF read from stdin; $1 is the -cpu spec, $2 the timeout in seconds.
QEMU_RUNNER = ('f=$(mktemp) && cat >"$f" && chmod +x "$f" && '
               'timeout "$2" qemu-riscv64 -cpu "$1" "$f"; rc=$?; rm -f "$f"; exit $rc')

# _start calls main, compares a0 to the expected value, then exits via the
# exit_group syscall with 0 (match) or MIR_MISMATCH_EXIT (mismatch).  {expected}
# is the interpreter's result, zero-extended to 64 bits, and {normalize}
# zero-extends a narrower a0 to match.
MIR_DRIVER_TEMPLATE = """\
.section .text._start
.globl _start
_start:
    call main
{normalize}    la   t0, expected
    ld   t1, 0(t0)
    li   t2, 0
    beq  a0, t1, .Lexit
    li   t2, {mismatch}
.Lexit:
    mv   a0, t2
    li   a7, 94           # exit_group
    ecall
.section .rodata
.align 3
expected: .dword {expected}
"""


@dataclass(frozen=True)
class ToolResult:
    returncode: int
    output: str


@dataclass(frozen=True)
class CheckResult:
    ok: bool
    before: str
    after: str | None
    message: str
    before_llvm: tuple[tuple[str, str], ...] = ()
    after_llvm: tuple[tuple[str, str], ...] = ()


@dataclass(frozen=True)
class LlvmInterpResult:
    veir_output: str
    final_value: str
    log: str


@dataclass(frozen=True)
class LlvmInterpreter:
    name: str
    bin: Path
    extra_args: tuple[str, ...]
    classifier: Callable[[str], tuple[str, str]]


class Spinner:
    MIN_INTERVAL = 0.2

    def __init__(self) -> None:
        self.frames = ("|", "/", "-", "\\")
        self.index = 0
        self.last_len = 0
        self.interactive = sys.stdout.isatty()
        self.last_update = 0.0

    def update(self, message: str) -> None:
        now = time.monotonic()
        if now - self.last_update < self.MIN_INTERVAL:
            return
        self.last_update = now
        frame = self.frames[self.index % len(self.frames)]
        self.index += 1
        text = f"{frame} {message}"
        if self.interactive:
            padding = " " * max(0, self.last_len - len(text))
            print(f"\r{text}{padding}", end="", flush=True)
            self.last_len = len(text)
        else:
            print(text, flush=True)

    def clear(self) -> None:
        if self.interactive and self.last_len:
            print(f"\r{' ' * self.last_len}\r", end="", flush=True)
            self.last_len = 0

    def finish(self, message: str) -> None:
        self.clear()
        print(message)


def run_tool(repo: Path, cmd: list[str]) -> ToolResult:
    result = subprocess.run(
        cmd,
        cwd=repo,
        stdin=subprocess.DEVNULL,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
    )
    return ToolResult(result.returncode, result.stdout + result.stderr)


KNOWN_INTERPRETER_FAILURES = ("Undefined behavior",)


def classify_failure(log: str) -> str:
    for pattern in KNOWN_INTERPRETER_FAILURES:
        if pattern in log:
            return pattern
    raise ValueError(f"interpreter produced no 'Program output:' line and no known failure pattern:\n{log}")


def outcome_key(output: str | None) -> str:
    if output is None:
        return "optimizer failed"
    if output.startswith("Program output:"):
        values = parse_output_values(output)
        if values is not None and any(v == "poison" for v in values):
            return "program output (poison)"
        return "program output"
    return output


def prebuild(repo: Path, mir_exec: bool = False) -> None:
    """Build veir-opt and veir-interpret so workers can exec them directly."""
    targets = ["veir-opt", "veir-interpret"]
    if mir_exec:
        targets.append("veir2mir")
    result = subprocess.run(
        ["lake", "build", *targets],
        cwd=repo,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )
    if result.returncode != 0:
        raise SystemExit(f"lake build failed:\n{result.stdout}")


def interpret(repo: Path, interp_bin: Path, path: Path) -> tuple[str, str]:
    """Return the last 'Program output: ...' line, or a classified failure string."""
    result = run_tool(repo, [str(interp_bin), str(path)])
    outputs = [line for line in result.output.splitlines() if line.startswith("Program output:")]
    output = outputs[-1] if outputs else classify_failure(result.output)
    return output, result.output


def translate_to_llvm(repo: Path, mlir_path: Path, ll_path: Path) -> tuple[bool, str]:
    """Translate VEIR's LLVM dialect to LLVM IR."""
    translator = repo / "Tools" / "veir2llvm"
    result = run_tool(repo, [str(translator), str(mlir_path)])
    if result.returncode != 0:
        return False, result.output
    if not result.output.strip().startswith("define "):
        return False, result.output
    # Aggregate offsets must use the same layout as veir-interpret. LLVM's
    # default layout has different alignments for i64 and wider integers.
    layout = 'target datalayout = "e-m:e-p:64:64-i64:64-i128:128-n32:64-S128"\n'
    ll_path.write_text(layout + result.output)
    return True, result.output


def format_int_output(width: int, value: int) -> str:
    mask = (1 << width) - 1
    unsigned = value & mask
    digits = max(1, (width + 3) // 4)
    return f"0x{unsigned:0{digits}x}#{width}"


def format_llubi_field(field: str) -> str:
    """Format one field of a struct value printed by llubi, e.g. `i32 8`."""
    if field == "poison":
        return "poison"
    if field in ("T", "F"):
        return "0x1#1" if field == "T" else "0x0#1"
    if re.fullmatch(r"ptr\s+(?:null|0x[0-9A-Fa-f]+(?:\s+\[[^\]]*\])?)", field):
        return "ptr"
    m = re.fullmatch(r"i(\d+)\s+(-?\d+)", field)
    if not m:
        raise ValueError(f"unsupported llubi struct field: {field}")
    return format_int_output(int(m.group(1)), int(m.group(2)))


def classify_llubi_output(log: str) -> tuple[str, str]:
    """Return llubi's result in veir-interpret format and llubi's final value.

    A function with several results returns them as a struct, which llubi
    prints field by field, e.g. `{ i32 8, poison, T }`."""
    for line in log.splitlines():
        if "Immediate UB detected" in line or "Exited with immediate UB" in line:
            return "Undefined behavior", line.strip()
    for line in log.splitlines():
        if "poison return value" in line:
            return "Program output: #[poison]", "poison"

    values: dict[str, str] = {}
    llubi_values: dict[str, str] = {}
    structs: dict[str, str] = {}
    for line in log.splitlines():
        struct_assignment = re.search(r"^\s*(%[-A-Za-z0-9_.$]+)\s*=.*[-=]>\s*\{(.*)\}\s*$", line)
        if struct_assignment:
            structs[struct_assignment.group(1)] = struct_assignment.group(2).strip()
            continue

        assignment = re.search(r"^\s*(%[-A-Za-z0-9_.$]+)\s*=.*[-=]>\s*(poison|i(\d+)\s+(-?\d+))\s*$", line)
        if assignment:
            name = assignment.group(1)
            if assignment.group(2) == "poison":
                values[name] = "poison"
                llubi_values[name] = "poison"
            else:
                values[name] = format_int_output(int(assignment.group(3)), int(assignment.group(4)))
                llubi_values[name] = assignment.group(2)
            continue

        # A pointer prints as its address and object; only that it is one matters here.
        ptr_assignment = re.search(
            r"^\s*(%[-A-Za-z0-9_.$]+)\s*=.*[-=]>\s*(ptr\s+0x[0-9A-Fa-f]+(?:\s+\[[^\]]*\])?)\s*$", line)
        if ptr_assignment:
            values[ptr_assignment.group(1)] = "ptr"
            llubi_values[ptr_assignment.group(1)] = ptr_assignment.group(2)
            continue

        bool_assignment = re.search(r"^\s*(%[-A-Za-z0-9_.$]+)\s*=.*[-=]>\s*([TF])\s*$", line)
        if bool_assignment:
            values[bool_assignment.group(1)] = "0x1#1" if bool_assignment.group(2) == "T" else "0x0#1"
            llubi_values[bool_assignment.group(1)] = bool_assignment.group(2)
            continue

        ret_const = re.search(r"^\s*ret\s+i(\d+)\s+(-?\d+)\s*$", line)
        if ret_const:
            width = int(ret_const.group(1))
            value = int(ret_const.group(2))
            return f"Program output: #[{format_int_output(width, value)}]", f"i{width} {value}"

        if re.fullmatch(r"\s*ret\s+(?:i\d+|ptr)\s+poison\s*", line):
            return "Program output: #[poison]", "poison"

        if re.fullmatch(r"\s*ret\s+ptr\s+null\s*", line):
            return "Program output: #[ptr]", "null"

        ret_bool = re.search(r"^\s*ret\s+i1\s+(true|false)\s*$", line)
        if ret_bool:
            value = ret_bool.group(1)
            return f"Program output: #[{'0x1#1' if value == 'true' else '0x0#1'}]", value

        ret_value = re.search(r"^\s*ret\s+(?:i\d+|ptr)\s+(%[-A-Za-z0-9_.$]+)\s*$", line)
        if ret_value:
            value = values.get(ret_value.group(1))
            llubi_value = llubi_values.get(ret_value.group(1))
            if value is None:
                raise ValueError(f"llubi return value was not recorded:\n{log}")
            if llubi_value is None:
                raise ValueError(f"llubi return value was not recorded:\n{log}")
            return f"Program output: #[{value}]", llubi_value

        ret_struct = re.search(r"^\s*ret\s+\{.*\}\s+(%[-A-Za-z0-9_.$]+)\s*$", line)
        if ret_struct:
            fields = structs.get(ret_struct.group(1))
            if fields is None:
                raise ValueError(f"llubi return value was not recorded:\n{log}")
            formatted = ", ".join(format_llubi_field(f.strip()) for f in fields.split(","))
            return f"Program output: #[{formatted}]", f"{{ {fields} }}"

        if re.fullmatch(r"\s*ret\s+void\s*", line):
            return "Program output: #[]", "void"

    raise ValueError(f"llubi produced no recognizable result:\n{log}")


ALIVE_EXEC_RETURNED_RE = re.compile(
    r"^Returned\s+#([xb])([0-9a-fA-F]+)\s+/\s+"
    r"non-poison=(?:#([xb])([0-9a-fA-F]+)|(true|false))\s*$"
)


def classify_alive_exec_output(log: str, return_type: str | None = None) -> tuple[str, str]:
    """Return alive-exec's result in veir-interpret format and alive-exec's final value.

    `return_type` is the LLVM return type of the executed function. It is
    needed to read a struct result: alive-exec adds the struct's padding as
    fields of its own, which cannot be told from real fields otherwise."""
    if "UB triggered!" in log:
        return "Undefined behavior", "UB triggered"
    if "ERROR: program returned poison" in log:
        return "Program output: #[poison]", "poison"

    final_line: str | None = None
    for line in log.splitlines():
        stripped = line.strip()
        if stripped.startswith("Returned "):
            final_line = stripped

    if final_line is None:
        raise ValueError(f"alive-exec produced no recognizable result:\n{log}")

    m = ALIVE_EXEC_RETURNED_RE.match(final_line)
    if not m:
        raise ValueError(
            f"alive-exec result line not recognized:\n  {final_line}\nfull log:\n{log}"
        )
    radix, digits, mask_radix, mask_digits, non_poison = m.groups()
    unsigned = int(digits, 16 if radix == "x" else 2)

    # A function with several results returns a struct. alive-exec echoes it as
    # one `define {i8, i1, ...}` line, padding included as extra fields, and
    # prints the fields' bits concatenated first-field-first with one
    # non-poison bit per field in the same order.
    alive_types: list[str] | None = None
    for line in log.splitlines():
        agg = re.match(r"define\s+\{([^}]*)\}\s+@", line)
        if agg:
            alive_types = [t.strip() for t in agg.group(1).split(",")]
    if alive_types is not None:
        real = re.fullmatch(r"\{([^}]*)\}", (return_type or "").strip())
        if not real:
            raise ValueError("alive-exec returned a struct, but its declared type is unknown")
        real_types = [t.strip() for t in real.group(1).split(",")]
        if not all(re.fullmatch(r"i\d+", t) for t in alive_types + real_types):
            raise ValueError(f"alive-exec returned a struct with a non-integer field: {return_type}")
        if non_poison is not None:
            mask = ((1 << len(alive_types)) - 1) if non_poison == "true" else 0
        else:
            mask = int(mask_digits, 16 if mask_radix == "x" else 2)
        # A padding field never has the type of the real field that follows
        # it: padding is narrower than that field's alignment, and so than the
        # field itself. So each alive-exec field that does not match the next
        # real field is padding.
        fields: list[str] = []
        offset = sum(int(t[1:]) for t in alive_types)
        index = 0
        for position, typ in enumerate(alive_types):
            width = int(typ[1:])
            offset -= width
            if index < len(real_types) and typ == real_types[index]:
                index += 1
                if (mask >> (len(alive_types) - 1 - position)) & 1:
                    fields.append(format_int_output(width, unsigned >> offset))
                else:
                    fields.append("poison")
        if index != len(real_types):
            raise ValueError(f"alive-exec's struct does not match {return_type}")
        return f"Program output: #[{', '.join(fields)}]", final_line

    poison = (non_poison == "false" if mask_digits is None
              else int(mask_digits, 16 if mask_radix == "x" else 2) == 0)
    if poison:
        return "Program output: #[poison]", "poison"
    # A pointer prints as its encoding; only that it is one matters here.
    if return_type == "ptr":
        return "Program output: #[ptr]", final_line
    width = 4 * len(digits) if radix == "x" else len(digits)
    return f"Program output: #[{format_int_output(width, unsigned)}]", final_line


def run_llvm_interp(repo: Path, interp: LlvmInterpreter, ll_path: Path) -> LlvmInterpResult:
    """Run an LLVM IR interpreter and return its classified result plus raw output."""
    result = run_tool(repo, [str(interp.bin), *interp.extra_args, str(ll_path)])
    veir_output, final_value = interp.classifier(result.output)
    return LlvmInterpResult(veir_output, final_value, result.output)


def check_llvm_interp_equiv(
    repo: Path,
    interp: LlvmInterpreter,
    mlir_path: Path,
    ll_path: Path,
    veir_output: str,
    label: str,
) -> tuple[bool, str, str | None]:
    """Check that an LLVM IR interpreter exactly agrees with veir-interpret."""
    try:
        result = run_llvm_interp(repo, interp, ll_path)
    except ValueError as e:
        return (
            False,
            f"{interp.name} output could not be classified for {label} program\n"
            f"source: {mlir_path}\nllvm: {ll_path}\n{e}",
            None,
        )
    if result.veir_output != veir_output:
        pad = max(len("veir-interpret"), len(interp.name))
        return (
            False,
            f"{interp.name} mismatch on {label} program\n"
            f"source: {mlir_path}\nllvm: {ll_path}\n"
            f"{'veir-interpret':<{pad}}: {veir_output}\n"
            f"{interp.name:<{pad}}: {result.veir_output}\n"
            f"{interp.name} log:\n{result.log}",
            result.final_value,
        )
    return True, "", result.final_value


def check_llvm_interps(
    repo: Path,
    interps: list[LlvmInterpreter],
    mlir_path: Path,
    veir_output: str,
    label: str,
) -> tuple[bool, str, tuple[tuple[str, str], ...]]:
    """Translate to LLVM IR once, then cross-check each interpreter against veir-interpret."""
    ll_path = mlir_path.with_suffix(".ll")
    ok, translate_log = translate_to_llvm(repo, mlir_path, ll_path)
    if not ok:
        return (
            False,
            f"veir2llvm failed on {label} program\n"
            f"source: {mlir_path}\nllvm: {ll_path}\n{translate_log}",
            (),
        )
    collected: list[tuple[str, str]] = []
    for interp in interps:
        ok, msg, value = check_llvm_interp_equiv(repo, interp, mlir_path, ll_path, veir_output, label)
        if value is not None:
            collected.append((interp.name, value))
        if not ok:
            return False, msg, tuple(collected)
    return True, "", tuple(collected)


LLVM_DIALECT_OP_RE = re.compile(r'"(llvm\.[A-Za-z0-9_.]+)"')

# These LLVM dialect ops are expected to remain after the RISC-V pipeline:
# the function wrapper and its terminator are not lowered away.
LLVM_DIALECT_OP_WHITELIST = frozenset({"llvm.func", "llvm.return"})


def find_llvm_dialect_ops(text: str) -> list[str]:
    """Return the sorted, unique LLVM dialect op names in generic-form MLIR."""
    ops = set(LLVM_DIALECT_OP_RE.findall(text))
    return sorted(ops - LLVM_DIALECT_OP_WHITELIST)


def optimize(repo: Path, opt_bin: Path, src: Path, dst: Path, passes: str) -> tuple[bool, str]:
    """Run veir-opt and write the result to dst. Returns (success, log)."""
    result = run_tool(repo, [str(opt_bin), str(src), f"-p={passes}"])
    output = result.output
    failed = (
        result.returncode != 0
        or "Error:" in output
        or "uncaught exception" in output
        or "panic" in output.lower()
        or not output.strip().startswith('"builtin.module"')
    )
    if failed:
        return False, output
    dst.write_text(output)
    return True, output


def parse_output_values(line: str) -> list[str] | None:
    """Parse 'Program output: #[a, b, c]' into a list of element strings."""
    prefix = "Program output:"
    if not line.startswith(prefix):
        return None
    body = line[len(prefix):].strip()
    if not (body.startswith("#[") and body.endswith("]")):
        return None
    inner = body[2:-1].strip()
    return [elem.strip() for elem in inner.split(",")] if inner else []


VALUE_RE = re.compile(r"^0x([0-9a-fA-F]+)#(\d+)$")


def parse_value(value: str) -> tuple[int, int] | None:
    """Parse a single output element like '0xdeadbeef#32' into (value, width)."""
    m = VALUE_RE.match(value)
    if not m:
        return None
    return int(m.group(1), 16), int(m.group(2))


def value_refines(riscv: bool, o: str, p: str) -> bool:
    """Return True if optimized element `p` is an acceptable refinement of original `o`."""
    if o == "poison" or o == p:
        return True
    if not riscv:
        return False
    # In --riscv mode, a function whose declared return type is narrower than a register
    # (e.g. i32) has that boundary coerced to `!riscv.reg` (see `coerceFunction` in
    # Reconciliation.lean), so the optimized value prints at the full register width.
    # The RISC-V calling convention does not guarantee the upper bits are zero- or
    # sign-extended, so they're garbage as far as this check is concerned: only the
    # original width's worth of low bits is meaningful and must be compared.
    orig = parse_value(o)
    opt = parse_value(p)
    if orig is None or opt is None:
        return False
    orig_value, orig_width = orig
    opt_value, opt_width = opt
    if opt_width != 64 or orig_width >= opt_width:
        return False
    return orig_value == (opt_value & ((1 << orig_width) - 1))


def refines(original: str, optimized: str, riscv: bool = False) -> bool:
    """Return True if optimized is a valid refinement of original.

    Undefined behavior in the original may map to any interpreter outcome.
    Poison values in the original may map to poison or concrete values.
    Concrete values in the original must be preserved exactly, except that in --riscv
    mode a narrower-than-register return value may come back truncated from garbage
    upper bits at register width (see `value_refines`).
    """
    if original == "Undefined behavior":
        return optimized == "Undefined behavior" or optimized.startswith("Program output:")

    orig = parse_output_values(original)
    opt = parse_output_values(optimized)
    if orig is None or opt is None or len(orig) != len(opt):
        return original == optimized
    return all(value_refines(riscv, o, p) for o, p in zip(orig, opt))


@functools.lru_cache(maxsize=1)
def resolve_mir_tools() -> dict[str, str] | None:
    """Resolve {llc, clang}; None if any is missing."""
    resolved = {name: shutil.which(name) for name in MIR_TOOLS}
    if any(path is None for path in resolved.values()):
        return None
    return resolved  # type: ignore[return-value]


@functools.lru_cache(maxsize=1)
def resolve_qemu_runner() -> list[str] | None:
    """Command prefix that runs QEMU_RUNNER locally if qemu-riscv64 is in PATH,
    else in the Lima VM RISCV_VM; None if neither qemu-riscv64 nor limactl is."""
    if shutil.which("qemu-riscv64"):
        return ["sh", "-c", QEMU_RUNNER, "sh"]
    limactl = shutil.which("limactl")
    if limactl:
        return [limactl, "shell", RISCV_VM, "sh", "-c", QEMU_RUNNER, "sh"]
    return None


def after_to_expected(after: str) -> tuple[int, int] | None:
    """Parse a single-integer (at most 64 bits) 'Program output:' line into
    its (value, width)."""
    values = parse_output_values(after)
    if values is None or len(values) != 1:
        return None
    parsed = parse_value(values[0])
    if parsed is None or parsed[1] > 64:
        return None
    return parsed


def _mir_run(cmd: list[str], stdin: Path | None = None) -> tuple[int, str]:
    with open(stdin, "rb") if stdin else open(os.devnull, "rb") as f:
        result = subprocess.run(cmd, stdin=f, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    return result.returncode, result.stdout.decode(errors="replace")


def mir_execution_check(repo: Path, opt: Path, after: str, keep: bool) -> tuple[bool, str]:
    """Lower `opt` to MIR, execute it under qemu-riscv64, and check the returned
    integer equals the interpreter's value `after`.  Returns (ok, message).

    For a result narrower than i64 only its low bits of a0 are compared: the
    upper bits are garbage as far as VeIR is concerned (see `value_refines`).
    Other results are skipped (treated as ok). A veir2mir failure or an op it
    can't lower is reported, as are value disagreements and toolchain failures.
    """
    parsed = after_to_expected(after)
    if parsed is None:
        return True, ""  # only single-integer returns are executable here
    expected, width = parsed

    tools = resolve_mir_tools()
    qemu = resolve_qemu_runner()
    assert tools is not None and qemu is not None  # guaranteed by the startup check
    veir2mir = repo / ".lake" / "build" / "bin" / "veir2mir"

    stem = opt.name[: -len(".opt.mlir")] if opt.name.endswith(".opt.mlir") else opt.stem
    d = opt.parent
    mir = d / f"{stem}.mir"
    lowered = d / f"{stem}.lowered.mir"
    asm = d / f"{stem}.s"
    func_o = d / f"{stem}.func.o"
    drv_s = d / f"{stem}.driver.s"
    drv_o = d / f"{stem}.driver.o"
    elf = d / f"{stem}.elf"
    artifacts = [mir, lowered, asm, func_o, drv_s, drv_o, elf]

    def cleanup() -> None:
        if not keep:
            for f in artifacts:
                f.unlink(missing_ok=True)

    march = ["--target=riscv64-unknown-linux-gnu", f"-march={MIR_MARCH}", f"-mabi={MIR_ABI}"]

    # 1. veir2mir: riscv-dialect MLIR -> pre-RA MIR
    rc, out = _mir_run([str(veir2mir), str(opt)])
    mir.write_text(out)
    if rc != 0:
        return False, (f"veir2mir failed (exit {rc})\n"
                       f"optimized: {opt}\nmir: {mir}\n{out}")
    if "UNHANDLED" in out:
        unhandled = [line.strip() for line in out.splitlines() if "UNHANDLED" in line]
        return False, (f"veir2mir could not lower some ops\n"
                       f"optimized: {opt}\nmir: {mir}\n" + "\n".join(unhandled))

    # 2a. llc, register-allocation stage: MIR -> fully-lowered MIR (see mir_llc)
    rc, out = _mir_run(regalloc_cmd(tools["llc"], str(mir), str(lowered)))
    if rc != 0:
        return False, (f"llc failed to lower MIR (regalloc stage)\n"
                       f"optimized: {opt}\nmir: {mir}\n{out}")

    # 2b. llc, emission stage: fully-lowered MIR -> RISC-V assembly (see mir_llc)
    rc, out = _mir_run(emit_asm_cmd(tools["llc"], str(lowered), str(asm)))
    if rc != 0:
        return False, (f"llc failed to emit assembly from lowered MIR\n"
                       f"optimized: {opt}\nlowered: {lowered}\n{out}")

    # 3. assemble the lowered function and the comparison driver, then link
    drv_s.write_text(MIR_DRIVER_TEMPLATE.format(
        expected=expected, mismatch=MIR_MISMATCH_EXIT,
        normalize=(f"    slli a0, a0, {64 - width}\n    srli a0, a0, {64 - width}\n"
                   if width < 64 else "")))
    for label, src_path, obj in (("function", asm, func_o), ("driver", drv_s, drv_o)):
        rc, out = _mir_run([tools["clang"], *march, "-c", str(src_path), "-o", str(obj)])
        if rc != 0:
            return False, (f"assembling MIR {label} failed\n"
                           f"optimized: {opt}\nmir: {mir}\n{out}")
    rc, out = _mir_run([tools["clang"], *march, "-fuse-ld=lld", "-nostdlib", "-static",
                        str(func_o), str(drv_o), "-o", str(elf)])
    if rc != 0:
        return False, (f"linking MIR executable failed\n"
                       f"optimized: {opt}\nmir: {mir}\n{out}")

    # 4. execute under qemu-riscv64: exit 0 = return value agrees,
    # MIR_MISMATCH_EXIT = disagrees, 124 = timeout, anything else = crash
    rc, out = _mir_run([*qemu, QEMU_CPU, str(QEMU_TIMEOUT)], stdin=elf)
    if rc == 0:
        cleanup()
        return True, ""
    if rc == 124:
        return False, (f"MIR execution under qemu-riscv64 timed out ({QEMU_TIMEOUT}s)\n"
                       f"optimized: {opt}\nmir: {mir}\nelf: {elf}")
    if rc == MIR_MISMATCH_EXIT:
        return False, (f"MIR execution disagrees with interpreter\n"
                       f"optimized: {opt}\nmir: {mir}\nelf: {elf}\n"
                       f"interpreter: {after}  (expected 0x{expected:016x})\n"
                       f"qemu-riscv64 returned a different i{width} (exit {rc})\n{out}")
    return False, (f"qemu-riscv64 crashed running MIR (exit {rc})\n"
                   f"optimized: {opt}\nmir: {mir}\nelf: {elf}\n{out}")


def check(
    repo: Path,
    interp_bin: Path,
    opt_bin: Path,
    llvm_interps: list[LlvmInterpreter],
    src: Path,
    opt: Path,
    passes: str,
    riscv: bool,
    keep: bool,
    mir_exec: bool,
) -> CheckResult:
    """Interpret, optimize, re-interpret, and compare."""
    before, _ = interpret(repo, interp_bin, src)
    before_llvm: tuple[tuple[str, str], ...] = ()
    if llvm_interps:
        ok, message, before_llvm = check_llvm_interps(repo, llvm_interps, src, before, "source")
        if not ok:
            return CheckResult(
                ok=False,
                before=before,
                after=None,
                message=message,
                before_llvm=before_llvm,
            )

    ok, opt_log = optimize(repo, opt_bin, src, opt, passes)
    if not ok:
        return CheckResult(
            ok=False,
            before=before,
            after=None,
            message=f"optimizer failed on {src}\n{opt_log}",
            before_llvm=before_llvm,
        )

    if riscv:
        leftover = find_llvm_dialect_ops(opt_log)
        if leftover:
            return CheckResult(
                ok=False,
                before=before,
                after=None,
                message=(
                    f"LLVM dialect ops survived the RISC-V pipeline\n"
                    f"source: {src}\noptimized: {opt}\n"
                    f"ops: {', '.join(leftover)}"
                ),
                before_llvm=before_llvm,
            )

        # A `builtin.unrealized_conversion_cast` can survive the RISC-V pipeline
        # inside a dead (unreachable) block: `reconcile-cast` doesn't yet handle
        # those, so isel-inserted boundary casts linger there. We tolerate this
        # rather than reporting it. It's harmless for differential testing: the
        # interpreter only executes reachable blocks (so a dead-block cast never
        # runs), and any cast that did survive in a *live* block would surface
        # anyway as an interpretation failure or a refinement mismatch below.

    after, _ = interpret(repo, interp_bin, opt)
    after_llvm: tuple[tuple[str, str], ...] = ()
    if llvm_interps:
        ok, message, after_llvm = check_llvm_interps(repo, llvm_interps, opt, after, "optimized")
        if not ok:
            return CheckResult(
                ok=False,
                before=before,
                after=after,
                message=message,
                before_llvm=before_llvm,
                after_llvm=after_llvm,
            )

    if not refines(before, after, riscv):
        return CheckResult(
            ok=False,
            before=before,
            after=after,
            before_llvm=before_llvm,
            after_llvm=after_llvm,
            message=(
                f"semantic mismatch\n"
                f"source: {src}\noptimized: {opt}\n"
                f"original:  {before}\n"
                f"optimized: {after}"
            ),
        )

    if riscv and mir_exec:
        mir_ok, mir_message = mir_execution_check(repo, opt, after, keep)
        if not mir_ok:
            return CheckResult(
                ok=False,
                before=before,
                after=after,
                before_llvm=before_llvm,
                after_llvm=after_llvm,
                message=mir_message,
            )

    return CheckResult(
        ok=True,
        before=before,
        after=after,
        message=before,
        before_llvm=before_llvm,
        after_llvm=after_llvm,
    )


def _worker(
    repo: Path,
    interp_bin: Path,
    opt_bin: Path,
    llvm_interps: list[LlvmInterpreter],
    workdir: Path,
    seed: int,
    count: int,
    start_index: int,
    passes: str,
    riscv: bool,
    keep: bool,
    mir_exec: bool,
    memory_only: bool,
    queue: "multiprocessing.Queue",
) -> None:
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    try:
        workdir.mkdir(parents=True, exist_ok=True)
        rng = random.Random(seed)
        iterator = itertools.count() if count < 0 else range(count)
        for i in iterator:
            idx = start_index + i
            src = workdir / f"case_{idx}.mlir"
            opt = workdir / f"case_{idx}.opt.mlir"
            generate(src, rng, riscv, memory_only)
            result = check(repo, interp_bin, opt_bin, llvm_interps, src, opt, passes, riscv, keep, mir_exec)
            if not result.ok:
                queue.put(("fail", result.message))
                return
            queue.put((
                "ok",
                outcome_key(result.before),
                outcome_key(result.after),
                idx,
                result.before,
                result.after,
                result.before_llvm,
                result.after_llvm,
            ))
            if not keep:
                src.unlink(missing_ok=True)
                opt.unlink(missing_ok=True)
                src.with_suffix(".ll").unlink(missing_ok=True)
                opt.with_suffix(".ll").unlink(missing_ok=True)
        queue.put(("done",))
    except BaseException:
        queue.put(("fail", f"worker crashed:\n{traceback.format_exc()}"))


def fuzz(
    repo: Path,
    interp_bin: Path,
    opt_bin: Path,
    llvm_interps: list[LlvmInterpreter],
    workdir: Path,
    iterations: int,
    passes: str,
    riscv: bool,
    keep: bool,
    mir_exec: bool,
    rng: random.Random,
    num_workers: int,
    spinner: Spinner,
    verbose: bool,
    memory_only: bool = False,
) -> bool:
    infinite = iterations < 0
    if not infinite and iterations == 0:
        spinner.finish("no iterations")
        return True

    if infinite:
        num_workers = max(1, num_workers)
        base, extra = 0, 0
    else:
        num_workers = max(1, min(num_workers, iterations))
        base, extra = divmod(iterations, num_workers)

    ctx = multiprocessing.get_context()
    queue: "multiprocessing.Queue" = ctx.Queue()
    processes: list[multiprocessing.Process] = []
    start = 0
    for w in range(num_workers):
        count = -1 if infinite else base + (1 if w < extra else 0)
        seed = rng.getrandbits(64)
        subdir = workdir / f"w{w}"
        p = ctx.Process(
            target=_worker,
            args=(repo, interp_bin, opt_bin, llvm_interps, subdir, seed, count, start, passes, riscv, keep, mir_exec, memory_only, queue),
        )
        p.start()
        processes.append(p)
        if not infinite:
            start += count

    source_counts: defaultdict[str, int] = defaultdict(int)
    opt_counts: defaultdict[str, int] = defaultdict(int)
    done = 0
    failure_message: str | None = None
    interrupted = False

    try:
        while infinite or done < num_workers:
            try:
                msg = queue.get()
            except KeyboardInterrupt:
                interrupted = True
                break
            tag = msg[0]
            if tag == "done":
                done += 1
            elif tag == "fail":
                failure_message = msg[1]
                break
            elif tag == "ok":
                _, src_key, opt_key, idx, before, after, before_llvm, after_llvm = msg
                source_counts[src_key] += 1
                opt_counts[opt_key] += 1
                if verbose:
                    print(f"case {idx}:")
                    print(f"  source:    {before}")
                    for name, value in before_llvm:
                        print(f"  source {name}: {value}")
                    print(f"  optimized: {after}")
                    for name, value in after_llvm:
                        print(f"  optimized {name}: {value}")
                else:
                    spinner.update(f"passed {sum(source_counts.values())} cases")
    finally:
        if failure_message is not None or interrupted:
            for p in processes:
                if p.is_alive():
                    p.terminate()
        for p in processes:
            p.join()

    if failure_message is not None:
        spinner.clear()
        print(failure_message)
        return False

    def fmt_counts(counts: defaultdict[str, int]) -> str:
        return ", ".join(f"{k}: {v}" for k, v in sorted(counts.items()))

    header = "interrupted, no mismatches detected" if interrupted else "no mismatches detected"
    spinner.finish(
        f"{header}\n"
        f"  source:    {fmt_counts(source_counts)}\n"
        f"  optimized: {fmt_counts(opt_counts)}"
    )
    return True


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--iterations", type=int, default=100,
                        help="number of random programs to test, or -1 to run forever (default: 100)")
    parser.add_argument("--seed", type=int, default=None,
                        help="random seed (default: fresh system entropy)")
    parser.add_argument("--passes", default=DEFAULT_PASSES,
                        help=f"veir-opt pass pipeline (default: {DEFAULT_PASSES})")
    parser.add_argument("--riscv", action="store_true",
                        help="use the RISC-V backend pass pipeline")
    parser.add_argument("--memory-only", action="store_true",
                        help="generate only defined GEP/load/store programs returning a memory checksum")
    parser.add_argument("--keep", action="store_true",
                        help="keep workdir on success (always kept on failure)")
    parser.add_argument("--no-mir-exec", action="store_true",
                        help="in --riscv mode, skip lowering to MIR and executing it under qemu-riscv64")
    parser.add_argument("--workers", type=int, default=os.cpu_count() or 1,
                        help="number of parallel workers (default: all available cores)")
    parser.add_argument("--verbose", action="store_true",
                        help="print per-case source/optimized outputs instead of the spinner")
    parser.add_argument("--llubi", action="store_true",
                        help="cross-check veir-interpret against llubi via Tools/veir2llvm")
    parser.add_argument("--alive-exec", action="store_true",
                        help="cross-check veir-interpret against alive-exec via Tools/veir2llvm")
    return parser.parse_args(argv)


def resolve_llvm_interp(
    name: str, flag: str, extra_args: tuple[str, ...], classifier: Callable[[str], tuple[str, str]]
) -> LlvmInterpreter:
    resolved = shutil.which(name)
    if resolved is None:
        raise SystemExit(f"{flag} was requested, but {name!r} was not found in PATH")
    return LlvmInterpreter(name=name, bin=Path(resolved), extra_args=extra_args, classifier=classifier)


def main(argv: list[str]) -> int:
    args = parse_args(argv)
    if args.riscv:
        args.passes = RISCV_PASSES
    repo = Path(__file__).resolve().parents[1]
    mir_exec = args.riscv and not args.no_mir_exec
    if mir_exec and (resolve_mir_tools() is None or resolve_qemu_runner() is None):
        missing = [name for name in MIR_TOOLS if shutil.which(name) is None]
        if resolve_qemu_runner() is None:
            missing.append(f"qemu-riscv64 (or limactl with a Lima VM named {RISCV_VM})")
        raise SystemExit(
            f"--riscv MIR execution needs {', '.join(missing)} in PATH; "
            f"install them or pass --no-mir-exec"
        )
    llvm_interps: list[LlvmInterpreter] = []
    if args.llubi:
        llvm_interps.append(resolve_llvm_interp("llubi", "--llubi", ("--verbose",), classify_llubi_output))
    if args.alive_exec:
        llvm_interps.append(resolve_llvm_interp("alive-exec", "--alive-exec", ("--func=veir_test_main",), classify_alive_exec_output))
    seed = args.seed if args.seed is not None else secrets.randbits(64)
    rng = random.Random(seed)

    workdir = Path(tempfile.mkdtemp(prefix="vsmith-"))

    print(f"workdir: {workdir}")
    print(f"seed:    {seed}")
    print(f"passes:  {args.passes}")
    print(f"workers: {args.workers}")
    if mir_exec:
        where = "locally" if resolve_qemu_runner()[0] == "sh" else f"in Lima VM {RISCV_VM}"
        print(f"mir-exec:    lower to MIR and execute under qemu-riscv64 {where} ({MIR_MARCH})")
    for interp in llvm_interps:
        print(f"{interp.name + ':':<12}{interp.bin}")

    prebuild(repo, mir_exec)
    interp_bin = repo / ".lake" / "build" / "bin" / "veir-interpret"
    opt_bin = repo / ".lake" / "build" / "bin" / "veir-opt"

    success = fuzz(
        repo, interp_bin, opt_bin, llvm_interps, workdir,
        args.iterations, args.passes, args.riscv, args.keep, mir_exec, rng,
        args.workers, Spinner(), args.verbose, args.memory_only,
    )

    if success and not args.keep:
        shutil.rmtree(workdir)
    else:
        print(f"kept workdir: {workdir}")

    return 0 if success else 1


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))