open-standards/wai/live-conformance/dred_verify.py

#!/usr/bin/env python3
"""Second implementation of extensions/opus-dred.md (`wai.audio.opus.dred.int`, the integer
profile `dred-int/1` of draft-ietf-mlcodec-opus-dred-07), over the corpus `live-conformance/dred/`.

Written from the extension's text and the draft's. It shares no code with wai-rs: it reads an Opus
packet's framing (RFC 6716 section 3.2) and its padding's extensions, decodes the DRED payload
with RFC 6716's range decoder (section 4.1), transcribes the draft's tables for itself and checks
them against the digest the extension states, reads the weights file, turns every weight into an
integer from its float32 bits, computes the two activation tables, runs the decoder network in
integer arithmetic as section 4.6 states it, rebuilds each stream's reconstruction under its loss
pattern, and verifies each player's log with WAI's strict Ed25519 rule
(claim-conformance/_ed25519.py) and BLAKE3 written from its specification
(claim-conformance/_blake3.py). It reads section 8's codes and the digests of sections 4.1, 4.2,
4.5 and 4.6 from the extension.

    python3 -I dred_verify.py <live-conformance/dred> <weights-file>

The weights file is the draft's ([DRED-Weights]); `dred/fetch.sh` fetches it. The verifier gives
every vector of the corpus the outcome it states, checks that every reason of section 8 refuses a
vector, and then mutates copies of the corpus and requires every copy to fail. Exit status 0 only
if all of that holds. Python 3.8 or later; standard library only.
"""

import copy
import hashlib
import json
import os
import re
import struct
import sys
from operator import mul

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "claim-conformance"))
import _ed25519 as ED  # noqa: E402
from _blake3 import blake3  # noqa: E402

EXTENSION = os.path.join(HERE, "..", "extensions", "opus-dred.md")
CAPABILITY = "wai.audio.opus.dred.int"
LOWERING = "dred-int/1"


class Refused(Exception):
    def __init__(self, code, at=None):
        super().__init__(code)
        self.code = code
        self.at = at


def sha(b):
    return hashlib.sha256(b).digest()


def le32(v):
    return struct.pack("<i", v)


# ---- section 4.2: the draft's Tables 2, 4, 5, 6, 8 and 9, transcribed from the draft ----

