Decode aligned LZSS banks for original-disc AOT coverage

docs/AOT_SHARDING.md

@@ -61,6 +61,7 @@ Profiles declare these reusable methods:
 | `images[].method = psx_exe` | Read the load address and image from a PS-X EXE header; the generic extractor may split resident and overlay floors. |
 | `images[].method = packed_sector_members` | Decode u32 count/offset sector descriptors across ordered payload files. Verify every member classification and requested full-payload coverage; use loader-established addresses for executable members. |
 | `images[].method = sector_extent_members` | Decode `{sector offset, byte size}` pairs in an archive. Use archive-relative offsets, verify the exact count and terminator, reject gaps/overlaps, and account for every member after sector rounding. Each member needs a verified load address or an explicit exclusion. |
+| `images[].method = aligned_lzss_banks` | Read aligned stored-size pairs, decode the metadata and its tagged bank members using parameterized LZSS, verify every container and bank classification, and deduplicate exact decoded images at verified destinations. External-RAM references fail extraction. |
 | `checks[].method = words` | Verify loader instructions or descriptors at explicit file offsets / virtual addresses. |
 | `checks[].method = pointer_strings` | Verify a pointer-indexed filename table against exact expected strings. |
 | `checks[].method = bcd_extent_table` | Verify an indexed BCD-MSF/size table against an ISO file's actual extent. |
@@ -86,6 +87,8 @@ release packaging. Pair counts are results, not inherited success thresholds.
 An inventory plus valid guards still does not prove complete static execution
 coverage or native semantics; keep fallback and perform gameplay spot checks.
 
+Jersey Devil uses `aligned_lzss_banks` for five renderer variants across 70 bank occurrences and `psx_exe` for five secondary executables. Its profile accounts for all 93 BZZ containers, including 23 without code banks.
+
 Current consumers:
 
 - Tomba! USA: `TombaRecomp/aot/overlays.json`, 25 fixed-address images, including

tools/aligned_lzss_banks.py

