#!/usr/bin/env python3
"""build.py: the checker corpus of the WAISM profile, built from source.

wai/extensions/waism.md §9.1 gives each module of the checker corpus by its
construction: `minimal-int` and, for every other case but `component`, `minimal-int` with one
change. This script is that construction, byte by byte. It writes each module as
`<case>.wasm` beside itself, and `cases.txt`, the verdicts §9.1 states for them.

A verdict's offset is taken from the construction, never from a checker: each change marks
the item its row of §9.1 names (an instruction, a value type, an entry, a section's id byte)
as the bytes are put together, and the offset in `cases.txt` is where that mark lands in the
module. `at the module's length` is the module's length.

`cases.txt` has one line a verdict: `<case> <level> <caps> <verdict>`, where `<level>` is
`integer` or `float`, `<caps>` is `-` for the profile's caps (§4.2) or the caps the row
states, comma-joined `name=value`, and `<verdict>` is §4.3's line. A row of §9.1 at both
levels gives a line at each. Lines §9.1 does not state but §4.3 decides are marked
`# derived` after the verdict: the other level of a row that breaks only R3, and the
profile's caps for the rows that state lower ones. `grow-swallowed`, §9.2's planted module,
is here too: it meets R1 to R12, and R13, which it breaks, is no checker's rule.

§9.3's probe is here as well, `probe.wasm`, held at the float level, with `probe.txt`: each
case's operands and the result the deterministic profile prescribes, and the probe's bytes,
the whole output its `wai_run` gives on an engine set to that profile's NaN rule. Beside it,
`probe-nan-branch.wasm` is the probe with a branch on whether each result is a NaN that
writes the canonical NaN itself: it meets R1 to R12 too, and gives the probe's bytes on every
engine, so it cannot tell an engine set to the NaN rule from one that is not. That is F1, a
property of the probe on engines, which no checker decides; `../wai-det-jpegai/host --f1`
shows the probe holds it and the planted copy fails it.

    python3 build.py               write the modules, cases.txt and probe.txt
    python3 build.py --check       exit 1 unless the files here are exactly what it would write
    python3 build.py --limits DIR  write the validation limits' corpus to DIR (below)

The validation limits of §4.2 are part of R2. `--limits` writes, for each, `minimal-int`
with the part it bounds at the limit and one past it, and DIR/cases.txt in the same form:
at the limit the verdict is that of the rule the part breaks, if any (R12 for a tag section,
R10 for an extra export, held for unused globals); one past it, `refused R2 -`. These are
large modules (up to 10 MB), so they are generated where they are checked and are not
committed.

Standard library only.
"""
import os
import struct
import sys

HERE = os.path.dirname(os.path.abspath(__file__))


# ------------------------------------------------------------------ bytes with marks


class Blob:
    """Bytes, with named offsets into them that follow the bytes when they are joined."""

    def __init__(self, data=b"", marks=None):
        self.data = bytes(data)
        self.marks = dict(marks or {})

    def __add__(self, other):
        marks = dict(self.marks)
        for k, v in other.marks.items():
            assert k not in marks, k
            marks[k] = v + len(self.data)
        return Blob(self.data + other.data, marks)

    def __len__(self):
        return len(self.data)


def b(*xs):
    return Blob(bytes(xs))


def cat(*blobs):
    data, marks, at = [], {}, 0
    for x in blobs:
        for k, v in x.marks.items():
            assert k not in marks, k
            marks[k] = v + at
        data.append(x.data)
        at += len(x.data)
    return Blob(b"".join(data), marks)


def mark(name, blob):
    assert name not in blob.marks
    return Blob(blob.data, {name: 0, **blob.marks})


def uleb(n):
    out = bytearray()
    while True:
        c = n & 0x7F
        n >>= 7
        out.append(c | (0x80 if n else 0))
        if not n:
            return Blob(out)


def sleb(n):
    out = bytearray()
    while True:
        c = n & 0x7F
        n >>= 7
        done = (n == 0 and not c & 0x40) or (n == -1 and c & 0x40)
        out.append(c | (0 if done else 0x80))
        if done:
            return Blob(out)


def vec(items):
    return uleb(len(items)) + cat(*items)


def name(s):
    raw = s.encode("utf-8")
    return uleb(len(raw)) + Blob(raw)


def sized(blob):
    return uleb(len(blob)) + blob


# ------------------------------------------------------------------ encodings

I32, I64, F32, F64, V128 = 0x7F, 0x7E, 0x7D, 0x7C, 0x7B
FUNCREF, EXTERNREF = 0x70, 0x6F
END, DROP = 0x0B, 0x1A


def functype(params, results, item_param=None):
    """A function type; with `item_param`, the value type of that parameter is marked `item`."""
    ps = [mark("item", b(t)) if i == item_param else b(t) for i, t in enumerate(params)]
    return b(0x60) + vec(ps) + vec([b(t) for t in results])


def limits(minimum, maximum=None, i64=False):
    flag = (0x04 if i64 else 0) | (0x01 if maximum is not None else 0)
    out = b(flag) + uleb(minimum)
    return out + uleb(maximum) if maximum is not None else out


