#!/usr/bin/env python3
"""wasmq_mutants.py: each rule of a wasmQ job's form (quantum-wasmq section 6.5), weakened in
each verifier, and held to the receipts differential: every weakening must be caught.

    python3 -I wasmq_mutants.py python --reference PATH/wai-quantum [--jobs N] [--rule N]
    python3 -I wasmq_mutants.py rust-schemata ENGINE_RS... OUT.json
    python3 -I wasmq_mutants.py rust --reference PATH/wai-quantum --mutants OUT.json [--jobs N]

The weakenings are made from the code, not written by hand: at each place a rule is checked,

  drop    the check never fails;
  bound   a comparison moved by one (`<` for `<=`, `>` for `>=`, and back), or a whole
          number in it one more;
  type    a count read as any number (0.5, -1, a number past 2^53), a digest as any string of
          its length, two values compared as written rather than as JSON values;
  null    a member that is `null` read as absent;
  absent  a member that is absent read as `null`;
  extreme the most of many read as the least, and back (a `max` over a list for a `min`).

Each is assigned the rule whose check it is (the second verifier numbers its checks; the
engine marks each rule's code `// Rule N.` or `// N:`). The second verifier's mutants are
whole files; the engine's are one build (`rust-schemata` writes every mutant into the engine
behind the environment variable WQ_MUTANT, each off unless it is named), run once a mutant.
The engine's files are engine/wasmq.rs, the first version's form (section 6.5), and
engine/wasmq2.rs, the second's (section 11.6: rules 4, 8, 11, 12, 30 and 31 as it reads
them, and 36, its energy), numbered as one list across both.

A mutant is caught when `receipt_differential.py --wasmq-only --first` fails with it: the
corpus, the rules stage, the labelled cases and the one-member changes. A mutant no input
can tell from the text (one that only changes how a verdict already refused is worded, say)
is listed in EQUIVALENT with the reason; any other survivor fails the run. Counts are
reported per rule. Standard library only, Python 3.9+.
"""
import argparse
import ast
import copy
import json
import os
import re
import subprocess
import sys
import tempfile
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor

HERE = os.path.dirname(os.path.abspath(__file__))
SECOND = os.path.join(HERE, "receipt_verify.py")
DIFFERENTIAL = os.path.join(HERE, "receipt_differential.py")