@@ -0,0 +1,104 @@
+"""Sector-aligned size-pair containers with tagged, compressed code banks.
+
+The compression stream has a parameter byte, a big-endian 24-bit token count
+minus one, LSB-first literal flags, and 16-bit big-endian backward references.
+The parameter selects distance/length bits and a piecewise length table.
+No emulated RAM, captured code, or game-specific names are inputs.
+"""
+import struct
+
+
+def require(condition, message):
+    if not condition:
+        raise ValueError(message)
+
+
+def decode(data, *, max_output=2_000_000):
+    require(len(data) >= 5 and max_output > 0, 'Truncated LZSS stream or invalid limit')
+    parameter = data[0]
+    require(parameter < 32, 'Unsupported LZSS parameter bits')
+    length_bits = 7 - (parameter & 7)
+    mask = 127 >> (parameter & 7)
+    threshold = 19 if mask >= 31 else mask // 2
+    shift = (parameter >> 3) & 3
+    lengths = [(i if i <= threshold else threshold + ((i - threshold) << shift)) + 3
+               for i in range(mask + 1)]
+    tokens = int.from_bytes(data[1:4], 'big') + 1
+    require(tokens <= max_output, 'LZSS token count exceeds output limit')
+    position, flags = 4, 0
+    output = bytearray()
+
+    def byte():
+        nonlocal position
+        require(position < len(data), 'Truncated LZSS token')
+        value = data[position]
+        position += 1
+        return value
+
+    for _ in range(tokens):
+        if flags <= 1:
+            flags = byte() | 256
+        literal = flags & 1
+        flags >>= 1
+        token = byte()
+        if literal:
+            require(len(output) < max_output, 'LZSS output exceeds limit')
+            output.append(token)
+        else:
+            token = token * 256 + byte()
+            distance, length = token >> length_bits, lengths[token & mask]
+            # A zero/out-of-image reference depends on pre-existing RAM. It
+            # cannot produce deterministic original-disc AOT code.
+            require(0 < distance <= len(output), 'LZSS reference depends on external RAM')
+            require(len(output) + length <= max_output, 'LZSS output exceeds limit')
+            for _ in range(length):
+                output.append(output[-distance])
+    return bytes(output), position
+
+
+def members(data, *, alignment=2048, version=1):
+    require(alignment >= 16 and len(data) >= alignment, 'Truncated size-pair archive')
+    actual_version, count = struct.unpack_from('<II', data)
+    require(actual_version == version, 'Unsupported size-pair archive version')
+    require(0 < count <= (alignment - 8) // 8, 'Invalid size-pair member count')
+    require(not any(data[8 + count * 8:alignment]), 'Unexpected archive header tail')
+    offset, result = alignment, []
+    for index in range(count):
+        stored, padded = struct.unpack_from('<II', data, 8 + index * 8)
+        require(stored > 0 and padded == (stored + alignment - 1) // alignment * alignment,
+                'Invalid aligned member size')
+        require(offset + padded <= len(data), 'Truncated aligned member')
+        result.append(dict(index=index, offset=offset, size=stored,
+                           body=data[offset:offset + stored]))
+        offset += padded
+    require(offset == len(data), 'Size-pair inventory does not cover archive')
+    return result
+
+
+def banks(data, *, alignment=2048, version=1, bank_tag=0x4B,
+          terminal_tag=0x31, data_tags=(0x30,), max_output=2_000_000):
+    """Decode only metadata and its explicit leading bank members.
+
+    Remaining members are data payloads, deliberately not decoded as code.
+    The caller supplies loader-established bank-index to RAM-address mappings.
+    """
+    inventory = members(data, alignment=alignment, version=version)
+    metadata, _ = decode(inventory[0]['body'], max_output=max_output)
+    position, result = 0, []
+    while True:
+        require(position + 4 <= len(metadata), 'Truncated bank selector metadata')
+        tag = struct.unpack_from('<I', metadata, position)[0]
+        if tag != bank_tag:
+            require(tag == terminal_tag if result else tag in data_tags,
+                    'Unknown bank metadata tag')
+            break
+        require(position + 8 <= len(metadata), 'Truncated bank selector')
+        bank = struct.unpack_from('<I', metadata, position + 4)[0]
+        index = len(result) + 1
+        require(index < len(inventory), 'Bank selector has no archive member')
+        member = inventory[index]
+        body, consumed = decode(member['body'], max_output=max_output)
+        result.append(dict(index=index, bank=bank, source_offset=member['offset'],
+                           stored_size=member['size'], consumed=consumed, body=body))
+        position += 8
+    return inventory, result

tools/aot_overlay_pipeline.py

@@ -21,6 +21,7 @@ import compile_overlays as compiler
 from aot_overlay_spike import extract_generic as extractor
 from packed_sector_table import extract_members as extract_sector_members
 from sector_extent_archive import extract_members as extract_extent_members
+from aligned_lzss_banks import banks as extract_lzss_banks
 
 FRAMEWORK = Path(__file__).resolve().parents[1]
 
@@ -146,6 +147,44 @@ def positioned_sources(disc, specifications):
     sources = []
     for spec in specifications:
         method = spec['method']
+        if method == 'aligned_lzss_banks':
+            grouped = {}
+            containers = spec['containers']
+            require(len({item['file'].upper() for item in containers}) == len(containers),
+                    'Duplicate compressed container')
+            for item in containers:
+                file = item['file'].upper()
+                inventory, banks = extract_lzss_banks(disc.read(file),
+                    alignment=number(spec.get('alignment', 2048)),
+                    version=number(spec.get('version', 1)),
+                    bank_tag=number(spec.get('bank_tag', 0x4B)),
+                    terminal_tag=number(spec.get('terminal_tag', 0x31)),
+                    data_tags=tuple(number(x) for x in spec.get('data_tags', [0x30])))
+                require(len(inventory) == item['member_count'] and
+                        [bank['bank'] for bank in banks] == item['bank_indices'],
+                        f'Compressed bank inventory changed: {file}')
+                for bank in banks:
+                    placement = item['placements'][str(bank['bank'])]
+                    base = number(placement['load_addr'])
+                    body = bank['body']
+                    require(hashlib.sha256(body).hexdigest() == placement['decoded_sha256'],
+                            f'Decoded bank changed: {file}')
+                    require(0x80000000 <= base < base + len(body) <= 0x80200000,
+                            'Decoded bank outside RAM')
+                    name = f"{file}:BANK_{bank['bank']:04X}"
+                    key = (base, body)
+                    if key in grouped:
+                        grouped[key]['aliases'].append(name)
+                        require(grouped[key]['spec']['entries'] == placement.get('entries', []),
+                                'Duplicate bank has conflicting declared entries')
+                        continue
+                    source = dict(name=name, base=base, body=body,
+                        spec={**spec, **placement, 'entries': placement.get('entries', []),
+                              'allow_missing': True},
+                        source_file=file, source_offset=bank['source_offset'], aliases=[])
+                    grouped[key] = source
+                    sources.append(source)
+            continue
         if method == 'packed_sector_members':
             sources.extend(sector_sources(disc, spec))
             continue
@@ -278,7 +317,7 @@ def prepare(profile, disc, records, output):
             if (lo, hi) == (source['base'], source['base'] + len(source['body'])):
                 singles[source['name']] = record
     for source in sources:
-        if source['spec']['method'] in ('fixed_address_files', 'packed_sector_members', 'sector_extent_members') and source['name'] not in singles:
+        if source['spec']['method'] in ('fixed_address_files', 'packed_sector_members', 'sector_extent_members', 'aligned_lzss_banks') and source['name'] not in singles:
             record = make_fixed_record(source, disc)
             singles[source['name']] = record
             records.append(record)
@@ -337,6 +376,9 @@ def prepare(profile, disc, records, output):
                      profile_sha256=hashlib.sha256(json.dumps(profile, sort_keys=True).encode()).hexdigest(),
                      original_disc_sha256=digest(disc.binary), required_images=sorted(by_name),
                      recipe_count=len(jobs), full_static_coverage_proven=False, jobs=jobs)