LATENT_SCALE = bytes.fromhex(
    "ffdbbfa8978a87030202010101010101ffd5b69e8b7b6c5b534a40362c262320ffc89c785a412e201401020000000000"
    "ffd9bba4948c98020202010101010101ffd8b9a18e7f6e040302020101010101ffd2af937e6d61564f554a0903020201"
    "ffd7b79e8b7b6c020201010101010000ffd0aa8c756354493e37302b26211d19ffd0ab8d746156000000000000000000"
    "ffd0aa8c746153463b332b25201b1714ffd0ab8d766355493e352d220c080706ffd5b39986776f6b67601c0000000000"
    "ffd5b39984756a615c5b580a05030202ffcfa98b736152463c342d28231f1b18ffd2ae927c6a5d514742370f03020201"
    "ffdabca594867e030202010101010101ffd6b59b8c9f75020202010101010100ffd2ae927c6b5d5147423d2a08010101"
    "ffd9bba3908172040302020101010101ffd4b398847469645f680b0201010101ffd5b49a86776c675d57400b04030202"
    "ffd3b0947e6d60534d51490502010101ffd6b69c887786030202010101010101ffc79a775a422e030202010101010101"
    "ffdabda79794a6030202010101010101"
)
LATENT_DECAY = bytes.fromhex(
    "33261c140e0a0300000000000000000005030202010101000000000000000000f8f6f3efeae2d7c6aa00000000000000"
    "322419110c08000000000000000000004333251a120b060000000000000000006e5d4c3c2e221a130e0d060000000000"
    "3d2c1e130c0703000000000000000000c1b7ab9e90827567594d41372c21171049342111060100000000000000000000"
    "e4ded8d0c7beb4a89b8f817465544638a899887765554537291e140800000000574636291e17120e0a04000000000000"
    "5c4d40352d251c130b05000000000000c5bbafa295887b6f635950483f33281f7d69564333251b120c0a050000000000"
    "332519100a060300000000000000000023170f0a0800000000000000000000008372604f4032271c14100b0100000000"
    "4334261b130c060000000000000000005f4c3a2b1f150f0c08030000000000004f41312318100a050301000000000000"
    "7664524133261c130e0c040000000000372719100905010000000000000000009c8a77634d3117000000000000000000"
    "2b1f17100d0700000000000000000000"
)
LATENT_P0 = bytes.fromhex(
    "abbecfdde9f3fdfffffffffffffffffffbfdfefeffffffffffffffffffffffff040506080b0f151f30ffffffffffffff"
    "9eb2c6d8e8f6ffffffffffffffffffff92a5b7c8d7e6f3ffffffffffffffffff738291a1b0becbd7e2eff8ffffffffff"
    "8c9fb2c5d6e6f5ffffffffffffffffff3f49535c666f79848f9ba6b0bdcad6e0788da4bdd8f2ffffffffffffffffffff"
    "0e11151a1e242a3139414a545f6d7b88424b55616d7a8796a4b7d1f2ffffffff8d9dadbccad7e2ecf4fcffffffffffff"
    "8ea0b1c1cfdbe4edf5fbffffffffffff333b454f5a65717e8c99a7b3c0cdd8e15462708090a1b1c2d2e3f2ffffffffff"
    "9eb1c3d3e2eefaffffffffffffffffffaec2d5e6f6ffffffffffffffffffffff64717f8d9ba9b6c3d0e0edfbffffffff"
    "98aabbcbd9e6f1ffffffffffffffffff758495a5b5c4d2e1effcffffffffffffacbfcfdde8f0f6fbfdffffffffffffff"
    "6e7b8998a6b5c2d0dfeffaffffffffff8a9eb3c7daebffffffffffffffffffff5d6d8199b3cfe9ffffffffffffffffff"
    "a9bdd0e0eef9ffffffffffffffffffff"
)
STATE_SCALE = bytes.fromhex(
    "ffd0ab8d7664564a3f36302a231e1a1763584f47403a36363affdcaf8d715f5cffd2ae927b6a5b50453d392e29292926"
    "ffd0ab8e776456493e352d272624211f6c60564f494544484f7298e3ffe1c2a8ffcea788705c4d41362f28231f1a1714"
    "ffcea789705d4e42372f29231e1a1714ffd1ad917b695c51484343403a36312dffd4b399847368504b4434221a171412"
    "ffd2ae927c6b5f4e453d322b27252322ffd2ae927d6b5e52504e453c352f2b27ffcfa98b746153473e37312c2723201e"
    "d1bead9e8f837a8080b4d1f9ffe4cab5ffcfaa8c756253473e36322f2c2a2725887b6f655c575da79ecaddffe9c6ab96"
    "ffcb9f7c5f48362925382c221a140f0cffd4b196817165584d413128221d1917ffdabea7968b8a9d8b7a6b5f5349423f"
    "2c2825221f1d1b1c1bffcca17f634d3dffcea889715f504439322c27221d1a1793867a6e655e5c7d7697b7fcffe0c3ac"
    "a99a8c80766e6a696e93ade1ffe8cdb3ffc69a775b45342720231b15100d0a08ffd1ac8f79675a4f463f3a35312d2b2b"
    "ffcda788705d4e4238312b26211e1a18f1dbc7b5a5978d867fb6d2ffeacab4aaffd3b0957f6e6151483c31281e1a1714"
    "ffd0ab8e776455493f38322d28232120ffcfa98b746254473c342d28221c1815ffd4b297837367564c4237291d171311"
    "ffdabca59387859d8a77685b4f443d38ffd2ae917b695d544d4c443c36322e2affcfa98a725f51453b342e28221e1a17"
    "ffd0ac9079685a4d402a251b1613110fffe3cdbaaca1a09b8c756a5f4a48453effcfa98a7360514439312a241e191612"
    "ffd1ad907a675a4d4137332e26211d1affd4b39985766f6f6566635a4c413932ffd9baa28f82819384786a5e5148474a"
    "ffd0a98b736050413328221e1a17151248423e39342f2d2f34ffdeb190746054ffcfaa8c756254483f37302a25201e1d"
    "ffd3b1968272695e4e3623100c0b0908ffcea889715e4f4238312a251f1a1613ffcda5876f5b4c40362e27211d181512"
    "3d38332f2b28252427ffd4a6856a5444ffd2ae937d6c60575551494039332d29ffcfa98b746356493f39312a25201d1b"
    "ffcda688705c4d41372f28221d191613ffd4ad8c735f5044393129231e191512"
)
STATE_DECAY = bytes.fromhex(
    "cfc5baaea19386796a5c51423324180e00000000000000000055402d190b03008d7c6b5a4a3c30251a130f0703010000"
    "8d7b6855433325190f08040100000000020100000000000000082a545f4d3e30e2dbd3cabfb4a89b8d8072645545372b"
    "e2dcd4ccc3b8ada09386786a5b4d3f32624f3c2c2016100b080501000000000067534131261b14070301000000000000"
    "4e443d312318100a06030201000000004f3d2c1d120b07030200000000000000a0928475675a4f433930271e16100b07"
    "010000000000000000020c2420140b058a796756473a3029251c110a0604020100000000000000060f2a37433425190f"
    "fbfaf8f6f3f0ebe4e1ebe6dfd6cabba972614f3f31231b170c0501000000000078685a4f484446524134271c12090301"
    "040201000000000002c0b3a38f796148cbc1b5a7998a7b6b5b4d3f33261a110b000000000000000000091c3539281c11"
    "01000000000000000000072b3b2c1f15e8e1d9cfc3b39f8978806b543d28180c86756556483b3027201a14100b070301"
    "c0b5a99c8f8275685c50453a30271f1901000000000000000007101e130a04016f5d4c3d2f241c0e0b05020000000000"
    "afa1918171615243362b2219110a0602dfd8d0c7bdb2a79b8e8073675645362868564536291f170b0603000000000000"
    "88796b6057514f5d4c3e3124180d06026656483c3229231d160e010000000000c9beb2a4958576655546382d1f140b06"
    "7764513f3022170f0701000000000000584b40383637363224180f0802000000f7f4f2efebe8e3dfd8d2cbc3baafa397"
    "725f4c3a2a1d130c050100000000000090807061554c463f3b403d3125180e075e4c3f353031303c2e251b1209030000"
    "826f5c483424160a0200000000000000000000000000000000634f3d28170a02b5a8998979685848392c20150c060201"
    "675340302318130a0400000000000000efece7e2dcd6cfc6bdb5aaa093837364f6f4f1eeeae5e0dbd4cec6bdb4a99e92"
    "0000000000000000008e7a654d3621124938281a0f0806040200000000000000bbad9e8f7f716553423a2c20150b0501"
    "e8e2dcd4ccc2b8ada195887a6c5d4f41fefefdfdfcfbfbfaf8f7f6f4f2efece9"
)
STATE_P0 = bytes.fromhex(
    "1b20262d353e47505b67707e8ea0b2c5ffffffffffffffffff6d8095afcbe6fd445160708090a0b0c0cbd6e6f4fdffff"
    "434f5e6d7e8fa0b1c3d2e1f2fefffffffdffffffffffffffffd2986d6473829113171a1e23293139424b55606d7b8997"
    "1d23282e343b41495159626d77828d986474869aaec1d2e1eef9feffffffffff5e6e8093a8bbcad5e6f7ffffffffffff"
    "98a9bccedde8f0f6fafdfeffffffffff738397aabed1e1f1fcffffffffffffff52606f7f8e9dabb9c6cfd9e2e9f0f5f9"
    "ffffffffffffffffffebc9a0a6b9ccdd4f5e6e7f91a3b3c5d6e4eef6fafcfeffffffffffffffffe6c2998a7d8d9fb0c2"
    "0203040506080b0e100b0d11161d253062707e8d9aa5adb4c8dbedfeffffffff515d6872787c7a6f7f8d9cabbdd1e3f4"
    "fbfefffffffffffffd222a3441506378353e444c545d68727c8893a0aebdcbd7ffffffffffffffffffd0ab8c879aacbd"
    "fffffffffffffffffffed6978696a7b71014192028324052615669819dbad6ed63728191a0aebcc9d4dce3ecf3f9fdff"
    "404b5764717e8b98a4b0bbc6d0d9e1e7ffffffffffffffffffd4c1a9bbcfe2f35867798b9eaebdc6d9ecfaffffffffff"
    "2f3945515e6c7b8a99aab9c6d5e3eef81115191e242a3139424b555e6b7b8b9b6a788897a6b3bdccddeefdffffffffff"
    "46505a636b707166748290a0b2c6d9ea7a8a9bacbbcad6e1e9f0f5fafdffffff1d232b333d47525f6c7a8895a7b9cada"
    "5160708192a3b3c2d7f2ffffffffffff6b7881888b8a8a8fa0b2c3d4e9feffff050607090b0d0f1115181c20262c343b"
    "5564758699aabac8dbeefcffffffffff404b57626d757a818580848f9fb2c4d76a75828b91909184949eadbdcfe4f8ff"
    "4a5767798ca0b4ceebffffffffffffffffffffffffffffffff6172839ab4cee83a444f5a65717d8a97a5b1c0d1dfecf6"
    "636f8193a4b4bdcddff6ffffffffffff080a0d0f13161a1f24292f363e495460050608090b0e1013171b1f242a30373e"
    "ffffffffffffffffff424f5f738ba5c178899bafc3d4e4f3fdffffffffffffff252d37414c565f6e7e8696a5b6cadef0"
    "0d1013171c21272d353d454f5a66727f01010102020203030404050607090a0c"
)


