Initial commit

ackit/codebooks.py

@@ -0,0 +1,118 @@
+"""Wwise packed codebook library (ww2ogg 'packed_codebooks.bin' format).
+
+The game's built-in table (runtime 0x0158E748, 598 entries) is byte-identical to
+ww2ogg's standard packed_codebooks.bin; see docs/wem_format.md.
+"""
+import os
+import struct
+
+from .bits import BitReader, BitWriter, ilog, book_maptype1_quantvals
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+DEFAULT_LIB = os.path.join(HERE, "data", "game_codebooks.bin")
+
+
+def load_library(path: str = DEFAULT_LIB) -> list:
+    b = open(path, "rb").read()
+    table = struct.unpack_from("<I", b, len(b) - 4)[0]
+    n = (len(b) - table) // 4
+    offs = struct.unpack_from("<%dI" % n, b, table)
+    return [b[offs[i]:offs[i + 1]] for i in range(n - 1)]
+
+
+def save_library(path: str, books: list):
+    out = bytearray()
+    offs = []
+    for bk in books:
+        offs.append(len(out))
+        out += bk
+    offs.append(len(out))  # table start == end of last book
+    out += struct.pack("<%dI" % len(offs), *offs)
+    open(path, "wb").write(bytes(out))
+
+
+def wwise_book_to_vorbis(blob: bytes, w: BitWriter):
+    """Expand a compact Wwise codebook into a standard Vorbis codebook (ww2ogg rebuild)."""
+    r = BitReader(blob)
+    w.put(24, 0x564342)
+    dims = r.get(4); w.put(16, dims)
+    entries = r.get(14); w.put(24, entries)
+    ordered = r.get(1); w.put(1, ordered)
+    if ordered:
+        w.put(5, r.get(5))
+        cur = 0
+        while cur < entries:
+            nb = ilog(entries - cur)
+            num = r.get(nb); w.put(nb, num)
+            cur += num
+    else:
+        cwll = r.get(3)
+        sparse = r.get(1); w.put(1, sparse)
+        if cwll == 0 or cwll > 5:
+            raise ValueError("bad codeword length length")
+        for _ in range(entries):
+            present = 1
+            if sparse:
+                present = r.get(1); w.put(1, present)
+            if present:
+                w.put(5, r.get(cwll))
+    lookup = r.get(1); w.put(4, lookup)
+    if lookup == 1:
+        w.put(32, r.get(32)); w.put(32, r.get(32))
+        vlen = r.get(4); w.put(4, vlen)
+        w.put(1, r.get(1))
+        for _ in range(book_maptype1_quantvals(entries, dims)):
+            w.put(vlen + 1, r.get(vlen + 1))
+    elif lookup != 0:
+        raise ValueError("bad lookup type")
+    return r.pos
+
+
+def skip_vorbis_book(r: BitReader):
+    """Advance reader past one standard Vorbis codebook; returns (dims, entries)."""
+    if r.get(24) != 0x564342:
+        raise ValueError("codebook sync lost")
+    dims = r.get(16)
+    entries = r.get(24)
+    if r.get(1):  # ordered
+        r.get(5)
+        cur = 0
+        while cur < entries:
+            cur += r.get(ilog(entries - cur))
+    else:
+        sparse = r.get(1)
+        for _ in range(entries):
+            if not sparse or r.get(1):
+                r.get(5)
+    lookup = r.get(4)
+    if lookup in (1, 2):
+        r.get(32); r.get(32)
+        vlen = r.get(4); r.get(1)
+        n = book_maptype1_quantvals(entries, dims) if lookup == 1 else entries * dims
+        r.pos += n * (vlen + 1)
+    elif lookup != 0:
+        raise ValueError("bad lookup type %d" % lookup)
+    return dims, entries
+
+
+def book_bits_key(data: bytes, start: int, end: int) -> tuple:
+    """Hashable representation of a bit range [start,end)."""
+    n = end - start
+    v = int.from_bytes(data, "little") >> start
+    return (n, v & ((1 << n) - 1))
+
+
+class CodebookLibrary:
+    def __init__(self, path: str = DEFAULT_LIB):
+        self.books = load_library(path)
+        self.by_vorbis = {}
+        self.vorbis = []
+        for i, blob in enumerate(self.books):
+            w = BitWriter()
+            wwise_book_to_vorbis(blob, w)
+            key = (w.pos, int.from_bytes(w.bytes(), "little") & ((1 << w.pos) - 1))
+            self.vorbis.append(key)
+            self.by_vorbis.setdefault(key, i)
+
+    def lookup(self, key):
+        return self.by_vorbis.get(key)