def i32_const(n):
    return b(0x41) + sleb(n)


def i64_const(n):
    return b(0x42) + sleb(n)


def f32_const(x):
    return b(0x43) + Blob(struct.pack("<f", x))


def f64_const(x):
    return b(0x44) + Blob(struct.pack("<d", x))


def v128_const(lanes=bytes(16)):
    return b(0xFD) + uleb(0x0C) + Blob(lanes)


def simd(op):
    return b(0xFD) + uleb(op)


def memarg(align, offset=0):
    return uleb(align) + uleb(offset)


def export(nm, kind, index):
    return name(nm) + b(kind) + uleb(index)


def body(instrs, locals_=None):
    """A function body: its local declarations (none by default), then its instructions,
    which end with `end`."""
    return sized((locals_ if locals_ is not None else uleb(0)) + instrs)


def section(sid, content, tag=None):
    head = b(sid)
    if tag:
        head = mark(tag, head)
    return head + sized(content)


# ------------------------------------------------------------------ minimal-int

T_ABI, T_ALLOC, T_RUN = 0, 1, 2       # () -> i32, (i32) -> i32, (i32 x 6) -> i32
EXPORTS = ["memory", "wai_abi", "wai_alloc", "wai_run", "wai_out_ptr", "wai_out_len"]


def store_and_return(addr=None):
    """wai_run's own instructions: store one fixed byte at offset 0, return 0."""
    return cat(addr or i32_const(0), i32_const(0x2A), b(0x3A), memarg(0), i32_const(0), b(END))


class Module:
    """minimal-int, as §9.1 states it, with every part open to a change."""

    def __init__(self):
        self.types = [functype([], [I32]), functype([I32], [I32]), functype([I32] * 6, [I32])]
        self.imports = []
        self.funcs = [T_ABI, T_ALLOC, T_RUN, T_ABI, T_ABI]
        self.tables = []
        self.memories = [limits(1, 2)]
        self.tags = None
        self.globals = []
        self.exports = None          # the six of §5 unless set
        self.func_base = 0           # the index of the first defined function (after imports)
        self.start = None
        self.elems = None
        self.run_prefix = Blob()     # added to wai_run before its store
        self.run_locals = None
        self.run_rest = None         # wai_run's instructions after the prefix, if changed
        self.abi_body = cat(i32_const(1), b(END))
        self.alloc_body = cat(i32_const(0), b(END))
        self.out_len_body = cat(i32_const(1), b(END))
        self.extra_bodies = []       # bodies of functions after the five
        self.customs = []
        self.datas = None
        self.no_memory_export = False
        self.drop_exports = set()

    def default_exports(self):
        out = []
        for i, nm in enumerate(EXPORTS):
            if nm in self.drop_exports or (nm == "memory" and self.no_memory_export):
                continue
            out.append(export(nm, 0x02, 0) if nm == "memory" else export(nm, 0x00, self.func_base + i - 1))
        return out

    def build(self):
        exports = self.exports if self.exports is not None else self.default_exports()
        run = self.run_prefix + (self.run_rest if self.run_rest is not None else store_and_return())
        bodies = [
            body(self.abi_body),
            body(self.alloc_body),
            body(run, self.run_locals),
            body(cat(i32_const(0), b(END))),
            body(self.out_len_body),
        ] + self.extra_bodies
        parts = [Blob(b"\0asm") + Blob(struct.pack("<I", 1))]
        parts.append(section(1, vec(self.types)))
        if self.imports:
            parts.append(section(2, vec(self.imports)))
        parts.append(section(3, vec([uleb(t) for t in self.funcs]), "function-section"))
        if self.tables:
            parts.append(section(4, vec(self.tables)))
        if self.memories:
            parts.append(section(5, vec(self.memories)))
        if self.tags is not None:
            parts.append(section(13, vec(self.tags), "tag-section"))
        if self.globals:
            parts.append(section(6, vec(self.globals)))
        parts.append(section(7, vec(exports)))
        if self.start is not None:
            parts.append(section(8, uleb(self.start), "start-section"))
        if self.elems is not None:
            parts.append(section(9, vec(self.elems)))
        parts.append(section(10, vec(bodies)))
        if self.datas is not None:
            parts.append(section(11, vec(self.datas)))
        for c in self.customs:
            parts.append(section(0, c))
        return cat(*parts)


# ------------------------------------------------------------------ the cases