# Mutants that no input can tell from the text, each by its rule, kind and what it does
# ("RULES|KIND|WHAT"), with the reason no input tells it from the text.
EQUIVALENT = {
    # The engine's.
    '6|bound|whole\'s upper bound one more in `!m["name"].is_string() || !is_digest(&m["blake3"]) || !whole(&m["bytes"], 1, MAX_EXACT)`': "a count is at most 2^53 - 1 by its own reading (count), so a bound one past that bounds nothing more",
    '6|bound|whole\'s upper bound one more in `!pr["name"].is_string() || !is_digest(&pr["blake3"]) || !whole(&pr["bytes"], 0, MAX_EXACT)`': "the same, for a prior's bytes",
    '8|bound|whole\'s upper bound one more in `!whole(&run["seed"], 0, MAX_EXACT)`': "the same, for the seed",
    "14|drop|the check never fails: `all.is_none()`": "rule 14's check of a call's largest against its fuel refuses a fuel that is no count already",
    '16|type|same read loosely in `mods.len() != got.len() || mods.iter().zip(got).any(|(a, b)| !same(&a["blake3"], &b["blake3"]) || !same(&a["name"], &b["name"]))`#blake3': "a blake3 that is no string breaks rule 16's own check that the result's is a digest, so comparing it as written changes no rule",
    '18|null|a null member read as absent: fuel["callouts"].get(name) in `fuel["callouts"].get(name)`': "both ways the stopped call is found missing are rule 18's",
    '18|absent|an absent member read as null: fuel["callouts"].get(name) in `fuel["callouts"].get(name)`': "the same",
    "19|type|count read loosely in `count(n).is_none_or(|n| n == 0)`": "rule 19's sum of the counts refuses a count that is no count already",
    '20-35|null|a null member read as absent: r.get("fault") in `if r["status"] == "fault" { r.get("fault").filter(|f| f.is_object()) } else { None }`': "the fault is read only as an object: null and absent are both none",
    '20-35|absent|an absent member read as null: r.get("fault") in `if r["status"] == "fault" { r.get("fault").filter(|f| f.is_object()) } else { None }`': "the same",
    '35|null|a null member read as absent: r.get("fault") in `if r["status"] == "fault" { r.get("fault").filter(|f| f.is_object()) } else { None }`': "the same, where rule 35 reads the fault",
    '35|absent|an absent member read as null: r.get("fault") in `if r["status"] == "fault" { r.get("fault").filter(|f| f.is_object()) } else { None }`': "the same",
    # The engine's, of the second version (engine/wasmq2.rs).
    "36|drop|the check never fails: `!o.contains_key(*name)`": "rule 36's other checks refuse what this one would: an absent run reads as null, which is no figure, and absent callouts read as no object, which rule 36 refuses beside an energy that is one",
    "36|null|a null member read as absent: contains_key in `!o.contains_key(*name)`": "the same: a run or callouts of null is no figure and no object",
    '8|null|a null member read as absent: c.get("run") in `c.get("run").is_some() && method.is_none()`': "a contract whose run is null breaks rule 8's other clauses, which the first version's code holds: the second version's clause on the method adds no rule there",
    '8|absent|an absent member read as null: c.get("run") in `c.get("run").is_some() && method.is_none()`': "the same, for a contract with no run",
    '12|type|count read loosely in `exact && count(&r["qubits"]).is_some_and(|q| q > u64::from(EXACT_MAX_QUBITS))`': "qubits that are no count break rule 12 by the first version's clause already: a loose reading of them adds no rule",
    '31|null|a null member read as absent: req.get("wasmq") in `one::envelope_bytes(req.get("wasmq"), r)`': "a request whose wasmq is null names an envelope of no frames, far below 2^27: the four bytes of null decide nothing",
    '31|absent|an absent member read as null: req.get("wasmq") in `one::envelope_bytes(req.get("wasmq"), r)`': "the same, for a request with no wasmq",
}


def _equivalent(rules, kind, what):
    return "%s|%s|%s" % (",".join(map(str, rules)), kind, what) in EQUIVALENT

# ------------------------------------------------------------------ the second verifier

# The type-check helpers: a value read through one of them is read the same whether it is
# absent or null (each refuses both), so a null-or-absent weakening under one changes nothing.
_CHECKERS = {"_count", "_count_in", "_text_is", "_lower_hex", "_obj", "_arr", "_or_zero", "_exactly", "isinstance", "_n"}
_LOOSE = {"_count": "_mut_count", "_count_in": "_mut_count_in", "_lower_hex": "_mut_lower_hex"}
_HELPERS_SOURCE = '''

def _mut_count(v):
    """A weakening: any number is a count."""
    return isinstance(v, Num) and v.value >= -1e300


def _mut_count_in(v, low, high):
    return _mut_count(v) and low <= int(v.value) <= high


def _mut_lower_hex(v, n):
    return isinstance(v, str) and len(v) == n


def _mut_get_null(v, name):
    """A weakening: a member that is null is read as absent."""
    x = _get(v, name)
    return _ABSENT if x is None else x


def _mut_get_absent(v, name):
    """A weakening: a member that is absent is read as null."""
    x = _get(v, name)
    return None if x is _ABSENT else x


def _mut_in_null(name, v):
    return isinstance(v, dict) and v.get(name) is not None


def _mut_in_absent(name, v):
    return isinstance(v, dict)
'''


