#!/usr/bin/env python3 -I
"""Part A (d): cargo-mutants-style mutant generator for small pure Rust functions, with a
differential-execution non-equivalence check.

For each function in corpus/functions.json, every (operator, site) pair yields one mutant.
A mutant is KEPT only if (1) it compiles, and (2) a concrete input exists on which the mutant's
observable outcome (return value, or panic) differs from the original's. That input is stored as
the witness, so non-equivalence is demonstrated, not assumed. Mutants with no witness within the
input budget are recorded as `no_witness` (likely equivalent) and excluded. Timeouts (mutant hangs)
are recorded as `timeout` and excluded (no witness input is available from a killed process).

Usage: mutate.py [--out DIR] [--budget N] [--only name,...] [--workdir DIR]
Workdir/out must be under /apps (never /tmp or /home).
"""
import argparse, itertools, json, os, re, subprocess, sys

HERE = os.path.dirname(os.path.abspath(__file__))
TOK = re.compile(r"\s+|//[^\n]*|\d+(?:_\d+)*(?:[ui](?:8|16|32|64|128|size))?|[A-Za-z_]\w*|<<=|>>=|<<|>>|<=|>=|==|!=|&&|\|\||->|=>|\+=|-=|\*=|/=|%=|&=|\|=|\^=|::|\.\.=?|.", re.S)

SWAP = {"+": ["-"], "-": ["+"], "*": ["+"], "/": ["*"], "%": ["/"], "<": ["<=", ">"], "<=": ["<", ">="], ">": [">=", "<"], ">=": [">", "<="],
        "==": ["!="], "!=": ["=="], "&&": ["||"], "||": ["&&"], "&": ["|"], "|": ["&"], "^": ["|"], "<<": [">>"], ">>": ["<<"],
        "+=": ["-="], "-=": ["+="], "*=": ["+="], "/=": ["*="], "<<=": [">>="], ">>=": ["<<="], "&=": ["|="], "|=": ["&="]}
BINARY_PREV_OK = re.compile(r"[\w\)\]]$")


def split_fn(src):
    i = src.index("{")  # first brace: end of the signature (no generics/where in corpus)
    return src[:i], src[i:]


def toks(body):
    return TOK.findall(body)


def sites(body):
    """Yield (token_index, kind, replacement) for every mutation site in the body."""
    t = toks(body)
    sig = [i for i, x in enumerate(t) if not x.isspace() and not x.startswith("//")]
    for n, i in enumerate(sig):
        x = t[i]
        prev = t[sig[n - 1]] if n else ""
        nxt = t[sig[n + 1]] if n + 1 < len(sig) else ""
        if x in SWAP:
            # a binary operator needs an operand token before it; skips unary -, &, and `|closure|`
            if x in ("-", "&", "*", "|") and not (BINARY_PREV_OK.search(prev) and prev not in ("as", "return", "in", "if", "while", "else", "let", "mut")):
                continue
            # `<`/`>` after an identifier followed by a type (generics/turbofish) are excluded by the corpus
            for r in SWAP[x]:
                yield i, "op", r
        elif re.fullmatch(r"\d+(?:_\d+)*", x) and prev not in ("::", "."):
            v = int(x.replace("_", ""))
            for r in sorted({v + 1, max(v - 1, 0), 0 if v else 1} - {v}):
                yield i, "const", str(r)
        elif re.fullmatch(r"\d+(?:_\d+)*[ui](?:8|16|32|64|128|size)", x):
            m = re.match(r"(\d+)(.*)", x); v = int(m.group(1))
            for r in sorted({v + 1, max(v - 1, 0)} - {v}):
                yield i, "const", f"{r}{m.group(2)}"
        elif x in ("true", "false"):
            yield i, "bool", "false" if x == "true" else "true"
        elif x in ("checked_add", "checked_sub", "checked_mul"):
            yield i, "checked", {"checked_add": "checked_sub", "checked_sub": "checked_add", "checked_mul": "checked_add"}[x]


def apply(body, idx, rep):
    t = toks(body)
    t[idx] = rep
    return "".join(t)


def rust_ty_gen(ty, k):
    """Expression (Rust) of an input of type ty drawn from index k of a combined stream."""
    return None


HARNESS = r'''
#![allow(unused, arithmetic_overflow, unconditional_panic)]
mod orig { include!("orig.rs"); }
mod mutant { include!("mutant.rs"); }
use std::panic::{catch_unwind, AssertUnwindSafe};
struct Rng(u64);
impl Rng { fn next(&mut self) -> u64 { self.0 ^= self.0 << 13; self.0 ^= self.0 >> 7; self.0 ^= self.0 << 17; self.0 } }
__GENS__
fn main() {
    std::panic::set_hook(Box::new(|_| {}));
    let mut rng = Rng(0x9E3779B97F4A7C15);
    let mut n: u64 = 0;
    __LOOPS__
}
'''


def gen_fn(ty):
    if ty == "bool":
        return "fn g_bool(r: &mut Rng, i: u64) -> bool { if i < 2 { i == 1 } else { r.next() & 1 == 1 } }"
    bits = {"u8": 8, "u16": 16, "u32": 32, "u64": 64, "usize": 64, "i32": 32, "i64": 64}[ty]
    signed = ty.startswith("i")
    edge = "[0, 1, 2, 3, 7, 8, 9, 10, 15, 16, 100, 127, 128, 255, 256, 1000, 65535, 65536, 0x7fff_ffff, 0x8000_0000, 0xffff_ffff, 0x7fff_ffff_ffff_ffff, 0xffff_ffff_ffff_ffff]"
    return (f"fn g_{ty}(r: &mut Rng, i: u64) -> {ty} {{ const E: [u64; 23] = {edge}; "
            f"if (i as usize) < E.len() * 2 {{ let e = E[(i as usize / 2) % E.len()]; "
            f"if i % 2 == 0 {{ e as {ty} }} else {{ (e as {ty}).wrapping_neg() }} }} else {{ let x = r.next(); "
            f"let sh = (x >> 58) as u32 % {bits}; ((r.next() >> sh) as {ty}) }} }}")