def case_table():
    """(case, module blob, [(level, caps, verdict-spec, derived)]) for every case.

    A verdict-spec is `held`, `R1`, `R2`, or (rule, mark) where mark names the item, or
    `end` for the module's length."""
    cases = []
    BOTH = ("integer", "float")

    def add(nm, blob, rows):
        cases.append((nm, blob, rows))

    def both(spec, caps="-"):
        return [(lvl, caps, spec, False) for lvl in BOTH]

    def int_only(spec):
        return [("integer", "-", spec, False), ("float", "-", "held", True)]

    m = Module()
    minimal = m.build()

    # -- planted
    add("oversize", minimal, both("R1", f"module_bytes={len(minimal) - 1}")
        + [(lvl, "-", "held", True) for lvl in BOTH])

    m = Module()
    m.abi_body = cat(i64_const(1), b(END))
    add("invalid", m.build(), both("R2"))

    m = Module()
    m.memories = [b(0x03) + uleb(1) + uleb(2)]                      # shared, with a maximum
    m.run_prefix = cat(i32_const(0), b(0xFE) + uleb(0x10) + memarg(2), b(DROP))  # i32.atomic.load
    add("shared-memory", m.build(), both("R2"))

    add("component", Blob(b"\0asm") + Blob(bytes([0x0D, 0x00, 0x01, 0x00])), both("R2"))

    m = Module()
    m.run_prefix = cat(mark("item", f32_const(1.0)), f32_const(2.0), b(0x92), b(DROP))
    add("scalar-float", m.build(), int_only(("R3", "item")))

    m = Module()
    m.run_prefix = cat(i32_const(0), b(0x0F), mark("item", b(0x92)), b(DROP))
    add("float-unreachable", m.build(), int_only(("R3", "item")))

    m = Module()
    # block (result f32) unreachable end, its result dropped: the item is the block type.
    m.run_prefix = cat(b(0x02), mark("item", b(F32)), b(0x00), b(END), b(DROP))
    add("float-block-type", m.build(), int_only(("R3", "item")))

    m = Module()
    # return, then select (result f32) of a polymorphic stack: the item is its value type.
    m.run_prefix = cat(i32_const(0), b(0x0F), i32_const(0), b(0x1C) + uleb(1), mark("item", b(F32)), b(DROP))
    add("float-select-type", m.build(), int_only(("R3", "item")))

    m = Module()
    m.globals = [mark("item", b(F64)) + b(0x00) + f64_const(0.0) + b(END)]
    add("float-constant", m.build(), int_only(("R3", "item")))

    m = Module()
    m.types.append(functype([F64], [], item_param=0))
    add("float-signature", m.build(), int_only(("R3", "item")))

    m = Module()
    m.run_prefix = cat(v128_const(), v128_const(), mark("item", simd(0xE4)), b(DROP))       # f32x4.add
    add("float-vector", m.build(), int_only(("R3", "item")))

    m = Module()
    m.run_prefix = cat(v128_const(), mark("item", simd(0xFA)), b(DROP))       # f32x4.convert_i32x4_s
    add("float-conversion", m.build(), int_only(("R3", "item")))

    m = Module()
    # i32.trunc_sat_f32_s of f32.reinterpret_i32
    m.run_prefix = cat(i32_const(0), mark("item", b(0xBE)), b(0xFC) + uleb(0), b(DROP))
    add("float-truncation", m.build(), int_only(("R3", "item")))

    m = Module()
    m.run_prefix = cat(i32_const(0), mark("item", b(0x2A) + memarg(2)), b(DROP))           # f32.load
    add("float-load", m.build(), int_only(("R3", "item")))

    for nm, op, arity in [("relaxed-swizzle", 0x100, 2), ("relaxed-laneselect", 0x10A, 3),
                          ("relaxed-dot", 0x112, 2)]:
        m = Module()
        m.run_prefix = cat(*[v128_const()] * arity, mark("item", simd(op)), b(DROP))
        add(nm, m.build(), both(("R4", "item")))

    m = Module()
    m.run_prefix = cat(*[v128_const()] * 3, mark("item", simd(0x105)), b(DROP))  # f32x4.relaxed_madd
    add("relaxed-madd", m.build(), [("float", "-", ("R4", "item"), False),
                                    ("integer", "-", ("R3", "item"), True)])

    m = Module()
    m.imports = [mark("item", name("env") + name("clock") + b(0x00) + uleb(T_ABI))]
    m.func_base = 1
    add("import-function", m.build(), both(("R5", "item")))

    m = Module()
    m.imports = [mark("item", name("env") + name("memory") + b(0x02) + limits(1, 2))]
    m.memories = []
    add("import-memory", m.build(), both(("R5", "item")))

    m = Module()
    m.imports = [mark("item", name("env") + name("g") + b(0x03) + b(I32, 0x00))]
    add("import-global", m.build(), both(("R5", "item")))

    m = Module()
    m.memories = []
    m.no_memory_export = True
    m.run_rest = cat(i32_const(0), b(END))
    add("no-memory", m.build(), both(("R6", "end")))

    m = Module()
    m.memories = [limits(1, 2), mark("item", limits(1, 2))]
    add("two-memories", m.build(), both(("R6", "item")))

    m = Module()
    m.memories = [mark("item", limits(1, 2, i64=True))]
    m.run_rest = store_and_return(addr=i64_const(0))
    add("memory64", m.build(), both(("R6", "item")))

    m = Module()
    m.memories = [mark("item", limits(1))]
    add("no-maximum", m.build(), both(("R6", "item")))

    m = Module()
    m.memories = [mark("item", limits(1, 17))]
    add("over-pages", m.build(), both(("R6", "item"), "memory_pages=16")
        + [(lvl, "-", "held", True) for lvl in BOTH])

    def fixed_table():
        return b(FUNCREF) + limits(0, 0)

    m = Module()
    m.tables = [fixed_table(), mark("item", fixed_table())]
    add("two-tables", m.build(), both(("R7", "item")))

    m = Module()
    m.tables = [mark("item", b(EXTERNREF) + limits(0, 0))]
    add("externref-table", m.build(), both(("R7", "item")))

    m = Module()
    m.tables = [mark("item", b(FUNCREF) + limits(0, 1))]
    add("table-unfixed", m.build(), both(("R7", "item")))

    ref_null_func = b(0xD0, FUNCREF)
    for nm, prefix in [
        ("table-grow", cat(ref_null_func, i32_const(0), mark("item", b(0xFC) + uleb(15) + uleb(0)), b(DROP))),
        ("table-set", cat(i32_const(0), ref_null_func, mark("item", b(0x26) + uleb(0)))),
        ("table-fill", cat(i32_const(0), ref_null_func, i32_const(0), mark("item", b(0xFC) + uleb(17) + uleb(0)))),
    ]:
        m = Module()
        m.tables = [fixed_table()]
        m.run_prefix = prefix
        add(nm, m.build(), both(("R7", "item")))

    m = Module()
    m.types.append(functype([], []))
    m.funcs.append(3)
    m.extra_bodies = [body(b(END))]
    m.start = 5
    add("start", m.build(), both(("R8", "start-section")))

    m = Module()
    m.types.append(functype([FUNCREF], [], item_param=0))
    m.funcs.append(3)
    m.extra_bodies = [body(b(END))]
    add("funcref-parameter", m.build(), both(("R9", "item")))

    m = Module()
    m.globals = [mark("item", b(EXTERNREF)) + b(0x00) + b(0xD0, EXTERNREF) + b(END)]
    add("externref-global", m.build(), both(("R9", "item")))

    m = Module()
    m.run_locals = uleb(1) + uleb(1) + mark("item", b(FUNCREF))
    add("funcref-local", m.build(), both(("R9", "item")))

    m = Module()
    m.types[T_RUN] = functype([I32] * 5 + [I64], [I32])
    ex = m.default_exports()
    ex[3] = mark("item", ex[3])
    m.exports = ex
    add("i64-export", m.build(), both(("R10", "item")))

    m = Module()
    m.exports = m.default_exports() + [mark("item", export("wai_extra", 0x00, 0))]
    add("extra-export", m.build(), both(("R10", "item")))

    m = Module()
    m.globals = [b(I32, 0x00) + i32_const(0) + b(END)]
    m.exports = m.default_exports() + [mark("item", export("wai_global", 0x03, 0))]
    add("exported-global", m.build(), both(("R10", "item")))

    m = Module()
    m.drop_exports = {"wai_run"}
    add("missing-run", m.build(), both(("R10", "end")))

    m = Module()
    m.funcs += [T_ABI] * 4
    m.extra_bodies = [body(cat(i32_const(0), b(END)))] * 4
    add("over-functions", m.build(), both(("R11", "function-section"), "functions=8")
        + [(lvl, "-", "held", True) for lvl in BOTH])

    m = Module()
    m.run_prefix = mark("item", b(0x12) + uleb(0))                   # return_call wai_abi
    add("tail-call", m.build(), both(("R12", "item")))

    m = Module()
    m.types.append(functype([], []))
    m.tags = [b(0x00) + uleb(3)]
    m.run_prefix = b(0x08) + uleb(0)                                  # throw 0
    add("exception", m.build(), both(("R12", "tag-section")))

    m = Module()
    m.types.append(mark("item", b(0x5F) + uleb(0)))                   # a struct type, no fields
    add("struct-type", m.build(), both(("R12", "item")))

    m = Module()
    m.elems = [uleb(3) + b(0x00) + vec([uleb(0)])]                    # declarative: wai_abi
    m.run_prefix = cat(b(0xD2) + uleb(0), mark("item", b(0x14) + uleb(T_ABI)), b(DROP))
    add("call-ref", m.build(), both(("R12", "item")))

    # -- conforming
    add("minimal-int", minimal, both("held"))

    m = Module()
    m.run_prefix = cat(
        v128_const(), v128_const(), simd(0xBA), b(DROP),              # i32x4.dot_i16x8_s
        v128_const(), v128_const(), simd(0x0E), b(DROP),              # i8x16.swizzle
        v128_const(), v128_const(), simd(0x0D) + Blob(bytes(range(16))), b(DROP),  # i8x16.shuffle
        v128_const(), v128_const(), simd(0x82), b(DROP),              # i16x8.q15mulr_sat_s
    )
    add("vector-int", m.build(), both("held"))

    m = Module()
    m.types.append(functype([], [I32, I32]))
    m.run_prefix = cat(
        i32_const(0), i32_const(0), i32_const(0), b(0xFC) + uleb(10) + uleb(0) + uleb(0),  # memory.copy
        i32_const(0), i32_const(0), i32_const(0), b(0xFC) + uleb(11) + uleb(0),            # memory.fill
        i32_const(0), b(0xC0), b(DROP),                                                    # i32.extend8_s
        b(0x02) + sleb(3), i32_const(1), i32_const(2), b(END), b(DROP), b(DROP),           # a block of two results
    )
    add("bulk-and-sign", m.build(), both("held"))

    m = Module()
    m.tables = [b(FUNCREF) + limits(1, 1)]
    m.elems = [uleb(0) + i32_const(0) + b(END) + vec([uleb(0)])]
    m.run_prefix = cat(i32_const(0), b(0x11) + uleb(T_ABI) + uleb(0), b(DROP))  # call_indirect
    add("fixed-table", m.build(), both("held"))

    m = Module()
    m.memories = [limits(1, 1)]
    add("fixed-memory", m.build(), both("held"))

    m = Module()
    m.customs = [name("profile") + Blob(b"waism/1-det")]
    add("custom-section", m.build(), both("held"))

    m = Module()
    m.run_prefix = cat(mark("item", f64_const(1.0)), f64_const(2.0), b(0xA0), b(DROP))
    add("minimal-float", m.build(), [("float", "-", "held", False), ("integer", "-", ("R3", "item"), False)])

    # -- §9.2: breaks R13 only, which no checker decides
    m = Module()
    m.run_prefix = cat(i32_const(1), b(0x40) + uleb(0), b(DROP))      # memory.grow 1, result dropped
    add("grow-swallowed", m.build(), both("held"))

    # -- §9.3: the probe, held at the float level; and the probe with a NaN-dependent
    # branch planted, which meets R1 to R12 too and fails F1 on an engine (no checker's rule)
    for nm, planted in [("probe", False), ("probe-nan-branch", True)]:
        blob, _ = probe_module(planted)
        add(nm, blob, [("float", "-", "held", False), ("integer", "-", ("R3", "item"), True)])

    return cases


