boxsim-fileformat-rev/boxsim/a2f.py

"""Pure Python implementation of Boxsim's A2F codec (decompression).

Reverse engineered from `ExpandA2F` at VA 0xa11880 in Boxsim 2.1. No part of
this needs Boxsim.exe; `boxsim._oracle` runs the original machine code and is
used only to differentially test this module.

Stream shape
------------
A bit stream, MSB first, read through a 31 bit accumulator. Symbols come from a
canonical Huffman code over 356 symbols decoded through a 65536 entry direct
lookup table:

    0x000..0x0FF  literal byte
    0x100         end of stream
    0x101         new code lengths follow, rebuild the tables
    0x102         define N entries of the 64 entry digram dictionary
    0x103..0x10A  match, length 3, distance in k+3 bits
    0x10B..0x114  match, length 4, distance in k+3 bits
    0x115..0x123  match, unary coded length class, distance in k+3 bits
    0x124..0x163  emit the two bytes of digram dictionary entry k

Distances use k = symbol - class base: k+3 raw bits biased by 2**(k+3) - 7, so
the classes tile 1.. without gaps.
"""

from __future__ import annotations

NSYM = 0x164  # 356 symbols
TABLE_BITS = 16
TABLE_SIZE = 1 << TABLE_BITS
DICT_SIZE = 0x40
INITIAL_CODE_LENGTH = 9
MAX_CODE_LENGTH = 16

SYM_EOS = 0x100
SYM_REBUILD = 0x101
SYM_DIGRAM_DEF = 0x102
SYM_MATCH3 = 0x103  # 8 symbols
SYM_MATCH4 = 0x10B  # 10 symbols
SYM_MATCHN = 0x115  # 15 symbols
SYM_DIGRAM = 0x124  # 64 symbols

# Unary length class -> (extra bits, base). Class 0 is unreachable because the
# unary read always consumes at least one bit.
LENGTH_CLASSES = {
    1: (2, 5),
    2: (4, 9),
    3: (6, 0x19),
    4: (8, 0x59),
    5: (10, 0x159),
    6: (12, 0x559),
}


class A2FError(ValueError):
    """Raised when the stream is malformed or uses an unreachable encoding."""


#: Bytes of zero padding tolerated past the end of the stream. Boxsim decodes
#: from a 0x4000 byte staging buffer and its refill reads on past the valid
#: length, so the final symbol's 16 bit lookup window routinely extends beyond
#: the last real byte. Those bits are never significant: the end-of-stream code
#: is fully contained in real bytes and the lookup table maps every window
#: sharing that prefix to the same symbol. The allowance stays small so a
#: genuinely truncated file still fails loudly instead of decoding garbage.
PAD_ALLOWANCE = 8


class _BitReader:
    """MSB-first bit reader with Boxsim's 23 bit accumulator mask.

    The mask is load bearing: Boxsim shifts `(acc & 0x7FFFFF) << 8` on every
    refill, which is safe only because a refill step runs when fewer than 19
    bits are buffered, so no live bit is ever above bit 22.
    """

    __slots__ = ("data", "pos", "acc", "navail", "padding")

    def __init__(self, data: bytes):
        self.data = data
        self.pos = 0
        self.acc = 0
        self.navail = 0
        self.padding = 0

    def refill(self, n: int) -> None:
        data, pos, acc, navail = self.data, self.pos, self.acc, self.navail
        end = len(data)
        while navail < n:
            if pos < end:
                byte = data[pos]
            else:
                self.padding += 1
                if self.padding > PAD_ALLOWANCE:
                    raise A2FError(
                        "input exhausted: needed more than "
                        f"{PAD_ALLOWANCE} bytes of padding past the stream end"
                    )
                byte = 0
            acc = ((acc & 0x7FFFFF) << 8) | byte
            pos += 1
            navail += 8
        self.pos, self.acc, self.navail = pos, acc, navail

    def take(self, n: int) -> int:
        """Consume the next n bits."""
        self.navail -= n
        return (self.acc >> self.navail) & ((1 << n) - 1)

    def peek16(self) -> int:
        return (self.acc >> (self.navail - TABLE_BITS)) & 0xFFFF