+    inventory['source_images'] = [dict(name=s['name'], aliases=s['aliases'],
+        method=s['spec']['method'], load_addr=hex(s['base']), size=len(s['body']),
+        sha256=hashlib.sha256(s['body']).hexdigest()) for s in sources]
     write_json(output / 'runtime-input-inventory.json', inventory)
     return inventory
 

tools/aot_overlay_spike/extract_generic.py

@@ -703,6 +703,11 @@ def optional_entry_delay_valid(body, base, entry):
         return False
     for at in range(offset, min(offset + 48, len(body) - 3), 4):
         word = _word(body, at)
+        # BLEZ/BGTZ require rt=zero. Library identification data following a
+        # return can resemble these opcodes with a nonzero reserved field.
+        # Such a prefix cannot establish an optional callable entry.
+        if word >> 26 in (6, 7) and (word >> 16) & 31:
+            return False
         if _is_control_flow_word(word):
             delay = _word(body, at + 4)
             return delay is not None and not _is_control_flow_word(delay)
@@ -1000,6 +1005,8 @@ def main():
                 aliases=[alias for alias in aliases_all
                          if va <= alias[0] < span_hi and
                          va <= alias[1] < alias[2] <= span_hi and
+                         (alias[0] == entry_pc or
+                          optional_entry_delay_valid(body, t_addr, alias[0])) and
                          (alias[0] != t_addr or alias[0] == entry_pc)]
                 records.append(rec(
                     va,seg,sd,static_alias_ranges=aliases))

tools/tests/test_aligned_lzss_banks.py