def rows(b, dim):
    return [list(b[16 * k:16 * k + 16]) for k in range(dim)]


L_SCALE, L_DECAY, L_P0 = rows(LATENT_SCALE, 25), rows(LATENT_DECAY, 25), rows(LATENT_P0, 25)
S_SCALE, S_DECAY, S_P0 = rows(STATE_SCALE, 50), rows(STATE_DECAY, 50), rows(STATE_P0, 50)
TABLES = LATENT_SCALE + LATENT_DECAY + LATENT_P0 + STATE_SCALE + STATE_DECAY + STATE_P0


# ---- RFC 6716 section 4.1: the range decoder ----

class Dec:
    def __init__(self, buf):
        self.buf, self.offs = buf, 0
        self.nbits = 9
        self.rng = 128
        self.rem = self.byte()
        self.val = self.rng - 1 - (self.rem >> 1)
        self.norm()

    def byte(self):
        if self.offs < len(self.buf):
            self.offs += 1
            return self.buf[self.offs - 1]
        return 0

    def norm(self):
        while self.rng <= 1 << 23:
            self.nbits += 8
            self.rng <<= 8
            sym = self.rem
            self.rem = self.byte()
            sym = ((sym << 8) | self.rem) >> 1
            self.val = ((self.val << 8) + (255 & ~sym)) & ((1 << 31) - 1)

    def decode(self, ft):
        self.ext = self.rng // ft
        return ft - min(self.val // self.ext + 1, ft)

    def update(self, fl, fh, ft):
        s = self.ext * (ft - fh)
        self.val -= s
        self.rng = self.ext * (fh - fl) if fl > 0 else self.rng - s
        self.norm()

    def uint(self, ft):
        s = self.decode(ft)
        self.update(s, s + 1, ft)
        return s

    def icdf16(self, icdf):
        r = self.rng >> 15
        s, k = self.rng, 0
        while True:
            t, s = s, r * icdf[k]
            if self.val < s:
                k += 1
            else:
                break
        self.val -= s
        self.rng = t - s
        self.norm()
        return k

    def tell(self):
        return self.nbits - self.rng.bit_length()


# ---- section 3.3: the payload ----

SLOPE = [0, 2, 3, 4, 6, 8, 12, 16]
CAP = 1 << 16


def coefficient(d, p0, r):
    if r == 0 or p0 == 255:
        return 0
    s0 = 32768 - 128 * p0
    sign = d.icdf16([s0, s0 // 2, 0])
    if sign == 0:
        return 0
    icdf = [max(7, 128 * r)]
    for i in range(1, 7):
        icdf.append(max(7 - i, (icdf[i - 1] * 128 * r) >> 15))
    icdf.append(0)
    mag = 1
    while True:
        v = d.icdf16(icdf)
        mag += v
        if mag > CAP:
            raise Refused("dred_latent_range")
        if v != 7:
            break
    return -mag if sign == 2 else mag


def vector(d, q, decay, p0, dim):
    return [coefficient(d, p0[k][q], decay[k][q]) for k in range(dim)]


def parse_payload(b):
    d = Dec(b)
    q0, dq = d.uint(16), d.uint(8)
    x = d.uint(2)
    e = d.uint(256) if x == 1 else 0
    i = d.uint(32)
    qmax = 15
    if q0 < 14 and dq != 0:
        n = 14 - q0
        s = d.decode(2 * n)
        if s >= n:
            qmax = q0 + (s - n) + 1
            d.update(s, s + 1, 2 * n)
        else:
            d.update(0, n, 2 * n)
    state = vector(d, q0, S_DECAY, S_P0, 50)
    lats = []
    while 8 * len(b) - d.tell() > 7:
        q = min(qmax, q0 + (SLOPE[dq] * len(lats) + 8) // 16)
        lats.append((vector(d, q, L_DECAY, L_P0, 25), q))
    return {"q0": q0, "dq": dq, "x": x, "e": e, "i": i, "qmax": qmax, "encoded": 32 * e + i, "state": state, "latents": lats}


def symbol_bytes(p):
    b = bytes([p["q0"], p["dq"], p["x"], p["e"], p["i"], p["qmax"]]) + struct.pack("<I", len(p["latents"]))
    b += b"".join(le32(v) for v in p["state"])
    for v, q in p["latents"]:
        b += bytes([q]) + b"".join(le32(x) for x in v)
    return b


# ---- section 3.1: the packet; section 3.2: its extensions ----

def frame_ticks(config):
    if config < 12:
        return [4, 8, 16, 24][config % 4]
    if config < 16:
        return [4, 8][config % 2]
    return [1, 2, 4, 8][config % 4]


def frame_len(p, i):
    if i >= len(p):
        raise Refused("packet_malformed")
    if p[i] < 252:
        return p[i], i + 1
    if i + 1 >= len(p):
        raise Refused("packet_malformed")
    return 4 * p[i + 1] + p[i], i + 2


def framing(p):
    """(frames, frame duration, padding) of a packet under RFC 6716 section 3.4's rules."""
    bad = Refused("packet_malformed")
    if not p:
        raise bad
    config, code = p[0] >> 3, p[0] & 3
    ticks = frame_ticks(config)
    pad = b""
    if code == 0:
        lens = [len(p) - 1]
    elif code == 1:
        if (len(p) - 1) % 2:
            raise bad
        lens = [(len(p) - 1) // 2] * 2
    elif code == 2:
        n1, i = frame_len(p, 1)
        if n1 > len(p) - i:
            raise bad
        lens = [n1, len(p) - i - n1]
    else:
        if len(p) < 2:
            raise bad
        vbr, padded, m = p[1] & 0x80, p[1] & 0x40, p[1] & 0x3F
        if m == 0 or m * ticks > 48:
            raise bad
        i, plen = 2, 0
        if padded:
            while True:
                if i >= len(p):
                    raise bad
                b = p[i]
                i += 1
                if b == 255:
                    plen += 254
                else:
                    plen += b
                    break
        if i + plen > len(p):
            raise bad
        end = len(p) - plen
        pad = p[end:]
        if vbr:
            lens = []
            for _ in range(m - 1):
                n, i = frame_len(p, i)
                if i > end:
                    raise bad
                lens.append(n)
            if i > end or sum(lens) > end - i:
                raise bad
            lens.append(end - i - sum(lens))
        else:
            if i > end or (end - i) % m:
                raise bad
            lens = [(end - i) // m] * m
    if any(n > 1275 for n in lens):
        raise bad
    return len(lens), ticks, pad


def skip(p, at, idb, trailing):
    """Where one extension's payload ends and its length bytes, or None past the padding."""
    i, L = idb >> 1, idb & 1
    if (i == 0 and L == 1) or i == 2:
        return at, 0
    if 0 < i < 32:
        return (at + L, 0) if len(p) - at >= L else None
    if L == 0:
        return (len(p) - trailing, 0) if len(p) - at >= trailing else None
    j, n, hdr = at, 0, 0
    while True:
        if j >= len(p):
            return None
        lace = p[j]
        j += 1
        hdr += 1
        n += lace
        if lace != 255:
            break
    if n > len(p) - j:
        return None
    return j + n, hdr


def extensions(p, frames):
    """Every extension instance (id, frame, data) of the padding, in bitstream order."""
    out = []
    at = frame = 0
    since, last_long, trailing = 0, None, 0
    while at < len(p):
        start, idb = at, p[at]
        i, L = idb >> 1, idb & 1
        r = skip(p, at + 1, idb, 0)
        if r is None:
            return out
        at, hdr = r
        if i == 1:
            inc = 1 if L == 0 else p[start + 1]
            if inc == 0:
                continue
            frame += inc
            if frame >= frames:
                return out
            since, last_long, trailing = at, None, 0
        elif i == 2:
            src, j = [], since
            while j < start:
                end_j, _ = skip(p, j + 1, p[j], 0)
                if p[j] > 3:
                    src.append((end_j, p[j]))
                j = end_j
            for f in range(frame + 1, frames):
                for end_j, b in src:
                    if L == 0 and f + 1 >= frames and end_j == last_long:
                        b &= ~1
                    r = skip(p, at, b, trailing)
                    if r is None:
                        return out
                    out.append((b >> 1, f, p[at + r[1]:r[0]]))
                    at = r[0]
            since, last_long = at, None
            if L == 0:
                frame += 1
                if frame >= frames:
                    return out
        elif i > 2:
            if i >= 32:
                last_long, trailing = at, 0
            else:
                trailing += L
            out.append((i, frame, p[start + 1 + hdr:at]))
    return out


def find_dred(p):
    frames, ticks, pad = framing(p)
    for i, f, data in extensions(pad, frames):
        if i == 32 and len(data) > 0:
            return data, f * ticks
        if i == 126 and len(data) > 2 and data[:2] == b"D\x0c":
            return data[2:], f * ticks
    return None


# ---- section 4.1 and 4.4: the weights, as integers ----

WEIGHTS_SHA256 = "394f0ecf19926cebd678019dc94a54c88ae078340c16325069707f286cdf28a6"
LAYERS = [("dense1", 26, 96, False), ("hidden_init", 50, 128, False), ("gru_init", 128, 320, True)]
for _k, _n in zip(range(1, 6), [96, 192, 288, 384, 480]):
    LAYERS += [("gru%d_input" % _k, _n, 192, True), ("gru%d_recurrent" % _k, 64, 192, False)]
LAYERS += [("glu%d" % _k, 64, 64, False) for _k in range(1, 6)]
LAYERS += [("conv_dense%d" % _k, _n, 32, True) for _k, _n in zip(range(1, 6), [160, 256, 352, 448, 544])]
LAYERS += [("conv%d" % _k, 64, 32, False) for _k in range(1, 6)]
LAYERS += [("output", 576, 80, True)]


def quantize(bits, f):
    """R_half_even of the float32's value times 2^f, from its bits (section 4.4)."""
    e, m = (bits >> 23) & 0xFF, bits & 0x7FFFFF
    if e == 255:
        raise Refused("model_mismatch")
    mant, exp = (m, -149) if e == 0 else (m | 0x800000, e - 150)
    k = exp + f
    if k >= 0:
        mag = mant << k
    elif -k >= 63:
        mag = 0
    else:
        s = -k
        q, r, h = mant >> s, mant & ((1 << s) - 1), 1 << (s - 1)
        mag = q + 1 if (r > h or (r == h and q & 1)) else q
    return -mag if bits >> 31 else mag


def records(blob):
    out, at = {}, 0
    while at < len(blob):
        if len(blob) - at < 64 or blob[at:at + 4] != b"DNNw":
            raise Refused("model_mismatch")
        version, kind, size, block = struct.unpack_from("<iiii", blob, at + 4)
        name = blob[at + 20:at + 64].split(b"\0")[0].decode("ascii")
        if version != 0 or not 0 <= size <= block or block % 64 or block - size >= 64 or name in out:
            raise Refused("model_mismatch")
        data = blob[at + 64:at + 64 + size]
        if kind in (0, 1):
            out[name] = list(struct.unpack("<%d%s" % (size // 4, "I" if kind == 0 else "i"), data))
        else:
            out[name] = data
        at += 64 + block
    return out


def layers(blob):
    """Each layer's float32 bit patterns: weights [out][in] and bias."""
    arr = records(blob)
    out = {}
    for name, nin, nout, sparse in LAYERS:
        raw, bias = arr["dec_%s_weights_float" % name], arr["dec_%s_bias" % name]
        w = [[0] * nin for _ in range(nout)]
        if sparse:
            idx, p, wp = arr["dec_%s_weights_idx" % name], 0, 0
            for og in range(0, nout, 8):
                n = idx[p]
                p += 1
                for _ in range(n):
                    pos = idx[p]
                    p += 1
                    if pos % 4 or pos >= nin:
                        raise Refused("model_mismatch")
                    for i in range(4):
                        for o in range(8):
                            w[og + o][pos + i] = raw[wp + 8 * i + o]
                    wp += 32
        else:
            for i in range(nin):
                for o in range(nout):
                    w[o][i] = raw[i * nout + o]
        out[name] = (w, bias)
    return out


# ---- section 4.5: the activation tables ----

C0 = 4609434768280048980


def exp_q62(u):
    c, x = C0, 1 << 62
    for k in range(15):
        if (u >> k) & 1:
            x = (x * c) >> 62
        c = (c * c) >> 62
    return x


def tables():
    tanh, sig = [0] * 65536, [0] * 65536
    for v in range(-32767, 32768):
        x = exp_q62(abs(v))
        b = (1 << 62) + x
        a = (1 << 15) * ((1 << 62) - x)
        t = min(32767, (2 * a + b) // (2 * b))
        tanh[v + 32768] = -t if v < 0 else t
        a = 1 << 77 if v >= 0 else (1 << 15) * x
        sig[v + 32768] = min(32767, (2 * a + b) // (2 * b))
    return tanh, sig


def table_digest(t):
    return hashlib.sha256(b"".join(struct.pack("<h", v) for v in t)).hexdigest()


# ---- sections 4.3, 4.6 and 4.7: the network ----

def rhu(p, s):
    return (p + (1 << (s - 1))) >> s


def clamp(v, lo=-32767, hi=32767):
    return lo if v < lo else hi if v > hi else v


class Net:
    def __init__(self, blob):
        L = layers(blob)
        self.tanh, self.sig = tables()

        def q(name, bias_frac):
            w, b = L[name]
            return [[quantize(x, 24) for x in row] for row in w], [quantize(x, bias_frac) for x in b]

        self.l = {}
        for name, _, _, _ in LAYERS:
            if name.startswith("gru") and name != "gru_init":
                continue
            self.l[name] = q(name, 40 if name in ("dense1", "hidden_init") else 39)
        for k in range(1, 6):
            wi, bi = q("gru%d_input" % k, 39)
            wr, br = q("gru%d_recurrent" % k, 39)
            self.l["gates%d" % k] = ([wi[o] + wr[o] for o in range(128)], [bi[o] + br[o] for o in range(128)])
            self.l["hin%d" % k] = (wi[128:], bi[128:])
            self.l["hrec%d" % k] = (wr[128:], br[128:])
        for name, (w, b) in self.l.items():
            if any(not -(1 << 31) <= x < 1 << 31 for row in w for x in row):
                raise Refused("model_mismatch")

    def dense(self, x, name, shift, act=None):
        w, b = self.l[name]
        y = [rhu(sum(map(mul, row, x)) + bo, shift) for row, bo in zip(w, b)]
        if act is not None:
            y = [act[clamp(v) + 32768] for v in y]
        return y

    def init(self, state):
        a = self.dense(state, "hidden_init", 28, self.tanh)
        s = self.dense(a, "gru_init", 27, self.tanh)
        return [s[64 * k:64 * k + 64] for k in range(5)]

    def step(self, z, hs, cs):
        t = self.dense(z, "dense1", 28, self.tanh)
        nh, nc = [], []
        for k in range(1, 6):
            h = hs[k - 1]
            g = self.dense(t + h, "gates%d" % k, 28, self.sig)
            zg, r = g[:64], g[64:]
            hin = self.dense(t, "hin%d" % k, 19)
            hrec = self.dense(h, "hrec%d" % k, 19)
            hc = [self.tanh[clamp(rhu(a + rhu(rr * b, 15), 8)) + 32768] for a, rr, b in zip(hin, r, hrec)]
            h2 = [clamp(c + rhu(zz * (hh - c), 15)) for c, zz, hh in zip(hc, zg, h)]
            gl = self.dense(h2, "glu%d" % k, 28, self.sig)
            t = t + [rhu(a * b, 15) for a, b in zip(h2, gl)]
            c = self.dense(t, "conv_dense%d" % k, 27, self.tanh)
            v = self.dense(cs[k - 1] + c, "conv%d" % k, 27, self.tanh)
            t = t + v
            nh.append(h2)
            nc.append(c)
        x = [clamp(v, -(1 << 31), (1 << 31) - 1) for v in self.dense(t, "output", 23)]
        return x, nh, nc

    def decode(self, p, n):
        """Rows 0 .. 4n-1 of a payload, newest first (section 4.7)."""
        def deq(x, s):
            s = max(1, s)
            return clamp((2 * x * (1 << 24) + s) // (2 * s), -(1 << 31) + 1, (1 << 31) - 1)
        q0 = p["q0"]
        hs = self.init([deq(x, S_SCALE[k][q0]) for k, x in enumerate(p["state"])])
        cs = [[0] * 32 for _ in range(5)]
        out = []
        for v, q in p["latents"][:n]:
            z = [deq(x, L_SCALE[k][q]) for k, x in enumerate(v)] + [(q - 8) << 13]
            x, hs, cs = self.step(z, hs, cs)
            out += [x[20 * j:20 * j + 20] for j in range(4)]
        return out


def row_bytes(rs):
    return b"".join(le32(v) for r in rs for v in r)


# ---- section 5: a stream under loss ----

def lost(seed, enter, stay, n):
    out, prev = [], False
    for i in range(n):
        u = int.from_bytes(sha(b"wai:dred-loss\x01" + seed + i.to_bytes(4, "big"))[:2], "big")
        prev = u < (stay if prev else enter)
        out.append(prev)
    return out


def dred_of(p):
    try:
        found = find_dred(p)
        if found is None:
            return None
        pl = parse_payload(found[0])
    except Refused:
        return None
    return pl, 16 - pl["encoded"] + found[1]


def slots(packets, d):
    out, m = [], 0
    while m < len(packets):
        if packets[m] is not None:
            m += 1
            continue
        a = m
        while m < len(packets) and packets[m] is None:
            m += 1
        rec = None
        if m < len(packets):
            got = dred_of(packets[m])
            if got is not None:
                rec = (m, len(got[0]["latents"]), m * d + got[1])
        for f in range(a, m):
            for h in (0, 1):
                e = f * d + 4 * (h + 1)
                src = None
                if rec is not None and rec[2] - e >= 0 and (rec[2] - e) // 4 < 4 * rec[1]:
                    src = (rec[0], (rec[2] - e) // 4)
                out.append((f, h, src))
    return out


class Rebuild:
    """Rows of each recovering packet, decoded through as many latents as asked (cached)."""

    def __init__(self, net):
        self.net, self.cache = net, {}

    def rows(self, packet, n):
        key = (bytes(packet), n)
        if key not in self.cache:
            self.cache[key] = self.net.decode(dred_of(packet)[0], n)
        return self.cache[key]


def buffer(ss, packets, rb):
    need = {}
    for _, _, src in ss:
        if src:
            need[src[0]] = max(need.get(src[0], 0), src[1] // 4 + 1)
    b = b""
    for f, h, src in ss:
        b += struct.pack("<IB", f, h)
        if src is None:
            b += b"\0"
        else:
            b += b"\1" + struct.pack("<II", src[0], src[1]) + b"".join(le32(v) for v in rb.rows(packets[src[0]], need[src[0]])[src[1]])
    return b


def runs(ss):
    out, i = [], 0
    while i < len(ss):
        j = i + 1
        while j < len(ss) and ss[j][0] <= ss[j - 1][0] + 1:
            j += 1
        out.append((ss[j - 1][0] + 1, i, j))
        i = j
    return out


def envelope_payload(e):
    if e[:4] != b"WAI1":
        raise Refused("record_malformed")
    ml = struct.unpack_from("<I", e, 4)[0]
    pl = struct.unpack_from("<I", e, 8 + ml)[0]
    body = e[12 + ml:]
    if len(body) != pl:
        raise Refused("record_malformed")
    return body


# ---- section 7.2: the player's log (synthesis section 10.3) ----

SIGN_DOMAIN = b"wai:synthesis\x01"
LINK_DOMAIN = b"wai:synthesis-link\x01"


def jstr(s):
    out = ['"']
    for c in s:
        if c == '"':
            out.append('\\"')
        elif c == "\\":
            out.append("\\\\")
        elif ord(c) < 0x20:
            out.append("\\u%04x" % ord(c))
        else:
            out.append(c)
    return "".join(out) + '"'


def jcs(v):
    if v is None:
        return "null"
    if isinstance(v, bool):
        raise Refused("record_malformed")
    if isinstance(v, int):
        return str(v)
    if isinstance(v, str):
        return jstr(v)
    if isinstance(v, list):
        return "[" + ",".join(jcs(x) for x in v) + "]"
    return "{" + ",".join(jstr(k) + ":" + jcs(v[k]) for k in sorted(v, key=lambda k: k.encode("utf-16-be"))) + "}"


def hexn(v, n):
    if not isinstance(v, str) or len(v) != 2 * n or not re.fullmatch("[0-9a-f]*", v):
        raise Refused("record_malformed")
    return bytes.fromhex(v)


PRESENT = ["type", "subject", "origin", "class", "media", "tier", "decoder", "variant", "output"]


def read_entry(text):
    def strict(pairs):
        d = {}
        for k, v in pairs:
            if k in d:
                raise Refused("record_malformed")
            d[k] = v
        return d
    try:
        v = json.loads(text, object_pairs_hook=strict, parse_float=lambda s: (_ for _ in ()).throw(Refused("record_malformed")))
    except ValueError:
        raise Refused("record_malformed")
    if not isinstance(v, dict) or sorted(v) != sorted(["kind", "session", "seq", "prev", "at", "entry", "signer_pubkey", "sig"]) or v["kind"] != "wai-player":
        raise Refused("record_malformed")
    e = v["entry"]
    if not isinstance(e, dict) or (e.get("type") == "present" and sorted(e) != sorted(PRESENT)):
        raise Refused("record_malformed")
    for k in ("seq", "at"):
        if not isinstance(v[k], int) or isinstance(v[k], bool) or not 0 <= v[k] < 2 ** 53:
            raise Refused("record_malformed")
    pk, sig = hexn(v["signer_pubkey"], 32), hexn(v["sig"], 64)
    hexn(v["session"], 32)
    if v["prev"] is not None:
        hexn(v["prev"], 32)
    if e["type"] == "present":
        hexn(e["subject"], 32)
        if e["output"] is not None:
            hexn(e["output"], 32)
    msg = SIGN_DOMAIN + jcs({k: x for k, x in v.items() if k != "sig"}).encode("utf-8")
    bad = ED.verify(pk, msg, sig)
    if bad:
        raise Refused(bad)
    return v, blake3(LINK_DOMAIN + jcs(v).encode("utf-8"))


def content_hash(e):
    return blake3(b"moq-jwp:object\x01" + e)


# ---- the checks of sections 5.5 and 7.2 ----

def check_case(case, stream, net, rb):
    rec = case["record"]
    if rec["weights_sha256"] != WEIGHTS_SHA256:
        raise Refused("weights_digest_mismatch", ["record"])
    objects = [None if o is None else bytes.fromhex(o) for o in stream["objects"]]
    loss = stream["loss"]
    pat = lost(bytes.fromhex(loss["seed"]), loss["enter"], loss["stay"], len(objects))
    want_lost = [i for i, x in enumerate(pat) if x]
    if rec["lost"] != want_lost or [i for i, o in enumerate(objects) if o is None] != want_lost:
        raise Refused("loss_pattern_mismatch", ["record"])
    packets = [None if o is None else envelope_payload(o) for o in objects]
    ss = slots(packets, stream["frame_ticks"])
    stated = [(s[0], s[1], None if s[2] is None else (s[2], s[3])) for s in rec["slots"]]
    if stated != ss:
        raise Refused("reconstruction_mismatch", ["record"])
    if sha(buffer(ss, packets, rb)).hex() != rec["output"]:
        raise Refused("reconstruction_mismatch", ["record"])
    rs = runs(ss)
    prev = None
    for i, text in enumerate(case["player"]):
        try:
            v, link = read_entry(text)
        except Refused as r:
            raise Refused(r.code, ["player", i])
        if prev is None:
            ok = v["seq"] == 0 and v["prev"] is None
        else:
            ok = (v["seq"] == prev[0]["seq"] + 1 and v["prev"] == prev[1].hex() and v["session"] == prev[0]["session"]
                  and v["signer_pubkey"] == prev[0]["signer_pubkey"] and v["at"] >= prev[0]["at"])
        if not ok:
            raise Refused("player_chain_broken", ["player", i])
        prev = (v, link)
        e = v["entry"]
        if e["type"] != "present":
            continue
        good = e["decoder"] == CAPABILITY and e["variant"] == LOWERING and e["tier"] == "decode-equivalence"
        good = good and any(
            n < len(objects) and objects[n] is not None and content_hash(objects[n]).hex() == e["subject"]
            and e["output"] == sha(buffer(ss[a:b], packets, rb)).hex()
            for n, a, b in rs
        )
        if not good:
            raise Refused("present_output_irreproducible", ["player", i])


def check_payload(v, net):
    p = bytes.fromhex(v["packet"])
    try:
        found = find_dred(p)
        if found is None:
            return "none", None
        pl = parse_payload(found[0])
    except Refused as r:
        return {"refuse": r.code}, None
    rs = net.decode(pl, len(pl["latents"]))
    got = {
        "frame_offset": found[1], "q0": pl["q0"], "dq": pl["dq"], "extended": pl["e"] if pl["x"] == 1 else None,
        "indicator": pl["i"], "qmax": pl["qmax"], "dred_offset": 16 - pl["encoded"] + found[1],
        "latents": len(pl["latents"]), "symbols_sha256": sha(symbol_bytes(pl)).hex(), "rows": len(rs),
        "features_sha256": sha(row_bytes(rs)).hex(), "newest": rs[0],
    }
    return "dred", got


def outcome_of(doc_case, stream, net, rb):
    try:
        check_case(doc_case, stream, net, rb)
        return "accept"
    except Refused as r:
        return {"refuse": r.code, "at": r.at}


# ---- the extension's text ----

def section(text, n):
    m = re.search(r"\n## %d\. [^\n]*\n(.*?)(?=\n## |\Z)" % n, text, re.S)
    if not m:
        raise SystemExit("extensions/opus-dred.md has no section %d" % n)
    return m.group(1)


def listed_codes(text):
    return re.findall(r"^\| `([a-z_]+)` \|", section(text, 8), re.M)


def verify(doc, net, rb, text, tanh, sig):
    """Every statement of the corpus, or the first that fails, as a string."""
    s4 = section(text, 4)
    for want in (WEIGHTS_SHA256, hashlib.sha256(TABLES).hexdigest(), table_digest(tanh), table_digest(sig), str(C0)):
        if want not in s4:
            return "section 4 does not state %s" % want
    if doc["weights"]["sha256"] != WEIGHTS_SHA256 or doc["tables_sha256"] != hashlib.sha256(TABLES).hexdigest():
        return "the corpus names other weights or tables"
    lw = doc["lowering"]
    if lw["exp_c0"] != str(C0) or lw["tanh_table_sha256"] != table_digest(tanh) or lw["sigmoid_table_sha256"] != table_digest(sig):
        return "the corpus states other tables"
    for k in ("init_sha256", "step_sha256"):
        if lw[k] not in s4:
            return "section 4.6 does not state the corpus's %s" % k
    given = set()
    for v in doc["payloads"]:
        kind, got = check_payload(v, net)
        if kind != v["outcome"]:
            return "payload %s: %r, not %r" % (v["name"], kind, v["outcome"])
        if isinstance(kind, dict):
            given.add(kind["refuse"])
        if got is not None and got != v["dred"]:
            return "payload %s: states %r, decodes to %r" % (v["name"], v["dred"], got)
    for c in doc["cases"]:
        got = outcome_of(c, doc["streams"][c["stream"]], net, rb)
        if got != c["outcome"]:
            return "case %s: %r, not %r" % (c["name"], got, c["outcome"])
        if isinstance(got, dict):
            given.add(got["refuse"])
    missing = [c for c in listed_codes(text) if c not in given]
    if missing:
        return "no vector refuses %s" % ", ".join(missing)
    return None


def mutations(doc):
    """Copies of the corpus with one defect each, every one of which must fail."""
    def flip(h, at):
        return h[:at] + ("0" if h[at] != "0" else "1") + h[at + 1:]
    out = []
    pv = [i for i, v in enumerate(doc["payloads"]) if v["outcome"] == "dred"]

    def m(label, f):
        d = copy.deepcopy(doc)
        f(d)
        out.append((label, d))
    m("a payload's packet byte", lambda d: d["payloads"][pv[0]].__setitem__("packet", flip(d["payloads"][pv[0]]["packet"], 60)))
    m("a payload's row digest", lambda d: d["payloads"][pv[1]]["dred"].__setitem__("features_sha256", flip(d["payloads"][pv[1]]["dred"]["features_sha256"], 0)))
    m("a payload's newest row", lambda d: d["payloads"][pv[2]]["dred"]["newest"].__setitem__(3, d["payloads"][pv[2]]["dred"]["newest"][3] + 1))
    m("a payload's symbol digest", lambda d: d["payloads"][pv[0]]["dred"].__setitem__("symbols_sha256", flip(d["payloads"][pv[0]]["dred"]["symbols_sha256"], 5)))
    first_refused = next(i for i, v in enumerate(doc["payloads"]) if isinstance(v["outcome"], dict))
    m("a refused payload stated as none", lambda d: d["payloads"][first_refused].__setitem__("outcome", "none"))
    m("the tables' digest", lambda d: d.__setitem__("tables_sha256", flip(d["tables_sha256"], 0)))
    m("the tanh table's digest", lambda d: d["lowering"].__setitem__("tanh_table_sha256", flip(d["lowering"]["tanh_table_sha256"], 9)))
    m("a step graph's identity", lambda d: d["lowering"].__setitem__("step_sha256", flip(d["lowering"]["step_sha256"], 9)))
    acc = [i for i, c in enumerate(doc["cases"]) if c["outcome"] == "accept"]
    m("an accepted record's output", lambda d: d["cases"][acc[0]]["record"].__setitem__("output", flip(d["cases"][acc[0]]["record"]["output"], 3)))
    m("an accepted record's lost list", lambda d: d["cases"][acc[1]]["record"]["lost"].pop())
    m("an accepted record's slot", lambda d: d["cases"][acc[0]]["record"]["slots"].__setitem__(0, [d["cases"][acc[0]]["record"]["slots"][0][0], 1 - d["cases"][acc[0]]["record"]["slots"][0][1]] + d["cases"][acc[0]]["record"]["slots"][0][2:]))
    m("a stream's loss seed", lambda d: d["streams"][doc["cases"][acc[0]]["stream"]]["loss"].__setitem__("seed", flip(d["streams"][doc["cases"][acc[0]]["stream"]]["loss"]["seed"], 2)))

    def recovering_object(d):
        objs = d["streams"][doc["cases"][acc[0]]["stream"]]["objects"]
        subject = json.loads(doc["cases"][acc[0]]["player"][0])["entry"]["subject"]
        i = next(i for i, o in enumerate(objs) if o is not None and content_hash(bytes.fromhex(o)).hex() == subject)
        objs[i] = flip(objs[i], len(objs[i]) - 9)
    m("a recovering packet's last bytes", recovering_object)

    def player_output(d):
        e = json.loads(d["cases"][acc[0]]["player"][0])
        e["entry"]["output"] = flip(e["entry"]["output"], 0)
        d["cases"][acc[0]]["player"][0] = json.dumps(e, separators=(",", ":"))
    m("an accepted log's output", player_output)
    m("a refusal's code", lambda d: d["cases"][-2]["outcome"].__setitem__("refuse", "reconstruction_mismatch"))
    first_record = next(i for i, c in enumerate(doc["cases"]) if c["outcome"] != "accept" and c["outcome"]["at"] == ["record"])
    m("a refusal's place", lambda d: d["cases"][first_record]["outcome"].__setitem__("at", ["player", 0]))
    return out


def main():
    if len(sys.argv) != 3:
        raise SystemExit("usage: python3 -I dred_verify.py <live-conformance/dred> <weights-file>")
    blob = open(sys.argv[2], "rb").read()
    if hashlib.sha256(blob).hexdigest() != WEIGHTS_SHA256:
        raise SystemExit("weights_digest_mismatch: %s is not the draft's weights file" % sys.argv[2])
    doc = json.load(open(os.path.join(sys.argv[1], "vectors.json")))
    text = open(EXTENSION).read()
    if listed_codes(text) != ["packet_malformed", "dred_latent_range", "weights_digest_mismatch", "loss_pattern_mismatch", "reconstruction_mismatch"]:
        raise SystemExit("section 8 lists %r" % listed_codes(text))
    net = Net(blob)
    rb = Rebuild(net)
    err = verify(doc, net, rb, text, net.tanh, net.sig)
    if err:
        print("FAIL: " + err)
        return 1
    print("dred: %d payload vectors and %d cases reach their outcomes; every reason of section 8 refuses one" % (len(doc["payloads"]), len(doc["cases"])))
    caught = 0
    for label, d in mutations(doc):
        if verify(d, net, rb, text, net.tanh, net.sig) is None:
            print("FAIL: the mutation '%s' was accepted" % label)
            return 1
        caught += 1
    print("dred: every mutated copy (%d) fails" % caught)
    return 0


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