# ------------------------------------------------------------------ §9.3: the probe

CANONICAL = {"f32": 0x7FC00000, "f64": 0x7FF8000000000000}

# Each case of §9.3's table: its number, its operation, the type of its operands, and the
# operands as bits (a vector operand as its lanes, lane 0 first).
PROBE = [
    (1, "f32.add", "f32", [0xFFA00001, 0x3F800000]),
    (2, "f32.mul", "f32", [0xFFA00001, 0x40000000]),
    (3, "f32.div", "f32", [0x00000000, 0x00000000]),
    (4, "f32.sub", "f32", [0x7F800000, 0x7F800000]),
    (5, "f32.sqrt", "f32", [0xBF800000]),
    (6, "f32.min", "f32", [0xFFA00001, 0x3F800000]),
    (7, "f32.nearest", "f32", [0xFFA00001]),
    (8, "f32.demote_f64", "f64", [0xFFF4000000000001]),
    (9, "f64.add", "f64", [0xFFF4000000000001, 0x3FF0000000000000]),
    (10, "f64.promote_f32", "f32", [0xFFA00001]),
    (11, "f32x4.add", "f32x4", [(0xFFA00001, 0x7FA00001, 0xFFC00000, 0x7F800000),
                                (0x3F800000, 0x3F800000, 0x3F800000, 0xFF800000)]),
    (12, "f64x2.mul", "f64x2", [(0xFFF4000000000001, 0x7FF8000000000001),
                                (0x4000000000000000, 0x4000000000000000)]),
]