def _wasmq_functions(tree):
    """The second verifier's functions that check a wasmQ job's form: wasmq_broken and each
    module-level function it calls that is no shared helper."""
    funcs = {n.name: n for n in tree.body if isinstance(n, ast.FunctionDef)}
    root = funcs["wasmq_broken"]
    called = {n.func.id for n in ast.walk(root) if isinstance(n, ast.Call) and isinstance(n.func, ast.Name)}
    shared = _CHECKERS | {"_get", "rule", "all", "any", "len", "set", "sorted", "sum", "min", "max", "abs", "int", "str", "list", "range", "isinstance", "_numbers", "_text_is"}
    out = [root]
    for name in sorted(called - shared):
        f = funcs.get(name)
        if f is not None and f.lineno > funcs.get("is_wasmq", root).lineno:
            out.append(f)
    return out


def _rule_calls(fn):
    """(line, number, the call) for each rule(N, ...) in fn, in order."""
    out = []
    for n in ast.walk(fn):
        if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id == "rule" and n.args and isinstance(n.args[0], ast.Constant):
            out.append((n.lineno, n.args[0].value, n))
    return sorted(out, key=lambda x: x[0])


def _rule_of(node, root_calls, helper_rules, fn):
    """The rule a node's check belongs to: the rule(N, ...) call around it; else, in
    wasmq_broken, for a value assigned, every rule whose own code (from the rule before it to
    its call) reads that name, or else the first rule called at or after its line; in a
    helper, the rules whose checks call it."""
    for line, number, call in root_calls:
        if any(x is node for x in ast.walk(call)):
            return [number]
    if fn.name != "wasmq_broken":
        return helper_rules.get(fn.name, ["?"])
    top = next((st for st in fn.body if any(x is node for x in ast.walk(st))), None)
    if isinstance(top, ast.Assign):
        names = {t.id for target in top.targets for t in ast.walk(target) if isinstance(t, ast.Name)}
        readers = []
        bounds = [0] + [line for line, _, _ in root_calls]
        for (line, number, _), start in zip(root_calls, bounds):
            block = [st for st in fn.body if start < st.lineno <= line and st is not top]
            if any(isinstance(x, ast.Name) and x.id in names for st in block for x in ast.walk(st)):
                readers.append(number)
        if readers:
            return sorted(set(readers))
    after = [number for line, number, _ in root_calls if line >= node.lineno]
    return after[:1] or ["?"]


