#!/usr/bin/env python3
"""Plant one instruction a WAISM rule refuses in a module, for the checkers to refuse.

    python3 -I plant.py float|relaxed IN OUT

Writes OUT: IN with instructions added at the start of the body of its `wai_abi`, the rest
of the module unchanged but for the two sizes that body changes (the body's and the code
section's). Prints the verdicts both checkers must then give, a line a level:

    float     `f32.const 0`, `drop`: refused R3 at the f32.const at the integer level; held at
              the float level (waism §9.1's scalar-float, in the module of record)
    relaxed   two `v128.const`s, `i8x16.relaxed_swizzle`, `drop`: refused R4 at the relaxed
              instruction at both levels (§9.1's relaxed-swizzle)

The module must be one both checkers hold at the integer level, with no imported function
(R5): its function index is then its body's index in the code section.
"""

import sys

PLANTS = {
    # f32.const 0.0, drop
    "float": (bytes([0x43, 0, 0, 0, 0, 0x1A]), 0, {"integer": "R3", "float": None}),
    # v128.const 0, v128.const 0, i8x16.relaxed_swizzle (0xfd 0x100), drop
    "relaxed": (bytes([0xFD, 0x0C]) + bytes(16) + bytes([0xFD, 0x0C]) + bytes(16) + bytes([0xFD, 0x80, 0x02, 0x1A]), 36, {"integer": "R4", "float": "R4"}),
}


def uleb(b, at):
    v, shift = 0, 0
    while True:
        c = b[at]
        at += 1
        v |= (c & 0x7F) << shift
        shift += 7
        if not c & 0x80:
            return v, at


def enc(v):
    out = bytearray()
    while True:
        c = v & 0x7F
        v >>= 7
        if v:
            out.append(c | 0x80)
        else:
            out.append(c)
            return bytes(out)


def sections(b):
    if b[:8] != b"\0asm\x01\0\0\0":
        raise SystemExit("plant.py: not a WebAssembly module")
    at, out = 8, []
    while at < len(b):
        sid = b[at]
        size, body = uleb(b, at + 1)
        out.append((sid, at, body, body + size))
        at = body + size
    return out


def name(b, at):
    n, at = uleb(b, at)
    return b[at:at + n].decode("utf-8"), at + n


def main(argv):
    if len(argv) != 3 or argv[0] not in PLANTS:
        print("usage: python3 -I plant.py float|relaxed IN OUT", file=sys.stderr)
        return 2
    code, rel, verdicts = PLANTS[argv[0]]
    with open(argv[1], "rb") as f:
        b = f.read()
    secs = sections(b)
    for sid, _, body, _ in secs:
        if sid == 2:
            raise SystemExit("plant.py: the module imports; a held module imports nothing (R5)")
    exports = next((s for s in secs if s[0] == 7), None)
    codesec = next((s for s in secs if s[0] == 10), None)
    if not exports or not codesec:
        raise SystemExit("plant.py: no export or code section")
    n, at = uleb(b, exports[2])
    func = None
    for _ in range(n):
        nm, at = name(b, at)
        kind = b[at]
        idx, at = uleb(b, at + 1)
        if nm == "wai_abi" and kind == 0:
            func = idx
    if func is None:
        raise SystemExit("plant.py: no exported function wai_abi")
    count, at = uleb(b, codesec[2])
    for i in range(count):
        size_at = at
        size, body_at = uleb(b, at)
        end = body_at + size
        if i == func:
            # Past the body's local declarations: a vector of (count, value type).
            groups, p = uleb(b, body_at)
            for _ in range(groups):
                _, p = uleb(b, p)
                t = b[p]
                p += 1
                if t in (0x63, 0x64):
                    _, p = uleb(b, p)
            new_body = b[body_at:p] + code + b[p:end]
            planted_in_body = len(b[body_at:p])
            new_entry = enc(len(new_body)) + new_body
            content = b[codesec[2]:size_at] + new_entry + b[end:codesec[3]]
            head = bytes([10]) + enc(len(content))
            out = b[:codesec[1]] + head + content + b[codesec[3]:]
            offset = codesec[1] + len(head) + (size_at - codesec[2]) + len(enc(len(new_body))) + planted_in_body + rel
            with open(argv[2], "wb") as f:
                f.write(out)
            for level in ("integer", "float"):
                r = verdicts[level]
                print(f"{level} {'held' if r is None else f'refused {r} {offset}'}")
            return 0
        at = end
    raise SystemExit("plant.py: wai_abi has no body")


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