# The result type of each operation, and its opcode.
PROBE_OPS = {
    "f32.add": ("f32", lambda: b(0x92)), "f32.mul": ("f32", lambda: b(0x94)),
    "f32.div": ("f32", lambda: b(0x95)), "f32.sub": ("f32", lambda: b(0x93)),
    "f32.sqrt": ("f32", lambda: b(0x91)), "f32.min": ("f32", lambda: b(0x96)),
    "f32.nearest": ("f32", lambda: b(0x90)), "f32.demote_f64": ("f32", lambda: b(0xB6)),
    "f64.add": ("f64", lambda: b(0xA0)), "f64.promote_f32": ("f64", lambda: b(0xBB)),
    "f32x4.add": ("f32x4", lambda: simd(0xE4)), "f64x2.mul": ("f64x2", lambda: simd(0xF2)),
}

# A type's size, and its load and store (natural alignment).
PROBE_TYPES = {
    "f32": (4, lambda: b(0x2A) + memarg(2), lambda: b(0x38) + memarg(2)),
    "f64": (8, lambda: b(0x2B) + memarg(3), lambda: b(0x39) + memarg(3)),
    "f32x4": (16, lambda: simd(0x00) + memarg(4), lambda: simd(0x0B) + memarg(4)),
    "f64x2": (16, lambda: simd(0x00) + memarg(4), lambda: simd(0x0B) + memarg(4)),
}

PROBE_DATA = 256        # where the operands lie
PROBE_HEAP = 1024       # where wai_alloc's first region begins