@@ -0,0 +1,100 @@
+import struct
+import hashlib
+import sys
+import unittest
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from aligned_lzss_banks import decode, members, banks
+from aot_overlay_spike.extract_generic import filter_full_discovery_seeds
+from aot_overlay_pipeline import positioned_sources
+
+
+def literal_stream(body):
+    out = bytearray(b'\0' + (len(body) - 1).to_bytes(3, 'big'))
+    for offset in range(0, len(body), 8):
+        block = body[offset:offset + 8]
+        out.append((1 << len(block)) - 1)
+        out.extend(block)
+    return bytes(out)
+
+
+def archive(*payloads):
+    head = bytearray(64)
+    struct.pack_into('<II', head, 0, 1, len(payloads))
+    result = bytearray()
+    for i, body in enumerate(payloads):
+        padded = (len(body) + 63) // 64 * 64
+        struct.pack_into('<II', head, 8 + i * 8, len(body), padded)
+        result.extend(body + bytes(padded - len(body)))
+    return bytes(head + result)
+
+
+class CompressedBankTests(unittest.TestCase):
+    def test_flags_cross_byte(self):
+        body = b'original disc code'
+        stream = literal_stream(body)
+        self.assertEqual(decode(stream), (body, len(stream)))
+
+    def test_overlapping_reference_and_length_curve(self):
+        # Parameter 8: seven length bits, x2 slope above index 19.
+        # Literal A then distance 1 / length index 20 => 24 more As.
+        self.assertEqual(decode(bytes.fromhex('0800000101410094'))[0], b'A' * 25)
+        # Parameter 7: zero length bits, distance occupies all 16 bits.
+        self.assertEqual(decode(bytes.fromhex('0700000101410001'))[0], b'AAAA')
+
+    def test_external_ram_and_truncation_rejected(self):
+        for stream in [b'', bytes.fromhex('000000010141'),
+                       bytes.fromhex('0000000101410000'),
+                       bytes.fromhex('0000000101410100')]:
+            with self.subTest(stream=stream), self.assertRaises(ValueError):
+                decode(stream)
+        with self.assertRaises(ValueError):
+            decode(literal_stream(b'ab'), max_output=1)
+
+    def test_selected_bank_and_uninterpreted_data(self):
+        meta = literal_stream(struct.pack('<III', 0x4B, 3, 0x31))
+        data = archive(meta, literal_stream(b'code'), b'not a compressed code stream')
+        inventory, selected = banks(data, alignment=64)
+        self.assertEqual(len(inventory), 3)
+        self.assertEqual([(x['index'], x['bank'], x['body']) for x in selected], [(1, 3, b'code')])
+
+    def test_data_only_container(self):
+        data = archive(literal_stream(struct.pack('<I', 0x30)), b'assets')
+        self.assertEqual(banks(data, alignment=64)[1], [])
+
+    def test_archive_and_selector_drift_rejected(self):
+        data = archive(literal_stream(struct.pack('<III', 0x4B, 0, 0x31)), literal_stream(b'code'))
+        for changed in [data[:-1], data + bytes(64), b'\2' + data[1:]]:
+            with self.subTest(size=len(changed)), self.assertRaises(ValueError):
+                members(changed, alignment=64)
+        with self.assertRaises(ValueError):
+            banks(archive(literal_stream(struct.pack('<I', 0x4B))), alignment=64)
+        with self.assertRaises(ValueError):
+            banks(archive(literal_stream(struct.pack('<I', 0x99))), alignment=64)
+
+    def test_reserved_branch_prefix_is_not_a_discovered_entry(self):
+        base = 0x80010000
+        data = struct.pack('<6I', 0, 0x19097350, 0x004162B3, 0x03e00008, 0, 0)
+        self.assertEqual(filter_full_discovery_seeds(data, base, [base+4, base+12], base), [base+12])
+        self.assertEqual(filter_full_discovery_seeds(data, base, [base+4], base+4), [base+4])
+
+    def test_pipeline_deduplicates_verified_banks_and_rejects_drift(self):
+        raw = archive(literal_stream(struct.pack('<III', 0x4B, 0, 0x31)), literal_stream(b'code'))
+        class Disc:
+            def read(self, name):
+                return raw
+        placement = dict(load_addr='0x80010200', decoded_sha256=hashlib.sha256(b'code').hexdigest())
+        spec = dict(method='aligned_lzss_banks', alignment=64, containers=[
+            dict(file=name, member_count=2, bank_indices=[0], placements={'0': placement})
+            for name in ['A.PAK', 'B.PAK']])
+        sources = positioned_sources(Disc(), [spec])
+        self.assertEqual(len(sources), 1)
+        self.assertEqual(sources[0]['aliases'], ['B.PAK:BANK_0000'])
+        spec['containers'][1]['bank_indices'] = []
+        with self.assertRaises(ValueError):
+            positioned_sources(Disc(), [spec])
+
+
+if __name__ == '__main__':
+    unittest.main()