def python_mutants(path=SECOND):
    """Every weakening of the second verifier's form check: (id, rules, what, the file's text)."""
    with open(path, encoding="utf-8") as f:
        source = f.read()
    tree = ast.parse(source)
    fns = _wasmq_functions(tree)
    root = fns[0]
    root_calls = _rule_calls(root)
    helper_rules = defaultdict(list)
    for line, number, call in root_calls:
        for n in ast.walk(call):
            if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and any(n.func.id == f.name for f in fns[1:]):
                helper_rules[n.func.id].append(number)
    # A helper called from code before a rule(N, ...) is that rule's.
    for f in fns[1:]:
        if f.name not in helper_rules:
            for n in ast.walk(root):
                if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id == f.name:
                    helper_rules[f.name] = _rule_of(n, root_calls, {}, root)

    parents = {}
    for fn in fns:
        for p in ast.walk(fn):
            for c in ast.iter_child_nodes(p):
                parents[c] = p
    sites = []  # (node, kind, what, transform)

    flip = {ast.Lt: ast.LtE, ast.LtE: ast.Lt, ast.Gt: ast.GtE, ast.GtE: ast.Gt}
    for fn in fns:
        for n in ast.walk(fn):
            if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id == "rule" and n.args and isinstance(n.args[0], ast.Constant):
                sites.append((n, "drop", "rule %s never fails" % n.args[0].value, lambda m: ast.Call(func=m.func, args=[m.args[0], ast.Constant(True)], keywords=[])))
            if isinstance(n, ast.Compare):
                for i, op in enumerate(n.ops):
                    if type(op) in flip:
                        def t(m, i=i):
                            m.ops[i] = flip[type(m.ops[i])]()
                            return m
                        sites.append((n, "bound", "%s for %s in %s" % (flip[type(op)].__name__, type(op).__name__, ast.unparse(n)), t))
                for c in ast.walk(n):
                    if isinstance(c, ast.Constant) and type(c.value) is int:
                        sites.append((c, "bound", "%d for %d in %s" % (c.value + 1, c.value, ast.unparse(n)), lambda m: ast.Constant(m.value + 1)))
            if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id in ("max", "min") and len(n.args) == 1:
                other = "min" if n.func.id == "max" else "max"
                sites.append((n, "extreme", "%s for %s in %s" % (other, n.func.id, ast.unparse(n)), lambda m, other=other: ast.Call(func=ast.Name(other, ast.Load()), args=m.args, keywords=m.keywords)))
            if isinstance(n, ast.Call) and isinstance(n.func, ast.Name):
                if n.func.id == "_count_in":
                    for c in n.args[1:]:
                        if isinstance(c, ast.Constant) and type(c.value) is int:
                            sites.append((c, "bound", "%d for %d in %s" % (c.value + 1, c.value, ast.unparse(n)), lambda m: ast.Constant(m.value + 1)))
                if n.func.id in _LOOSE:
                    sites.append((n, "type", "%s read loosely" % ast.unparse(n), lambda m: ast.Call(func=ast.Name(_LOOSE[m.func.id], ast.Load()), args=m.args, keywords=m.keywords)))
                if n.func.id == "_get":
                    p = parents.get(n)
                    checked = isinstance(p, ast.Call) and isinstance(p.func, ast.Name) and p.func.id in _CHECKERS
                    if not checked:
                        sites.append((n, "null", "a null member read as absent: %s" % ast.unparse(n), lambda m: ast.Call(func=ast.Name("_mut_get_null", ast.Load()), args=m.args, keywords=[])))
                        sites.append((n, "absent", "an absent member read as null: %s" % ast.unparse(n), lambda m: ast.Call(func=ast.Name("_mut_get_absent", ast.Load()), args=m.args, keywords=[])))
            if isinstance(n, ast.Compare) and len(n.ops) == 1 and isinstance(n.ops[0], (ast.In, ast.NotIn)) and isinstance(n.left, ast.Constant) and isinstance(n.left.value, str):
                for kind, helper in (("null", "_mut_in_null"), ("absent", "_mut_in_absent")):
                    def t(m, helper=helper):
                        call = ast.Call(func=ast.Name(helper, ast.Load()), args=[m.left, m.comparators[0]], keywords=[])
                        return ast.UnaryOp(ast.Not(), call) if isinstance(m.ops[0], ast.NotIn) else call
                    sites.append((n, kind, "`%s` in: %s" % (n.left.value, "null read as absent" if kind == "null" else "absent read as null"), t))

    out = []
    for k, (node, kind, what, transform) in enumerate(sites):
        fn = next(f for f in fns if any(x is node for x in ast.walk(f)))
        rules = _rule_of(node, root_calls, helper_rules, fn)
        mutant_tree = copy.deepcopy(tree)
        # The same node in the copy: by its place in a walk of the same function.
        index = [i for i, x in enumerate(ast.walk(fn)) if x is node][0]
        mfn = next(f for f in mutant_tree.body if isinstance(f, ast.FunctionDef) and f.name == fn.name)
        target = list(ast.walk(mfn))[index]
        replaced = transform(copy.deepcopy(target))
        _replace(mfn, target, replaced)
        ast.fix_missing_locations(mutant_tree)
        text = ast.unparse(mutant_tree) + _HELPERS_SOURCE
        ident = "P%04d %s:%d %s" % (k, fn.name, node.lineno, kind)
        out.append((ident, rules, what, text))
    return out


def _replace(tree, old, new):
    for p in ast.walk(tree):
        for field, value in ast.iter_fields(p):
            if value is old:
                setattr(p, field, new)
                return
            if isinstance(value, list):
                for i, x in enumerate(value):
                    if x is old:
                        value[i] = new
                        return
    raise ValueError("no place to replace")


# ------------------------------------------------------------------ the engine

_FUNCS = ("fn contract_form(", "pub fn problems(", "fn across(", "fn members(", "pub fn envelope_bytes(", "pub fn bare_transcript(",
          # engine/wasmq2.rs: the energy's figures and members (rule 36).
          "fn exactly(", "fn figure(", "fn energy(")
