Dragon-s-Lair-X68k/tools/encoder/dlx.py

#!/usr/bin/env python3
"""Reference DLX reader/decoder -- the ground truth the 68000 player is checked against.

This is deliberately a *decoder*, not a re-run of the encoder: it parses the
container byte-for-byte the way `src/player/` must, so that any disagreement
between it and the 68000 is a decoder bug rather than two encoders differing.
`tools/analysis/09_ratectl_drift.py` already gates the encoder against its own
reconstruction; this gates the container against the player.

Everything is big-endian (see the `encode.py` docstring).  Block raster order,
2-bit modes packed MSB-first: 00=SKIP 01=V1 10=V4 11=RAW.

V4 sub-block order is (sub_y, sub_x) row-major -- TL, TR, BL, BR -- matching
`vq_hybrid.paint`'s reshape(-1,2,2,2,2).transpose(0,1,3,2,4).

DLX3 adds the v7 LITERAL SPAN section between the mode header and the block
payload (FINDINGS 40, tools/encoder/spans.py).  A spanned block reads SKIP in
the mode header and is painted by a span instead, so a reader that ignores the
section does not merely lose the spans -- it displays stale pixels wherever one
was.  The section is at a KNOWN offset (header end) rather than behind the
block payload precisely so that the 68000 can paint it before it has parsed
anything of variable length.
"""
import os, sys, struct
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import spans as SP

MODE_SKIP, MODE_V1, MODE_V4, MODE_RAW = 0, 1, 2, 3