def build_harness(fn):
    tys = sorted(set(fn["args"]))
    gens = "\n".join(gen_fn(t) for t in tys)
    n = len(fn["args"])
    # edge grid over the first two args, random for the rest; then pure random
    names = [f"a{k}" for k in range(n)]
    loop = []
    loop.append("let budget: u64 = BUDGET;")
    loop.append("let mut it: u64 = 0;")
    loop.append("while it < budget {")
    for k, t in enumerate(fn["args"]):
        idx = {0: "it % 64", 1: "(it / 64) % 64"}.get(k, "it")
        loop.append(f"  let a{k}: {t} = if it < 4096 {{ g_{t}(&mut rng, {idx}) }} else {{ g_{t}(&mut rng, 1000 + it) }};")
    call = ", ".join(names)
    loop.append(f"  let o = catch_unwind(AssertUnwindSafe(|| format!(\"{{:?}}\", orig::{fn['name']}({call}))));")
    loop.append(f"  let m = catch_unwind(AssertUnwindSafe(|| format!(\"{{:?}}\", mutant::{fn['name']}({call}))));")
    loop.append("  let os = match &o { Ok(s) => s.clone(), Err(_) => \"PANIC\".to_string() };")
    loop.append("  let ms = match &m { Ok(s) => s.clone(), Err(_) => \"PANIC\".to_string() };")
    fmt = ",".join(f'"{{:?}}"' for _ in names)
    loop.append(f"  if os != ms {{ println!(\"DIFF [{{}}] orig={{}} mutant={{}}\", format!(\"{{:?}}\", ({call},)), os, ms); return; }}")
    loop.append("  it += 1;")
    loop.append("}")
    loop.append('println!("SAME");')
    return HARNESS.replace("__GENS__", gens).replace("__LOOPS__", "\n    ".join(loop))


def run(cmd, cwd, timeout):
    try:
        p = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True, timeout=timeout)
        return p.returncode, p.stdout, p.stderr
    except subprocess.TimeoutExpired:
        return -9, "", "timeout"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="/apps/development/rlhf/bugfind_data/out/mutants.jsonl")
    ap.add_argument("--workdir", default="/apps/development/rlhf/bugfind_data/work/mutants")
    ap.add_argument("--budget", type=int, default=2_000_000)
    ap.add_argument("--only", default="")
    ap.add_argument("--rustc", default="/usr/bin/rustc")  # system rustc 1.88 (the rustup shim has no default toolchain)
    a = ap.parse_args()
    for p in (a.out, a.workdir):
        if not os.path.abspath(p).startswith("/apps/"): sys.exit("paths must be under /apps")
    os.makedirs(a.workdir, exist_ok=True); os.makedirs(os.path.dirname(a.out), exist_ok=True)
    corpus = json.load(open(os.path.join(HERE, "corpus", "functions.json")))
    only = set(filter(None, a.only.split(",")))
    stats = {"generated": 0, "compile_fail": 0, "kept": 0, "no_witness": 0, "timeout": 0}
    per_fn = {}
    out = open(a.out, "w")
    for fn in corpus:
        if only and fn["name"] not in only: continue
        sig, body = split_fn(fn["src"])
        d = os.path.join(a.workdir, fn["name"]); os.makedirs(d, exist_ok=True)
        open(os.path.join(d, "orig.rs"), "w").write(fn["src"] + "\n")
        open(os.path.join(d, "main.rs"), "w").write(build_harness(fn).replace("BUDGET", str(a.budget)))
        ms = list(sites(body))
        per_fn[fn["name"]] = {"generated": len(ms), "kept": 0}
        for k, (idx, kind, rep) in enumerate(ms):
            stats["generated"] += 1
            mut_body = apply(body, idx, rep)
            orig_tok = toks(body)[idx]
            msrc = sig + mut_body
            open(os.path.join(d, "mutant.rs"), "w").write(msrc + "\n")
            rc, so, se = run(["nice", "-n", "15", a.rustc, "--edition", "2021", "-C", "opt-level=1", "-C", "overflow-checks=on", "-C", "debuginfo=0",
                              "-A", "warnings", "main.rs", "-o", "h.bin"], d, 120)
            rec = dict(item_id=f"mut-{fn['name']}-{k:03d}", kind="mutant", function=fn["name"], operator=kind, orig_token=orig_tok, new_token=rep,
                       token_index=idx, original_source=fn["src"], mutant_source=msrc, spec_doc=fn["doc"])
            if rc != 0:
                stats["compile_fail"] += 1; rec["status"] = "compile_fail"; out.write(json.dumps(rec) + "\n"); continue
            rc, so, se = run(["nice", "-n", "15", "./h.bin"], d, 60)
            if rc == -9:
                stats["timeout"] += 1; rec["status"] = "timeout"
            elif so.startswith("DIFF"):
                stats["kept"] += 1; per_fn[fn["name"]]["kept"] += 1
                rec["status"] = "kept_nonequivalent"; rec["witness"] = so.strip()
            else:
                stats["no_witness"] += 1; rec["status"] = "no_witness"
            out.write(json.dumps(rec) + "\n"); out.flush()
    json.dump({"totals": stats, "per_function": per_fn}, open(os.path.join(os.path.dirname(a.out), "mutant_funnel.json"), "w"), indent=1)
    print(json.dumps(stats))

if __name__ == "__main__":
    main()