_MARK = re.compile(r"^\s*// (?:Rule )?(\d+)[.:]")
_CMP = re.compile(r"(?<=[\w\)\]]) (>=|<=|>|<) (?=[\w\(&!*-])")
_FLIP = {">=": ">", ">": ">=", "<=": "<", "<": "<="}
_HELPERS_RS = '''

/// The weakening this build is, when wasmq_mutants.py made it: none otherwise.
fn mutant(k: u32) -> bool {
    static WHICH: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
    *WHICH.get_or_init(|| std::env::var("WQ_MUTANT").ok().and_then(|v| v.parse().ok()).unwrap_or(u32::MAX)) == k
}

/// A weakening: any number is a count.
fn loose_count(v: &Value) -> Option<u64> {
    v.as_f64().map(|x| x.abs() as u64)
}

fn loose_whole(v: &Value, lo: u64, hi: u64) -> bool {
    loose_count(v).is_some_and(|n| lo <= n && n <= hi)
}

fn loose_digest(v: &Value) -> bool {
    v.as_str().is_some_and(|s| s.len() == 64)
}

fn raw_same(a: &Value, b: &Value) -> bool {
    a == b
}
'''


def _calls(expr, name):
    """Each place `name(` is called in expr: (start, end) of the call, its parentheses matched."""
    out = []
    for m in re.finditer(r"(?<![\w.])" + re.escape(name) + r"\(", expr):
        depth, i = 0, m.end() - 1
        while i < len(expr):
            if expr[i] == "(":
                depth += 1
            elif expr[i] == ")":
                depth -= 1
                if depth == 0:
                    out.append((m.start(), i + 1))
                    break
            i += 1
    return out


def _args(call):
    """A call's arguments, split at its top-level commas."""
    inner = call[call.index("(") + 1:-1]
    out, depth, start = [], 0, 0
    for i, ch in enumerate(inner):
        if ch in "([{":
            depth += 1
        elif ch in ")]}":
            depth -= 1
        elif ch == "," and depth == 0:
            out.append(inner[start:i].strip())
            start = i + 1
    out.append(inner[start:].strip())
    return out


def _expr_mutants(expr):
    """(kind, what, the expression weakened) for each weakening of one expression, `what`
    naming the expression it weakens."""
    return [(kind, "%s in `%s`" % (what, expr), mut) for kind, what, mut in _expr_mutants_bare(expr)]


def _expr_mutants_bare(expr):
    out = []
    for m in _CMP.finditer(expr):
        out.append(("bound", "%s for %s" % (_FLIP[m.group(1)], m.group(1)), expr[:m.start(1)] + _FLIP[m.group(1)] + expr[m.end(1):]))
    for m in re.finditer(r"(?<![\w.<])Some\((\d+)\)", expr):
        out.append(("bound", "Some(%d) for Some(%s)" % (int(m.group(1)) + 1, m.group(1)), expr[:m.start(1)] + str(int(m.group(1)) + 1) + expr[m.end(1):]))
    for a, b in _calls(expr, "whole"):
        args = _args(expr[a:b])
        if len(args) == 3:
            out.append(("bound", "whole's upper bound one more", expr[:a] + "whole(%s, %s, (%s).saturating_add(1))" % (args[0], args[1], args[2]) + expr[b:]))
            if args[1].isdigit() and int(args[1]) > 0:
                out.append(("bound", "whole's lower bound one less", expr[:a] + "whole(%s, %d, %s)" % (args[0], int(args[1]) - 1, args[2]) + expr[b:]))
            out.append(("type", "whole read loosely", expr[:a] + "loose_whole(" + expr[a + len("whole("):]))
    for name, loose in (("count", "loose_count"), ("is_digest", "loose_digest"), ("same", "raw_same")):
        for a, b in _calls(expr, name):
            out.append(("type", "%s read loosely" % name, expr[:a] + loose + expr[a + len(name):]))
    for m in re.finditer(r"(?<![\w.])(\w+(?:\[[^\[\]]*\])*)\.get\(([^()]*)\)", expr):
        whole_get = m.group(0)
        recv, key = whole_get[: whole_get.rindex(".get(")], m.group(2)
        out.append(("null", "a null member read as absent: %s" % whole_get, expr[:m.start()] + "%s.get(%s).filter(|v| !v.is_null())" % (recv, key) + expr[m.end():]))
        out.append(("absent", "an absent member read as null: %s" % whole_get, expr[:m.start()] + "Some(%s.get(%s).unwrap_or(&Value::Null))" % (recv, key) + expr[m.end():]))
    for m in re.finditer(r"\.(max|min)\(\)", expr):
        other = "min" if m.group(1) == "max" else "max"
        out.append(("extreme", "%s() for %s()" % (other, m.group(1)), expr[:m.start(1)] + other + expr[m.end(1):]))
    for m in re.finditer(r"(\w+)\.contains_key\(([^()]*)\)", expr):
        out.append(("null", "a null member read as absent: contains_key", expr[:m.start()] + "%s.get(%s).is_some_and(|v| !v.is_null())" % (m.group(1), m.group(2)) + expr[m.end():]))
    return out