def bits(ty, value):
    """An operand or a result as bytes, little-endian, lane 0 first."""
    if ty == "f32":
        return struct.pack("<I", value)
    if ty == "f64":
        return struct.pack("<Q", value)
    lane = "<I" if ty == "f32x4" else "<Q"
    return b"".join(struct.pack(lane, v) for v in value)


def expected(ty):
    """What the deterministic profile prescribes for a NaN result of type `ty`: the positive
    canonical NaN, in every lane."""
    if ty in CANONICAL:
        return bits(ty, CANONICAL[ty])
    lanes = 4 if ty == "f32x4" else 2
    return bits(ty, (CANONICAL["f32" if ty == "f32x4" else "f64"],) * lanes)


def nan_branch(ty):
    """The planted fault: a branch on whether the result is a NaN that writes the canonical
    NaN itself, so the module's output no longer shows the engine's NaN. With the result on
    the stack, it leaves the result, or the canonical NaN in its place. Locals 6, 7 and 8 of
    `wai_run` are its `f32`, `f64` and `v128`."""
    if ty in ("f32", "f64"):
        idx, ne, vt = (6, 0x5C, F32) if ty == "f32" else (7, 0x62, F64)
        const = b(0x43 if ty == "f32" else 0x44) + Blob(bits(ty, CANONICAL[ty]))
        return cat(b(0x21, idx), b(0x20, idx), b(0x20, idx), b(ne),
                   b(0x04, vt), const, b(0x05), b(0x20, idx), b(END))
    ne = simd(0x42 if ty == "f32x4" else 0x48)
    mask = cat(b(0x20, 8), b(0x20, 8), ne)
    return cat(b(0x21, 8), mask, simd(0x53),                                  # v128.any_true
               b(0x04, V128), v128_const(expected(ty)), b(0x20, 8), mask, simd(0x52),  # v128.bitselect
               b(0x05), b(0x20, 8), b(END))


def probe_module(planted=False):
    """§9.3's probe: `minimal-int` whose `wai_run`, for any input, writes the result of each
    case at the output's offset, in order, and returns 0. Its operands lie in a data segment
    and are loaded at run time, so no compiler folds an operation. Its memory is one page,
    its minimum its maximum, with no `memory.grow` (R13); `wai_alloc` hands out regions from
    a bump pointer, a mutable `i32` global, above the operands, and 0 past the page.

    With `planted`, each result passes through `nan_branch` before it is stored."""
    data, operand_at, out = bytearray(), [], 0
    instrs, out_at = [], []
    for n, op, ty, operands in PROBE:
        size, load, _ = PROBE_TYPES[ty]
        ats = []
        for v in operands:
            while len(data) % size:
                data.append(0)
            ats.append(PROBE_DATA + len(data))
            data += bits(ty, v)
        operand_at.append(ats)
    first = True
    for (n, op, ty, operands), ats in zip(PROBE, operand_at):
        rty, opcode = PROBE_OPS[op]
        size, _, store = PROBE_TYPES[rty]
        _, load, _ = PROBE_TYPES[ty]
        while out % size:
            out += 1
        out_at.append(out)
        loads = []
        for at in ats:
            ld = load()
            if first and not planted:
                ld = mark("item", ld)
                first = False
            loads.append(cat(i32_const(at), ld))
        instrs.append(cat(i32_const(out), *loads, opcode(),
                          nan_branch(rty) if planted else Blob(), store()))
        out += size
    m = Module()
    m.memories = [limits(1, 1)]
    m.globals = [b(I32, 0x01) + i32_const(PROBE_HEAP) + b(END)]
    m.datas = [uleb(0) + i32_const(PROBE_DATA) + b(END) + uleb(len(data)) + Blob(bytes(data))]
    m.alloc_body = cat(
        b(0x20, 0x00), i32_const(65536), b(0x23, 0x00), b(0x6B), b(0x4B),   # len > 65536 - top
        b(0x04, I32), i32_const(0),
        b(0x05), b(0x23, 0x00), b(0x23, 0x00), b(0x20, 0x00), b(0x6A), b(0x24, 0x00),
        b(END), b(END))
    m.out_len_body = cat(i32_const(out), b(END))
    m.run_rest = cat(*instrs, i32_const(0), b(END))
    if planted:
        m.run_locals = uleb(3) + uleb(1) + mark("item", b(F32)) + uleb(1) + b(F64) + uleb(1) + b(V128)
    assert len(data) + PROBE_DATA <= PROBE_HEAP and out <= PROBE_DATA
    return m.build(), out


def probe_vectors():
    """probe.txt: §9.3's cases, each with its operands and the result the deterministic
    profile prescribes, and the probe's bytes, the whole output."""
    lines = [
        "# The probe of waism.md §9.3, written by build.py: `<case> <operation> <operands>",
        "# <result>`, one line a case, each operand and result as its bits, a vector's lanes",
        "# comma-joined with lane 0 first, the result the one the deterministic profile",
        "# prescribes, little-endian hex. Last, `bytes <hex>`: the probe's bytes, the output",
        "# of its wai_run, every result in order. probe.wasm is the probe; probe-nan-branch.wasm",
        "# is it with a NaN-dependent branch planted after each operation.",
    ]
    whole = b""
    for n, op, ty, operands in PROBE:
        rty, _ = PROBE_OPS[op]
        width = 8 if ty in ("f32", "f32x4") else 16
        ops = " ".join(
            ",".join(f"0x{lane:0{width}x}" for lane in (v if isinstance(v, tuple) else (v,)))
            for v in operands)
        want = expected(rty)
        whole += want
        lines.append(f"{n} {op} {ops} {want.hex()}")
    lines.append(f"bytes {whole.hex()}")
    return ("\n".join(lines) + "\n").encode("utf-8"), whole


