libretro/scripts/derivations.py

"""Catalog entries whose bytes a held dump already contains.

A catalog lists chips as they come off the board; a collection often holds
the image an emulator reads. The Amiga 3.1 Kickstart pair 391773/391774 is
the high and low 16-bit words of amiga-os-310-a1200.rom, byte-swapped; the
TOSEC "Basic V7" is basiclo and basichi end to end. Listed as missing, they
send a maintainer after files that are already here. Each transform below is
one the source boards use; a match is by sha1, never by size or name.
"""

from __future__ import annotations

import array
import functools
import hashlib
import os
import zipfile
from collections import defaultdict
from collections.abc import Callable, Iterable

# A held file larger than this is not read: no chip image is near it.
MAX_SOURCE_SIZE = 64 * 1024 * 1024


def _swap16(data: bytes) -> bytes:
    swapped = bytearray(data)
    swapped[0::2], swapped[1::2] = data[1::2], data[0::2]
    return bytes(swapped)


def _words(data: bytes, start: int) -> bytes:
    """Every other 16-bit word, from word `start`, bytes kept in order.

    Read as an array of 2-byte units and strided: joining slices one word
    at a time made this pass take most of an hour over the romset chips.
    """
    words = array.array("H")
    words.frombytes(data[: len(data) - len(data) % 2])
    return words[start::2].tobytes()