def rust_schemata(engine_rss, out_json):
    """Writes every weakening of the engine's form check into each of engine_rss, each behind
    mutant(K), K numbered across the files, and the list of them to out_json."""
    mutants = []
    for engine_rs in engine_rss:
        _rust_schemata_one(engine_rs, mutants)
    with open(out_json, "w", encoding="utf-8") as f:
        json.dump(mutants, f, indent=1)
    print("rust-schemata: %d weakenings written into %s" % (len(mutants), ", ".join(engine_rss)))


def _rust_schemata_one(engine_rs, mutants):
    with open(engine_rs, encoding="utf-8") as f:
        lines = f.read().split("\n")
    first = len(mutants)
    rule, inside, depth = "?", False, 0
    new = []
    for no, line in enumerate(lines, 1):
        if not inside and any(line.startswith(f) or line.startswith("    " + f) or line.startswith("pub(crate) " + f) for f in _FUNCS):
            # What across reads before rule 20's mark is the fault, which rules 20 to 35 share.
            inside, depth, rule = True, 0, ("members" if "fn members(" in line else "20-35" if "fn across(" in line else "?")
            new.append(line)
            depth += line.count("{") - line.count("}")
            continue
        if inside:
            m = _MARK.match(line)
            if m:
                rule = m.group(1)
            if line.strip().startswith("fn ") and "fn priors(" in line:
                new.append(line)
                depth += line.count("{") - line.count("}")
                continue
            done = _mutate_line(line, lines[no] if no < len(lines) else "", rule, no, mutants)
            new.append(done)
            depth += line.count("{") - line.count("}")
            if depth <= 0:
                inside = False
            continue
        new.append(line)
    text = "\n".join(new)
    test_at = text.find("\n#[cfg(test)]")
    text = text[:test_at] + _HELPERS_RS + text[test_at:] if test_at >= 0 else text + _HELPERS_RS
    with open(engine_rs, "w", encoding="utf-8") as f:
        f.write(text)
    for m in mutants[first:]:
        m["file"] = os.path.basename(engine_rs)
    print("rust-schemata: %d weakenings in %s" % (len(mutants) - first, engine_rs))