class _Huffman:
    """Canonical Huffman decoder built exactly like Boxsim's VA 0xa10548."""

    __slots__ = ("lengths", "symtab", "lentab")

    def __init__(self) -> None:
        self.lengths = [INITIAL_CODE_LENGTH] * NSYM
        self.symtab = [0] * TABLE_SIZE
        self.lentab = [0] * TABLE_SIZE
        self.build()

    def build(self) -> None:
        lengths = self.lengths
        for sym, length in enumerate(lengths):
            if not 0 < length <= MAX_CODE_LENGTH:
                raise A2FError(f"symbol {sym} has invalid code length {length}")

        # Boxsim accumulates codes in a 24 bit register, adding 2**(24-len) per
        # assigned symbol, which is a plain canonical Huffman walk.
        codes = [0] * NSYM
        code = 0
        for bitlen in range(1, MAX_CODE_LENGTH + 1):
            step = 1 << (24 - bitlen)
            shift = 24 - bitlen
            for sym in range(NSYM):
                if lengths[sym] == bitlen:
                    codes[sym] = code >> shift
                    code += step
        if code > 0x1000000:
            raise A2FError("over-subscribed Huffman code")

        symtab, lentab = self.symtab, self.lentab
        for sym in range(NSYM):
            length = lengths[sym]
            start = codes[sym] << (TABLE_BITS - length)
            count = 1 << (TABLE_BITS - length)
            symtab[start : start + count] = [sym] * count
            lentab[start : start + count] = [length] * count

    def decode(self, bits: _BitReader) -> int:
        bits.refill(TABLE_BITS)
        idx = bits.peek16()
        length = self.lentab[idx]
        if length == 0:
            raise A2FError("incomplete Huffman code: undefined table entry")
        bits.navail -= length
        return self.symtab[idx]

    def read_lengths(self, bits: _BitReader) -> None:
        """Symbol 0x101 payload: 356 code lengths.

        Two zero bits repeat the previous length. Otherwise a 4 bit nibble plus
        one gives the length, and the value 5 is an escape meaning a second
        nibble follows and replaces it.
        """
        lengths = self.lengths
        prev = -1
        for i in range(NSYM):
            bits.refill(8)
            if (bits.acc >> (bits.navail - 2)) & 3 == 0:
                lengths[i] = prev
                bits.navail -= 2
            else:
                value = bits.take(4) + 1
                if value == 5:
                    value = bits.take(4) + 1
                lengths[i] = value
                prev = value
        self.build()


def _distance(bits: _BitReader, k: int) -> int:
    """Distance for match class index k: k+3 raw bits biased by 2**(k+3) - 7."""
    bits.refill(0x13)
    nbits = k + 3
    return bits.take(nbits) + (1 << nbits) - 7


def expand(data: bytes, *, skip_magic: int = 8) -> bytes:
    """Decompress an A2F stream, skipping `skip_magic` leading bytes."""
    bits = _BitReader(data[skip_magic:])
    huff = _Huffman()
    digrams = [0] * DICT_SIZE
    dptr = 0
    out = bytearray()

    while True:
        sym = huff.decode(bits)

        if sym < SYM_EOS:
            out.append(sym)
            continue

        if sym == SYM_EOS:
            break

        if sym == SYM_REBUILD:
            huff.read_lengths(bits)
            continue

        if sym == SYM_DIGRAM_DEF:
            bits.refill(6)
            count = bits.take(5) + 1
            for _ in range(count):
                bits.refill(0x11)
                high = huff.decode(bits)
                bits.refill(0x11)
                low = huff.decode(bits)
                digrams[dptr] = ((high << 8) | low) & 0xFFFF
                dptr = (dptr + 1) & (DICT_SIZE - 1)
            continue

        if sym >= SYM_DIGRAM:
            pair = digrams[sym - SYM_DIGRAM]
            out.append(pair >> 8)
            out.append(pair & 0xFF)
            continue

        # Match symbols.
        if sym < SYM_MATCH4:
            length = 3
            dist = _distance(bits, sym - SYM_MATCH3)
        elif sym < SYM_MATCHN:
            length = 4
            dist = _distance(bits, sym - SYM_MATCH4)
        else:
            k = sym - SYM_MATCHN
            bits.refill(0x11)
            unary = 0
            while True:
                bit = (bits.acc >> (bits.navail - 1)) & 1
                unary += 1
                bits.navail -= 1
                if not bit:
                    break
            bits.refill(0xD)
            try:
                nbits, base = LENGTH_CLASSES[unary]
            except KeyError:
                raise A2FError(f"unsupported length class {unary}") from None
            length = bits.take(nbits) + base
            dist = _distance(bits, k)

        if dist > len(out):
            raise A2FError(f"match distance {dist} exceeds {len(out)} bytes of output")
        start = len(out) - dist
        if dist >= length:
            out += out[start : start + length]
        else:
            # Overlapping copy, byte by byte, as the original does.
            for i in range(start, start + length):
                out.append(out[i])

    return bytes(out)