# (name, size of the output for an input of size n or None, transform)
TRANSFORMS: tuple[tuple[str, Callable[[int], int | None], Callable[[bytes], bytes]], ...] = (
    ("first half", lambda n: n // 2 if n % 2 == 0 else None, lambda d: d[:len(d) // 2]),
    ("second half", lambda n: n // 2 if n % 2 == 0 else None, lambda d: d[len(d) // 2:]),
    *(
        (f"quarter {q + 1}", lambda n: n // 4 if n % 4 == 0 else None,
         lambda d, q=q: d[q * len(d) // 4:(q + 1) * len(d) // 4])
        for q in range(4)
    ),
    ("even bytes", lambda n: n // 2 if n % 2 == 0 else None, lambda d: d[0::2]),
    ("odd bytes", lambda n: n // 2 if n % 2 == 0 else None, lambda d: d[1::2]),
    ("even words", lambda n: n // 2 if n % 4 == 0 else None, lambda d: _words(d, 0)),
    ("odd words", lambda n: n // 2 if n % 4 == 0 else None, lambda d: _words(d, 1)),
    ("even words, byte-swapped", lambda n: n // 2 if n % 4 == 0 else None,
     lambda d: _swap16(_words(d, 0))),
    ("odd words, byte-swapped", lambda n: n // 2 if n % 4 == 0 else None,
     lambda d: _swap16(_words(d, 1))),
    ("byte-swapped", lambda n: n if n % 2 == 0 else None, _swap16),
    ("mirrored twice", lambda n: n * 2, lambda d: d + d),
)


def _gf2_times(matrix: list[int], vector: int) -> int:
    total, row = 0, 0
    while vector:
        if vector & 1:
            total ^= matrix[row]
        vector >>= 1
        row += 1
    return total


def _gf2_square(matrix: list[int]) -> list[int]:
    return [_gf2_times(matrix, matrix[n]) for n in range(32)]


def crc32_combine(crc1: int, crc2: int, len2: int) -> int:
    """crc32 of A + B from crc32(A), crc32(B) and len(B), as zlib computes it."""
    if len2 <= 0:
        return crc1
    return _gf2_times(_shift_operator(len2), crc1) ^ crc2


@functools.lru_cache(maxsize=None)
def _shift_operator(len2: int) -> tuple[int, ...]:
    """The linear map zlib applies to crc(A) for len2 bytes appended.

    Built once per length, as zlib's crc32_combine walks it: the chain of
    squared operators depends on len2 alone, so it is squared once and the
    operators the set bits select are composed. Rebuilding the chain for
    each of the 32 columns spent 57 s of the provenance step in GF(2)
    products.
    """
    odd = [0xEDB88320] + [1 << (n - 1) for n in range(1, 32)]
    even = _gf2_square(odd)
    odd = _gf2_square(even)
    selected: list[list[int]] = []
    while True:
        even = _gf2_square(odd)
        if len2 & 1:
            selected.append(even)
        len2 >>= 1
        if not len2:
            break
        odd = _gf2_square(even)
        if len2 & 1:
            selected.append(odd)
        len2 >>= 1
        if not len2:
            break
    columns = [1 << bit for bit in range(32)]
    for matrix in selected:
        columns = [_gf2_times(matrix, column) for column in columns]
    return tuple(columns)


def _read(held: dict, archives: dict[str, zipfile.ZipFile] | None = None) -> bytes | None:
    """The bytes of a held file, or of a member read out of its archive.

    `archives` keeps each archive open across its members: reopened for
    every one, a MAME set of a hundred chips parsed its directory a
    hundred times.
    """
    try:
        if held.get("member"):
            if archives is None:
                with zipfile.ZipFile(held["archive"]) as zf:
                    return zf.read(held["member"])
            zf = archives.get(held["archive"])
            if zf is None:
                zf = archives[held["archive"]] = zipfile.ZipFile(held["archive"])
            return zf.read(held["member"])
        with open(held["path"], "rb") as handle:
            return handle.read()
    except (OSError, KeyError, zipfile.BadZipFile):
        return None


def _held_sources(db: dict) -> list[dict]:
    """Every file the collection holds, the members of its archives included.

    A romset holds chips a catalog lists on their own: cdi910.rom inside
    cdi910.zip is the CD-i 910 BIOS byte-swapped, and the two halves of the
    Salora Manager sit in manager.zip. Read from the top level alone, they
    stayed MISSING. A member's group is its archive, as a file's is its
    directory: the pairs a catalog splits come from one machine.
    """
    sources: list[dict] = []
    for record in db.get("files", {}).values():
        path = record.get("path", "")
        if not os.path.exists(path):
            continue
        if 0 < record.get("size", 0) <= MAX_SOURCE_SIZE:
            sources.append({**record, "group": os.path.dirname(path)})
        if not path.lower().endswith(".zip"):
            continue
        try:
            with zipfile.ZipFile(path) as zf:
                infos = [i for i in zf.infolist() if not i.is_dir()]
        except (OSError, zipfile.BadZipFile):
            continue
        for info in infos:
            if 0 < info.file_size <= MAX_SOURCE_SIZE:
                sources.append({
                    "path": f"{path}/{info.filename}", "archive": path,
                    "member": info.filename, "size": info.file_size,
                    "crc32": f"{info.CRC:08x}", "group": path,
                })
    return sources


def derivable(entries: Iterable[dict], db: dict) -> dict[str, dict]:
    """Missing sha1 to the held file or member and the transform that
    produce its bytes; archives stay open for the whole pass."""
    archives: dict[str, zipfile.ZipFile] = {}
    try:
        return _derivable(entries, db, archives)
    finally:
        for zf in archives.values():
            zf.close()


def _derivable(
    entries: Iterable[dict], db: dict, archives: dict[str, zipfile.ZipFile]
) -> dict[str, dict]:
    """Only held files whose size can produce a wanted size are read, and a
    concatenation is only tried within one directory or one archive: the
    chip pairs a catalog splits come from one machine.
    """
    wanted: dict[tuple[int, str], str] = {}
    by_size_crc: dict[tuple[int, int], list[str]] = defaultdict(list)
    for entry in entries:
        sha1, size = (entry.get("sha1") or "").lower(), entry.get("size")
        if sha1 and isinstance(size, int) and size > 0:
            wanted[(size, sha1)] = sha1
            crc = entry.get("crc32")
            if crc:
                by_size_crc[(size, int(str(crc), 16))].append(sha1)
    sizes = {size for size, _sha1 in wanted}
    found: dict[str, dict] = {}

    files = _held_sources(db)
    for held in files:
        n = held["size"]
        useful = [t for t in TRANSFORMS if t[1](n) in sizes]
        if not useful:
            continue
        data = _read(held, archives)
        if data is None:
            continue
        for name, out_size, transform in useful:
            out = transform(data)
            sha1 = wanted.get((out_size(n), hashlib.sha1(out).hexdigest()))
            if sha1 and sha1 not in found:
                found[sha1] = {"from": [held["path"]], "transform": name}

    by_dir: dict[str, list[dict]] = defaultdict(list)
    for held in files:
        if held.get("crc32"):
            by_dir[held["group"]].append(held)
    for group in by_dir.values():
        # A pair is only worth a crc when its sizes add up to a wanted one:
        # indexed by size, a group is walked once per wanted size instead
        # of once per member, and a romset of a few hundred chips stays
        # linear.
        by_size: dict[int, list[dict]] = defaultdict(list)
        for held in group:
            by_size[held["size"]].append(held)
        for first in group:
            for total in sizes:
                for second in by_size.get(total - first["size"], ()):
                    if first is second:
                        continue
                    crc = crc32_combine(int(first["crc32"], 16),
                                        int(second["crc32"], 16), second["size"])
                    candidates = [s for s in by_size_crc.get((total, crc), [])
                                  if s not in found]
                    if not candidates:
                        continue
                    a, b = _read(first, archives), _read(second, archives)
                    if a is None or b is None:
                        continue
                    sha1 = hashlib.sha1(a + b).hexdigest()
                    if sha1 in candidates:
                        found[sha1] = {"from": [first["path"], second["path"]],
                                       "transform": "concatenated"}
    return found