def _mutate_line(line, following, rule, no, mutants):
    """The line with each weakening of its check behind mutant(K)."""
    def chain(expr, extra=()):
        variants = list(extra) + _expr_mutants(expr)
        if not variants:
            return None
        arms = []
        seen = defaultdict(int)
        for kind, what, mut in variants:
            k = len(mutants)
            seen[what] += 1
            if seen[what] > 1 or sum(1 for _, w, _ in variants if w == what) > 1:
                # The same weakening at two places of one expression: named by what it weakens.
                at = re.search(r'&\w\["(\w+)"\], &\w\["\1"\]', mut[mut.find("raw_same(") if "raw_same(" in mut else 0:])
                what = "%s#%s" % (what, at.group(1) if at else seen[what])
            mutants.append({"id": k, "rule": rule, "kind": kind, "what": what, "line": no})
            arms.append("if mutant(%d) { %s }" % (k, mut))
            if kind == "drop":
                mutants[-1]["what"] = "the check never fails: `%s`" % expr
        return "(" + " else ".join(arms) + " else { %s })" % expr
    pushes = following.strip().startswith(("out.push(", "problem("))
    m = re.match(r"^(\s*)if let (.+) = (.+) \{$", line)
    if m and not pushes and re.search(r"\.(max|min)\(\)", m.group(3)):
        # What a check after it reads, found as the most or the least of many.
        c = chain(m.group(3))
        return "%sif let %s = %s {" % (m.group(1), m.group(2), c) if c else line
    if m and pushes:
        # A check of what is found, as rule 1's: never found.
        k = len(mutants)
        mutants.append({"id": k, "rule": rule, "kind": "drop", "what": "the check never fails: `%s`" % m.group(3), "line": no})
        return "%sif let %s = (%s).filter(|_| !mutant(%d)) {" % (m.group(1), m.group(2), m.group(3), k)
    m = re.match(r"^(\s*)&& (.+)$", line)
    if m and following.strip() == "{":
        c = chain(m.group(2), [("drop", "the check never fails", "false")])
        return "%s&& %s" % (m.group(1), c) if c else line
    m = re.match(r"^(\s*)if (?!let )(.+) \{$", line)
    if m and pushes and " let " not in m.group(2):
        c = chain(m.group(2), [("drop", "the check never fails", "false")])
        return "%sif %s {" % (m.group(1), c) if c else line
    m = re.match(r"^(\s*)((?:Some\(.*?\)|None)) if (.+?) => (.*)$", line)
    if m:
        c = chain(m.group(3), [("drop", "the check never fails", "false")])
        return "%s%s if %s => %s" % (m.group(1), m.group(2), c, m.group(4)) if c else line
    m = re.match(r"^(\s*)match (.+) \{$", line)
    if m and ".get(" in m.group(2):
        c = chain(m.group(2))
        return "%smatch %s {" % (m.group(1), c) if c else line
    m = re.match(r"^(\s*)let (mut )?([\w(), ]+?)(: [^=]+)? = (.+);$", line)
    if m and not m.group(5).lstrip().startswith(("|", "move |")) and any(x in m.group(5) for x in ("count(", "whole(", ".get(", " > ", " < ", " >= ", " <= ", "Some(")):
        c = chain(m.group(5))
        return "%slet %s%s%s = %s;" % (m.group(1), m.group(2) or "", m.group(3), m.group(4) or "", c) if c else line
    return line


# ------------------------------------------------------------------ running them

def caught(env_extra, reference, second=None):
    cmd = [sys.executable, "-I", DIFFERENTIAL, "--reference", reference, "--wasmq-only", "--first"]
    if second:
        cmd += ["--second", second]
    env = dict(os.environ)
    env.update(env_extra)
    p = subprocess.run(cmd, capture_output=True, text=True, env=env)
    first = next((l for l in p.stdout.splitlines() if l.startswith("DISAGREE")), "")
    if p.returncode not in (0, 1):
        return None, (p.stderr.strip().splitlines() or ["status %d" % p.returncode])[-1]
    return p.returncode == 1, first