class DLX:
    def __init__(self, path):
        self.raw = open(path, "rb").read()
        b = self.raw
        # DLX2 pads every frame record up to a 4-byte boundary; DLX1 lays them
        # end to end.  On a 68000 that is not a slow read but an ADDRESS ERROR
        # (FINDINGS 28.3), so the padding is part of the format, not a loader
        # convenience -- but DLX1 containers stay readable, because every
        # measurement in FINDINGS 28-31 was taken on one.
        if b[:4] not in (b"DLX1", b"DLX2", b"DLX3", b"DLX4"):
            raise ValueError(f"{path}: not a DLX container")
        self.version = int(b[3:4])
        self.aligned = self.version >= 2
        self.has_spans = self.version >= 3
        self.has_index = self.version >= 4
        (self.W, self.H, self.fps, self.nframes,
         self.k1, self.k4) = struct.unpack(">HHHHHH", b[4:16])
        off_pal, off_cb1, off_cb4, off_frm = struct.unpack(">IIII", b[16:32])
        off_idx = struct.unpack(">I", b[32:36])[0] if self.has_index else None

        self.pal = np.frombuffer(b, np.uint8, 256 * 3, off_pal).reshape(256, 3)
        self.cb1 = np.frombuffer(b, np.uint8, self.k1 * 16,
                                 off_cb1).reshape(self.k1, 4, 4)
        self.cb4 = np.frombuffer(b, np.uint8, self.k4 * 4,
                                 off_cb4).reshape(self.k4, 2, 2)

        self.idx_bytes = 1 if max(self.k1, self.k4) <= 256 else 2
        self.nbx, self.nby = self.W // 4, self.H // 4
        self.nb = self.nbx * self.nby
        self.mode_bytes = (self.nb * 2 + 7) // 8

        # frame directory: (offset of the mode header, payload length)
        if self.aligned and off_frm % 4:
            raise ValueError(f"{path}: DLX2 frame stream starts at {off_frm}, "
                             f"which is not 4-byte aligned")
        self.frames = []
        p = off_frm
        for _ in range(self.nframes):
            (n,) = struct.unpack(">I", b[p:p + 4])
            self.frames.append((p + 4, n))
            p += 4 + n
            if self.aligned:
                p += -p % 4          # skip the pad to the next record
        # The writer does not pad after the LAST record -- nothing follows it --
        # so `p` may have advanced past the end by up to 3 bytes there.
        slack = len(b) - p
        if not (slack == 0 or (self.aligned and -3 <= slack < 0)):
            raise ValueError(f"{path}: {slack} trailing bytes after "
                             f"{self.nframes} frames")

        # DLX4's record index, and it is CHECKED rather than trusted.  The
        # walk above is what every reader in this tree did before there was an
        # index -- read a record's length word to find the next one -- and it is
        # exactly what a player streaming off a disc cannot do, because the
        # length word of record i+1 is one of the bytes it has not fetched.  So
        # the two are computed independently here and required to agree: the
        # index is the producer's only source of record geometry, and an index
        # that disagrees with the stream places records at wrong addresses,
        # which the block loop reads without a bounds check (49.2).
        self.index = None
        if self.has_index:
            self.index = list(struct.unpack(
                f">{self.nframes}H", b[off_idx:off_idx + 2 * self.nframes]))
            walked = [(-(4 + n) % 4 + 4 + n) // 4 for _, n in self.frames]
            if self.index != walked:
                bad = next(i for i in range(self.nframes)
                           if self.index[i] != walked[i])
                raise ValueError(
                    f"{path}: record index disagrees with the frame stream at "
                    f"frame {bad}: index says {self.index[bad]} longwords, the "
                    f"stream is {walked[bad]}")
            if off_frm + 4 * sum(self.index) != len(b):
                raise ValueError(
                    f"{path}: the index accounts for "
                    f"{off_frm + 4 * sum(self.index)} bytes and the file is "
                    f"{len(b)} -- a producer trusting it would run off the end")

    def modes(self, f):
        o, _ = self.frames[f]
        h = np.frombuffer(self.raw, np.uint8, self.mode_bytes, o)
        m = np.stack([(h >> 6) & 3, (h >> 4) & 3, (h >> 2) & 3, h & 3], axis=1)
        return m.reshape(-1)[:self.nb].copy()

    def spans(self, f):
        """The frame's literal spans: (list of (y, x, pixels), payload offset).

        Layout, big-endian, at the end of the mode header (DLX3 only):
            u16 nspans
            nspans * { u32 GVRAM address, u16 coarse disp, c*48 B pixels,
                       u16 fine disp,     f*4 B pixels }
        The displacements are JUMP offsets into the decoder's two unrolled copy
        chains, so the pixel counts are read back out of them -- which is the
        strongest available check that the encoder and `blit.s` agree about the
        chain geometry, because a wrong displacement lands mid-chain and paints
        the wrong number of pixels rather than failing loudly.
        """
        o, n = self.frames[f]
        p = o + self.mode_bytes
        if not self.has_spans:
            return [], p
        b = self.raw
        (ns,) = struct.unpack(">H", b[p:p + 2])
        p += 2
        out = []
        for _ in range(ns):
            addr, cd = struct.unpack(">IH", b[p:p + 6])
            p += 6
            c = SP.COARSE_N - cd // SP.COARSE_CODE
            if cd % SP.COARSE_CODE or not 0 <= c <= SP.COARSE_N:
                raise ValueError(f"frame {f}: coarse displacement {cd} is not "
                                 f"an entry point in an {SP.COARSE_N}-unit chain")
            px = list(np.frombuffer(b, ">u2", c * SP.COARSE_PX, p) & 0xFF)
            p += c * SP.COARSE_PX * 2
            (fd,) = struct.unpack(">H", b[p:p + 2])
            p += 2
            fu = SP.FINE_N - fd // SP.FINE_CODE
            if fd % SP.FINE_CODE or not 0 <= fu <= SP.FINE_N:
                raise ValueError(f"frame {f}: fine displacement {fd} is not "
                                 f"an entry point in a {SP.FINE_N}-unit chain")
            px += list(np.frombuffer(b, ">u2", fu * SP.FINE_PX, p) & 0xFF)
            p += fu * SP.FINE_PX * 2
            a = addr - SP.GVRAM
            y, x = divmod(a, SP.STRIDE)
            y -= SP.YOFF
            if x % 2 or not (0 <= y < self.H) or not (0 <= x // 2 < self.W):
                raise ValueError(f"frame {f}: span destination {addr:#x} is "
                                 f"not a pixel of the {self.W}x{self.H} picture")
            # blit.s tolerates a span running past the visible 256 pixels (the
            # line stride is 1024 bytes and only the first 512 are displayed),
            # but nothing an encoder emits should need to: a span is a run of
            # whole blocks.  numpy would truncate it here in silence.
            if x // 2 + len(px) > self.W:
                raise ValueError(f"frame {f}: span at ({x//2},{y}) of "
                                 f"{len(px)} px overruns the picture width")
            out.append((y, x // 2, np.array(px, np.uint8)))
        return out, p

    def blocks(self, f):
        """Decoded 4x4 palette-index blocks for the non-SKIP blocks of frame f.

        Returns (mode, dict{block index -> (4,4) uint8}).  SKIP blocks are
        absent by construction -- the player must leave those pixels alone,
        and a decoder that materialises them is hiding the very bug this
        module exists to catch.
        """
        mode = self.modes(f)
        o, n = self.frames[f]
        _, p = self.spans(f)
        end = o + n
        ib, out = self.idx_bytes, {}
        b = self.raw
        for i, mo in enumerate(mode):
            if mo == MODE_SKIP:
                continue
            if mo == MODE_V1:
                v = b[p] if ib == 1 else (b[p] << 8) | b[p + 1]
                p += ib
                out[i] = self.cb1[v]
            elif mo == MODE_V4:
                sub = []
                for _ in range(4):
                    v = b[p] if ib == 1 else (b[p] << 8) | b[p + 1]
                    p += ib
                    sub.append(self.cb4[v])
                blk = np.empty((4, 4), np.uint8)
                blk[0:2, 0:2], blk[0:2, 2:4] = sub[0], sub[1]
                blk[2:4, 0:2], blk[2:4, 2:4] = sub[2], sub[3]
                out[i] = blk
            else:
                out[i] = np.frombuffer(b, np.uint8, 16, p).reshape(4, 4)
                p += 16
        if p != end:
            raise ValueError(f"frame {f}: payload consumed {p - o} of {n} bytes")
        return mode, out

    def paint(self, canvas, f):
        """Apply frame f in place to a (H,W) index canvas.  SKIP = untouched."""
        mode, blks = self.blocks(f)
        for i, blk in blks.items():
            by, bx = divmod(i, self.nbx)
            canvas[by * 4:by * 4 + 4, bx * 4:bx * 4 + 4] = blk
        sp, _ = self.spans(f)
        for y, x, pix in sp:
            # A span paints blocks the mode header calls SKIP.  If it ever
            # overlaps a coded block the two disagree about the same pixels and
            # the 68000's answer depends on which it does last -- so this is a
            # format invariant, not a courtesy check.
            b0, b1 = x // 4, -(-(x + len(pix)) // 4)
            bad = [b for b in range(b0, b1)
                   if mode[(y // 4) * self.nbx + b] != MODE_SKIP]
            if bad:
                raise ValueError(f"frame {f}: span at ({x},{y}) covers "
                                 f"non-SKIP block(s) {bad} of block row {y//4}")
            canvas[y, x:x + len(pix)] = pix
        return mode

    def decode_all(self):
        """The true reconstruction sequence: what any correct player displays."""
        c = np.zeros((self.H, self.W), np.uint8)
        out = []
        for f in range(self.nframes):
            self.paint(c, f)
            out.append(c.copy())
        return out