def verdict(spec, blob):
    if spec == "held":
        return "held"
    if spec == "R1":
        return "refused R1 0"
    if spec == "R2":
        return "refused R2 -"
    rule, where = spec
    at = len(blob) if where == "end" else blob.marks[where]
    return f"refused {rule} {at}"


def outputs():
    files = {}
    lines = [
        "# The checker corpus of waism.md §9.1, §9.2's grow-swallowed and §9.3's probe,",
        "# written by build.py: `<case> <level> <caps> <verdict>`, one line a verdict.",
    ]
    for nm, blob, rows in case_table():
        files[f"{nm}.wasm"] = blob.data
        for level, caps, spec, derived in rows:
            line = f"{nm} {level} {caps} {verdict(spec, blob)}"
            lines.append(line + ("  # derived" if derived else ""))
    files["cases.txt"] = ("\n".join(lines) + "\n").encode("utf-8")
    text, whole = probe_vectors()
    assert len(whole) == probe_module()[1]
    files["probe.txt"] = text
    return files


def limit_cases():
    """(case, blob, verdict-spec) for each validation limit of §4.2, at it and one past it."""
    out = []
    nop_type = functype([], [])

    def both(nm, at_limit, past, spec):
        out.append((f"{nm}-at-limit", at_limit, spec))
        out.append((f"{nm}-past-limit", past, "R2"))

    def module(change):
        m = Module()
        change(m)
        return m.build()

    # types: entries of the type section, and the types they define
    def plain_types(n):
        return lambda m: m.types.extend([nop_type] * (n - 3))
    both("types", module(plain_types(1_000_000)), module(plain_types(1_000_001)), "held")

    def empty_groups(n):
        def change(m):
            m.types.append(mark("item", b(0x4E) + uleb(0)))
            m.types.extend([b(0x4E) + uleb(0)] * (n - 4))
        return change
    both("type-entries", module(empty_groups(1_000_000)), module(empty_groups(1_000_001)), ("R12", "item"))

    def one_group(n):
        return lambda m: m.types.append(mark("item", b(0x4E) + vec([nop_type] * (n - 3))))
    both("types-in-a-group", module(one_group(1_000_000)), module(one_group(1_000_001)), ("R12", "item"))

    def functions(n):
        def change(m):
            m.funcs += [T_ABI] * (n - 5)
            m.extra_bodies = [body(cat(i32_const(0), b(END)))] * (n - 5)
        return change
    both("functions", module(functions(1_000_000)), module(functions(1_000_001)), ("R11", "function-section"))

    def globals(n):
        return lambda m: m.globals.extend([b(I32, 0x00) + i32_const(0) + b(END)] * n)
    both("globals", module(globals(1_000_000)), module(globals(1_000_001)), "held")

    def tags(n):
        def change(m):
            m.types.append(nop_type)
            m.tags = [b(0x00) + uleb(3)] * n
        return change
    both("tags", module(tags(1_000_000)), module(tags(1_000_001)), ("R12", "tag-section"))

    def tables(n):
        return lambda m: m.tables.extend([b(FUNCREF) + limits(0, 0), mark("item", b(FUNCREF) + limits(0, 0))]
                                         + [b(FUNCREF) + limits(0, 0)] * (n - 2))
    both("tables", module(tables(100_000)), module(tables(100_001)), ("R7", "item"))

    def memories(n):
        return lambda m: m.memories.extend([mark("item", limits(1, 2))] + [limits(1, 2)] * (n - 2))
    both("memories", module(memories(100)), module(memories(101)), ("R6", "item"))

    def elements(n):
        return lambda m: setattr(m, "elems", [uleb(1) + b(0x00) + uleb(0)] * n)
    both("element-segments", module(elements(100_000)), module(elements(100_001)), "held")

    def datas(n):
        return lambda m: setattr(m, "datas", [uleb(1) + uleb(0)] * n)
    both("data-segments", module(datas(100_000)), module(datas(100_001)), "held")

    def items(n):
        return lambda m: setattr(m, "elems", [uleb(1) + b(0x00) + uleb(n) + Blob(bytes(n))])
    both("segment-items", module(items(10_000_000)), module(items(10_000_001)), "held")

    def export_name(n):
        return lambda m: setattr(m, "exports", m.default_exports() + [mark("item", export("x" * n, 0x00, 0))])
    both("name-of-an-export", module(export_name(100_000)), module(export_name(100_001)), ("R10", "item"))

    def import_name(n):
        def change(m):
            m.imports = [mark("item", name("env") + name("g" * n) + b(0x03) + b(I32, 0x00))]
        return change
    both("name-of-an-import", module(import_name(100_000)), module(import_name(100_001)), ("R5", "item"))

    def custom_name(n):
        return lambda m: setattr(m, "customs", [name("c" * n)])
    both("name-of-a-custom-section", module(custom_name(100_000)), module(custom_name(100_001)), "held")

    def body_bytes(n):
        # wai_run's body: its declaration of no locals (1 byte), nops, then its own 13 bytes.
        def change(m):
            m.run_prefix = Blob(b"\x01" * (n - 1 - len(store_and_return())))
        return change
    both("body", module(body_bytes(7_654_321)), module(body_bytes(7_654_322)), "held")

    def locals_(n):
        # wai_run's six parameters and n - 6 locals.
        return lambda m: setattr(m, "run_locals", uleb(1) + uleb(n - 6) + b(I32))
    both("locals", module(locals_(50_000)), module(locals_(50_001)), "held")

    def params(n):
        return lambda m: m.types.append(functype([I32] * n, []))
    both("parameters", module(params(1_000)), module(params(1_001)), "held")

    def results(n):
        return lambda m: m.types.append(functype([], [I32] * n))
    both("results", module(results(1_000)), module(results(1_001)), "held")

    def fields(n):
        return lambda m: m.types.append(mark("item", b(0x5F) + vec([b(I32, 0x00)] * n)))
    both("struct-fields", module(fields(10_000)), module(fields(10_001)), ("R12", "item"))

    def catches(n):
        # A block, and in it a try_table whose catch_all clauses all name the block.
        return lambda m: setattr(m, "run_prefix", cat(
            b(0x02, 0x40), mark("item", b(0x1F, 0x40)), uleb(n), Blob(b"\x02\x00" * n), b(END), b(END)))
    both("catches", module(catches(10_000)), module(catches(10_001)), ("R12", "item"))

    def depth(d):
        # A chain of struct subtypes: the first declares no supertype (depth 0), each next
        # names the one before.
        def change(m):
            m.types.append(mark("item", b(0x50) + uleb(0) + b(0x5F) + uleb(0)))
            for k in range(1, d + 1):
                m.types.append(b(0x50) + uleb(1) + uleb(3 + k - 1) + b(0x5F) + uleb(0))
        return change
    both("subtype-depth", module(depth(63)), module(depth(64)), ("R12", "item"))

    def interface(total):
        # The six exports of §5 weigh 23 (memory 1; wai_abi, wai_out_ptr, wai_out_len 3 each;
        # wai_alloc 4; wai_run 9); each further export of the memory weighs 1.
        def change(m):
            extra = total - 23
            names = [f"m{i:x}" for i in range(extra)]
            ex = m.default_exports()
            ex.append(mark("item", export(names[0], 0x02, 0)))
            ex.extend(export(nm, 0x02, 0) for nm in names[1:])
            m.exports = ex
        return change
    both("interface", module(interface(999_998)), module(interface(999_999)), ("R10", "item"))
    return out