def report(results):
    """results: (id, rules, kind, what, caught, why). Prints counts per rule; returns the
    survivors not listed as equivalent."""
    per = defaultdict(lambda: defaultdict(int))
    survivors = []
    for ident, rules, kind, what, got, why in results:
        for r in rules:
            per[str(r)]["generated"] += 1
            per[str(r)]["caught" if got else "equivalent" if _equivalent(rules, kind, what) else "survived"] += 1
        if not got and not _equivalent(rules, kind, what):
            survivors.append((ident, rules, kind, what, why))
    def key(r):
        return (0, int(r)) if r.isdigit() else (1, r)
    print("%-10s %9s %7s %10s %8s" % ("rule", "generated", "caught", "equivalent", "survived"))
    for r in sorted(per, key=key):
        c = per[r]
        print("%-10s %9d %7d %10d %8d" % (r, c["generated"], c["caught"], c["equivalent"], c["survived"]))
    for s in survivors:
        print("SURVIVED %s (rules %s) %s: %s %s" % (s[0], ",".join(map(str, s[1])), s[2], s[3], s[4]))
    return survivors


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("what", choices=["python", "rust", "rust-schemata"])
    ap.add_argument("paths", nargs="*")
    ap.add_argument("--reference")
    ap.add_argument("--mutants")
    ap.add_argument("--jobs", type=int, default=4)
    ap.add_argument("--rule", help="only the weakenings of this rule")
    ap.add_argument("--list", action="store_true", help="list the weakenings, run none")
    a = ap.parse_args()
    if a.what == "rust-schemata":
        if len(a.paths) < 2:
            ap.error("rust-schemata ENGINE_RS... OUT.json")
        return rust_schemata(a.paths[:-1], a.paths[-1])
    if a.what == "python" and a.list:
        ms = [m for m in python_mutants() if not a.rule or a.rule in map(str, m[1])]
        counts = defaultdict(lambda: defaultdict(int))
        for ident, rules, what, _ in ms:
            print(ident, rules, what)
            for r in rules:
                counts[str(r)][ident.split()[-1]] += 1
        for r in sorted(counts, key=lambda r: (0, int(r)) if r.isdigit() else (1, r)):
            print("rule %s: %s" % (r, dict(counts[r])))
        print("%d weakenings" % len(ms))
        return 0
    if not a.reference:
        ap.error("--reference is required")
    cache = tempfile.mkdtemp(prefix="wasmq-mutants-cache-")
    base = {"RD_CACHE": cache}
    # The unweakened pair agrees: else nothing below means anything.
    ok, why = caught(base, a.reference)
    if ok is not False:
        print("the unweakened verifiers do not agree, or the differential failed: %s" % why)
        return 2
    results = []
    if a.what == "python":
        ms = [m for m in python_mutants() if not a.rule or a.rule in map(str, m[1])]
        if a.list:
            for ident, rules, what, _ in ms:
                print(ident, rules, what)
            return 0
        tmp = tempfile.mkdtemp(prefix="wasmq-mutants-py-")

        def one(m):
            ident, rules, what, text = m
            path = os.path.join(tmp, "m%s.py" % ident.split()[0])
            with open(path, "w", encoding="utf-8") as f:
                f.write(text)
            got, why = caught(base, a.reference, second=path)
            os.remove(path)
            return (ident, rules, ident.split()[-1], what, bool(got), why)
        with ThreadPoolExecutor(a.jobs) as pool:
            for r in pool.map(one, ms):
                results.append(r)
                print("%s %-8s %s %s" % ("caught  " if r[4] else "SURVIVED", r[2], r[0], r[3]), flush=True)
    else:
        with open(a.mutants, encoding="utf-8") as f:
            ms = [m for m in json.load(f) if not a.rule or m["rule"] == a.rule]

        def one(m):
            env = dict(base, WQ_MUTANT=str(m["id"]))
            got, why = caught(env, a.reference)
            ident = "R%04d %s:%d %s" % (m["id"], m.get("file", "wasmq.rs"), m["line"], m["kind"])
            return (ident, [m["rule"]], m["kind"], m["what"], bool(got), why)
        with ThreadPoolExecutor(a.jobs) as pool:
            for r in pool.map(one, ms):
                results.append(r)
                print("%s %-8s %s %s" % ("caught  " if r[4] else "SURVIVED", r[2], r[0], r[3]), flush=True)
    survivors = report(results)
    print("%s: %d weakenings, %d caught, %d survived" % (a.what, len(results), sum(1 for r in results if r[4]), len(survivors)))
    return 1 if survivors else 0


if __name__ == "__main__":
    sys.exit(main())