def write_limits(directory):
    os.makedirs(directory, exist_ok=True)
    lines = ["# The validation limits of waism.md §4.2, written by build.py --limits."]
    for nm, blob, spec in limit_cases():
        assert len(blob) <= 16_777_216, nm
        with open(os.path.join(directory, f"{nm}.wasm"), "wb") as f:
            f.write(blob.data)
        for level in ("integer", "float"):
            lines.append(f"{nm} {level} - {verdict(spec, blob)}")
    with open(os.path.join(directory, "cases.txt"), "w", encoding="utf-8") as f:
        f.write("\n".join(lines) + "\n")
    print(f"build.py: wrote {len(lines) - 1} verdicts of the validation limits to {directory}")
    return 0


def main(argv):
    if len(argv) == 3 and argv[1] == "--limits":
        return write_limits(argv[2])
    check = argv[1:] == ["--check"]
    if argv[1:] and not check:
        print("usage: build.py [--check | --limits DIR]", file=sys.stderr)
        return 2
    want = outputs()
    have = {n for n in os.listdir(HERE) if n.endswith(".wasm") or n in ("cases.txt", "probe.txt")}
    if check:
        bad = []
        for n, data in sorted(want.items()):
            p = os.path.join(HERE, n)
            if not os.path.isfile(p):
                bad.append(f"{n}: missing")
            elif open(p, "rb").read() != data:
                bad.append(f"{n}: differs from what build.py writes")
        bad += [f"{n}: not written by build.py" for n in sorted(have - set(want))]
        for line in bad:
            print(f"build.py --check: {line}", file=sys.stderr)
        if not bad:
            print(f"build.py --check: {len(want) - 2} modules, cases.txt and probe.txt are what build.py writes")
        return 1 if bad else 0
    for n in sorted(have - set(want)):
        os.remove(os.path.join(HERE, n))
    for n, data in want.items():
        with open(os.path.join(HERE, n), "wb") as f:
            f.write(data)
    print(f"build.py: wrote {len(want) - 2} modules, cases.txt and probe.txt")
    return 0


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