From 32cc2e81f426436ff4bb169114f2e603d11cb0a7 Mon Sep 17 00:00:00 2001
From: Matthew Stanley <1379tech@gmail.com>
Date: Sat, 18 Jul 2026 08:27:09 -0700
Subject: [PATCH] aot: bound header dispatch fallback discovery

---
 docs/AOT_OVERLAY_PLAN.md                      |  7 ++
 recompiler/tests/test_aot_base_recovery.py    | 21 ++++++
 .../tests/test_aot_overlay_discovery.py       | 26 +++++++
 tools/aot_overlay_spike/README.md             | 12 ++++
 tools/aot_overlay_spike/extract_generic.py    | 70 ++++++++++++++++---
 tools/compile_overlays.py                     | 62 ++++++++++++----
 6 files changed, 174 insertions(+), 24 deletions(-)

diff --git a/docs/AOT_OVERLAY_PLAN.md b/docs/AOT_OVERLAY_PLAN.md
index 651df507..d58e43a9 100644
--- a/docs/AOT_OVERLAY_PLAN.md
+++ b/docs/AOT_OVERLAY_PLAN.md
@@ -50,6 +50,13 @@ continue to fail safely to the interpreter.
   100% combined BIOS code-range coverage for both vault and live history.
 - Large archive reads are now linear rather than quadratic sector concatenation.
   These validations wrote only `%TEMP%` caches; title caches remain untouched.
+- Tomba 1 header exports are authoritative mid-function dispatch entries rather
+  than callable roots. Their conservative retry now uses only the nearest
+  preceding prologue host, hard-stops it at the next prologue, and rejects calls
+  outside those producer ranges. The clean 25-record corpus built 18 unique DLLs
+  with zero failures (the rest deduplicated); across 24 legacy sessions it covers
+  86.2% of played PCs by overlay code range and 95.1% with base BIOS/kernel code,
+  leaving 22 genuine range misses rather than admitting embedded data as code.
 
 ---
 
diff --git a/recompiler/tests/test_aot_base_recovery.py b/recompiler/tests/test_aot_base_recovery.py
index 7447a372..3647d015 100644
--- a/recompiler/tests/test_aot_base_recovery.py
+++ b/recompiler/tests/test_aot_base_recovery.py
@@ -38,6 +38,27 @@ def main():
     struct.pack_into("<I", data, 4, 0x0C000000 | ((0x80105678 >> 2) & 0x03FFFFFF))
     assert MOD.jal_targets(data) == {0x80105678}
 
+    # A JAL target becomes a conservative shard root only when the target is
+    # local and has independent callable-boundary evidence.
+    rooted = bytearray(0x80)
+    struct.pack_into('<I', rooted, 0, 0x0C000000 |
+                     (((0x80100000 + 0x40) >> 2) & 0x03FFFFFF))
+    struct.pack_into('<I', rooted, 0x40, 0x27BDFFF0)
+    assert MOD.direct_jal_roots(rooted, 0x80100000) == [0x80100040]
+    struct.pack_into('<I', rooted, 0x40, 0x24020001)
+    assert MOD.direct_jal_roots(rooted, 0x80100000) == []
+
+    bounded = bytearray(0x90)
+    for off in (0x10, 0x40, 0x70):
+        struct.pack_into('<I', bounded, off, 0x27BDFFF0)
+    recipe = MOD.bounded_dispatch_fallback(
+        bounded, 0x80100000, [0x80100024, 0x80100050])
+    assert recipe['function_entry_pcs'] == ['0x80100010', '0x80100040']
+    assert recipe['producer_ranges'] == [
+        {'start': '0x80100010', 'end': '0x80100040'},
+        {'start': '0x80100040', 'end': '0x80100070'},
+    ]
+
     wrapped = MOD.make_psx_exe(b'\x11\x22\x33\x44', 0x80102000,
                                0x80102040)
     assert wrapped[:8] == b'PS-X EXE'
diff --git a/recompiler/tests/test_aot_overlay_discovery.py b/recompiler/tests/test_aot_overlay_discovery.py
index d72899b0..bd6e60cc 100644
--- a/recompiler/tests/test_aot_overlay_discovery.py
+++ b/recompiler/tests/test_aot_overlay_discovery.py
@@ -65,6 +65,12 @@ def check_composite_call_boundaries():
         data, LOAD, len(data), LOAD, LOAD + 0x100, ranges)
     assert cross_target in scoped['direct_jals']
 
+    strict = MOD._walk_overlay_function(
+        data, LOAD, len(data), LOAD, LOAD + 0x100, ranges,
+        allow_cross_producer_calls=False)
+    assert cross_target not in strict['direct_jals']
+    assert cross_target in strict['rejected_cross_producer_calls']
+
     # Independent target-local boundary evidence makes a real cross-producer
     # export safe to retain.
     put(data, 0x50, 0x27BDFFE0)
@@ -141,6 +147,26 @@ def check_optional_enrichment_fallback():
     assert 'optional_enrichment_fallback_entry_pcs' not in fallback
     assert MOD.optional_enrichment_fallback_capture(fallback) is None
 
+    dispatch = LOAD + 0x30
+    cap['optional_enrichment_fallback_entry_pcs'] = {
+        'function_entry_pcs': [],
+        'dispatch_entry_pcs': [f'0x{dispatch:08X}'],
+        'static_discovery_entry_pcs': [],
+        'producer_ranges': [{
+            'start': f'0x{LOAD + 0x20:08X}',
+            'end': f'0x{LOAD + 0x40:08X}',
+        }],
+        'strict_producer_ranges': True,
+    }
+    fallback = MOD.optional_enrichment_fallback_capture(cap)
+    assert fallback['function_entry_pcs'] == []
+    assert fallback['dispatch_entry_pcs'] == [f'0x{dispatch:08X}']
+    assert fallback['static_discovery_entry_pcs'] == []
+    assert fallback['seeds'] == []
+    assert fallback['producer_ranges'] == cap[
+        'optional_enrichment_fallback_entry_pcs']['producer_ranges']
+    assert fallback['strict_producer_ranges'] is True
+
 
 def check_forward_branch_root():
     data = bytearray(0x80)
diff --git a/tools/aot_overlay_spike/README.md b/tools/aot_overlay_spike/README.md
index b0b32464..8df675c4 100644
--- a/tools/aot_overlay_spike/README.md
+++ b/tools/aot_overlay_spike/README.md
@@ -227,6 +227,18 @@ including the separately recompiled, live-byte-guarded base BIOS raises combined
 native code-range recall to **100.0%** (677/677). Validation used only disposable
 `%TEMP%` caches.
 
+Tomba 1's header-table exports require a different conservative retry. Those
+authoritative addresses are commonly switch-case entries in the middle of a
+function, so a direct-JAL-only retry loses the dispatcher while a broad scan can
+walk into embedded data. The extractor now pairs each export with its nearest
+preceding prologue, hard-bounds that host at the next prologue, and forbids calls
+from escaping those proven producer ranges. A clean 25-record build produced 18
+unique DLLs with zero failures; the other records were safely deduplicated. Across
+24 legacy capture files, the disposable cache covers **86.2%** of exercised PCs
+by overlay code range and **95.1%** with the separately generated BIOS/kernel
+ranges, leaving 22 true range misses for later producer recovery. The generated-C
+audit remains the gate: rejected broad discovery is never installed.
+
 Ape Escape's `KKIIDDZZ.HED` encodes contiguous
 `size_sectors:12 | logical_sector:20` runs spanning sibling DAT then BNS files.
 The same cross-member vote gate proves `0x80136000` (10 anchors), `0x8013D000`
diff --git a/tools/aot_overlay_spike/extract_generic.py b/tools/aot_overlay_spike/extract_generic.py
index 12e6d49a..7e05fb09 100644
--- a/tools/aot_overlay_spike/extract_generic.py
+++ b/tools/aot_overlay_spike/extract_generic.py
@@ -149,6 +149,47 @@ def jal_targets(data):
             ts.add(0x80000000 | ((w & 0x03FFFFFF) << 2))
     return ts
 
+def direct_jal_roots(data, base):
+    """Locally bounded call targets with independent boundary evidence."""
+    lo,hi=base,base+len(data)
+    return sorted(t for t in jal_targets(data)
+                  if lo <= t < hi and (t&3)==0 and
+                  (_callable_boundary(data,t-base) or
+                   ((_word(data,t-base) or 0)>>16)==0x27BD))
+
+def bounded_dispatch_fallback(data, base, entries):
+    """Bound authoritative mid-function entries to their nearest hosts.
+
+    Header-table exports are real dispatch targets but are commonly switch
+    cases rather than function starts.  For a conservative retry, pair each
+    with only its nearest preceding prologue and stop that host at the very
+    next prologue.  A host without a provable upper boundary is omitted.
+    """
+    starts=prologues(data,base)
+    if len(starts)<2:
+        return None
+    hosts=set(); dispatch=set(); ranges=set()
+    for entry in sorted(set(entries)):
+        before=[start for start in starts if start<=entry]
+        if not before:
+            continue
+        host=before[-1]
+        end=next((start for start in starts if start>host),None)
+        if end is None or not (host<=entry<end):
+            continue
+        hosts.add(host); dispatch.add(entry); ranges.add((host,end))
+    if not hosts or not dispatch:
+        return None
+    return {
+        'function_entry_pcs':[f"0x{addr:08X}" for addr in sorted(hosts)],
+        'dispatch_entry_pcs':[f"0x{addr:08X}" for addr in sorted(dispatch)],
+        'static_discovery_entry_pcs':[],
+        'producer_ranges':[
+            {'start':f"0x{lo:08X}",'end':f"0x{hi:08X}"}
+            for lo,hi in sorted(ranges)],
+        'strict_producer_ranges':True,
+    }
+
 def _select_base_vote(hist, trusted_bases=()):
     """Select a sharp vote peak, or an exact independently trusted near-tie.
 
@@ -330,17 +371,13 @@ def recover_consensus_members(members, base_lo, base_hi=0x80200000):
         runner=ranked[1][0] if len(ranked)>1 else 0
         if score < 2 or score < runner*2:
             continue
-        lo,hi=base,base+len(body)
-        direct={t for t in jal_targets(body)
-                if lo <= t < hi and (t&3)==0 and
-                (_callable_boundary(body,t-base) or
-                 ((_word(body,t-base) or 0)>>16)==0x27BD)}
-        seeds=sorted(direct)
+        direct=direct_jal_roots(body,base)
+        seeds=list(direct)
         if len(direct) < 2:
             continue
         recovered.append({
             'id':ident, 'file_offset':off, 'data':body, 'base':base,
-            'score':score, 'runner':runner, 'direct_seeds':sorted(direct),
+            'score':score, 'runner':runner, 'direct_seeds':direct,
             'seeds':seeds,
         })
     return recovered
@@ -815,7 +852,17 @@ def main():
         page_base = base & ~0xFFF
         fill = base - page_base
         region = b'\x00'*fill + data
-        records.append(rec(page_base, region, seeds, dispatch_extra=hdr_entries))
+        direct=direct_jal_roots(data,base)
+        record=rec(page_base, region, seeds, dispatch_extra=hdr_entries,
+                   static_discovery=direct)
+        fallback=bounded_dispatch_fallback(data,base,hdr_entries)
+        if fallback:
+            # The on-disc export table is authoritative dispatch evidence but
+            # its entries are commonly mid-function.  On audit rejection,
+            # classify those entries from scratch without retaining broad
+            # prologue/pointer guesses or trusting unrooted JAL-shaped data.
+            record['optional_enrichment_fallback_entry_pcs'] = fallback
+        records.append(record)
         positioned.append((p,data,base,seeds,hdr_entries))
         nh+=1
         extra=len(set(hdr_entries)-set(seeds))
@@ -838,7 +885,12 @@ def main():
             supplemental=supplemental_callable_seeds(data,base)
             seeds=sorted(set(prologue_seeds)|supplemental)
             page_base,region=page_aligned_region(base,data)
-            records.append(rec(page_base,region,seeds))
+            direct=direct_jal_roots(data,base)
+            record=rec(page_base,region,seeds,static_discovery=direct)
+            if direct and set(seeds) != set(direct):
+                record['optional_enrichment_fallback_entry_pcs'] = [
+                    f"0x{addr:08X}" for addr in direct]
+            records.append(record)
             positioned.append((p,data,base,seeds,()))
             nr+=1
             print(f"  [raw-code] {p}: {len(data)}B file@0x{base:08X} "
diff --git a/tools/compile_overlays.py b/tools/compile_overlays.py
index 2e976222..712c75a3 100644
--- a/tools/compile_overlays.py
+++ b/tools/compile_overlays.py
@@ -364,15 +364,38 @@ def optional_enrichment_fallback_capture(cap: dict) -> dict | None:
     shard while all optional aliases and normal-mode discoveries are dropped.
     """
     key = 'optional_enrichment_fallback_entry_pcs'
-    entries = sorted(_parse_addr_list(cap.get(key, [])))
-    if not entries:
-        return None
-    encoded = [f'0x{entry:08X}' for entry in entries]
+    recipe = cap.get(key, [])
     fallback = dict(cap)
-    fallback['function_entry_pcs'] = encoded
-    fallback['dispatch_entry_pcs'] = encoded
-    fallback['static_discovery_entry_pcs'] = encoded
-    fallback['seeds'] = encoded
+    if isinstance(recipe, dict):
+        functions = sorted(_parse_addr_list(
+            recipe.get('function_entry_pcs', [])))
+        dispatch = sorted(_parse_addr_list(
+            recipe.get('dispatch_entry_pcs', [])))
+        static = sorted(_parse_addr_list(
+            recipe.get('static_discovery_entry_pcs', [])))
+        if not (functions or dispatch or static):
+            return None
+        fallback['function_entry_pcs'] = [
+            f'0x{entry:08X}' for entry in functions]
+        fallback['dispatch_entry_pcs'] = [
+            f'0x{entry:08X}' for entry in dispatch]
+        fallback['static_discovery_entry_pcs'] = [
+            f'0x{entry:08X}' for entry in static]
+        fallback['seeds'] = fallback['function_entry_pcs']
+        if 'producer_ranges' in recipe:
+            fallback['producer_ranges'] = recipe['producer_ranges']
+        if 'strict_producer_ranges' in recipe:
+            fallback['strict_producer_ranges'] = bool(
+                recipe['strict_producer_ranges'])
+    else:
+        entries = sorted(_parse_addr_list(recipe))
+        if not entries:
+            return None
+        encoded = [f'0x{entry:08X}' for entry in entries]
+        fallback['function_entry_pcs'] = encoded
+        fallback['dispatch_entry_pcs'] = encoded
+        fallback['static_discovery_entry_pcs'] = encoded
+        fallback['seeds'] = encoded
     fallback.pop('static_alias_ranges', None)
     fallback.pop('_prior_aliases', None)
     fallback.pop(key, None)
@@ -693,7 +716,8 @@ def _plausible_callable_target(data: bytes, load_addr: int, size: int,
 
 def _walk_overlay_function(data: bytes, load_addr: int, size: int,
                            entry: int, hard_cap: int,
-                           producer_ranges=()) -> dict:
+                           producer_ranges=(),
+                           allow_cross_producer_calls=True) -> dict:
     lo = load_addr
     hi = load_addr + size
 
@@ -723,6 +747,9 @@ def _walk_overlay_function(data: bytes, load_addr: int, size: int,
         target_range = producer_for(target)
         if source_range is not None and source_range == target_range:
             return True
+        if not allow_cross_producer_calls:
+            rejected_cross_producer_calls.add(target)
+            return False
         # Adjacent-producer composites are a byte envelope, not proof that a
         # call encoded by producer A is linked to arbitrary bytes at the same
         # address in every producer-B variant. Cross-boundary calls therefore
@@ -878,6 +905,7 @@ def classify_overlay_seeds(cap: dict, data: bytes, load_addr: int, size: int,
         raise RuntimeError('producer_ranges overlap')
 
     static_alias_ranges = set()
+    allow_cross_producer_calls = not bool(cap.get('strict_producer_ranges'))
     for raw_alias in cap.get('static_alias_ranges', []) or []:
         if not isinstance(raw_alias, dict):
             raise RuntimeError('static_alias_ranges entries must be objects')
@@ -929,7 +957,8 @@ def classify_overlay_seeds(cap: dict, data: bytes, load_addr: int, size: int,
     for i, entry in enumerate(pre_roots_sorted):
         hard_cap = pre_roots_sorted[i + 1] if i + 1 < len(pre_roots_sorted) else hi
         walk = _walk_overlay_function(
-            data, load_addr, size, entry, hard_cap, producer_ranges)
+            data, load_addr, size, entry, hard_cap, producer_ranges,
+            allow_cross_producer_calls)
         jump_table_targets.update(walk['jump_table_targets'])
 
     def impossible_entry_start(addr: int) -> bool:
@@ -1030,7 +1059,8 @@ def classify_overlay_seeds(cap: dict, data: bytes, load_addr: int, size: int,
             sorted_known = sorted(known)
             hard_cap = next((x for x in sorted_known if x > entry), hi)
             walk = _walk_overlay_function(
-                data, load_addr, size, entry, hard_cap, producer_ranges)
+                data, load_addr, size, entry, hard_cap, producer_ranges,
+                allow_cross_producer_calls)
             all_branch_targets.update(walk['branch_targets'])
             all_branch_targets.update(walk['jump_table_targets'])
             for target in sorted(walk['direct_jals']):
@@ -1084,7 +1114,8 @@ def classify_overlay_seeds(cap: dict, data: bytes, load_addr: int, size: int,
         for i, entry in enumerate(sorted_known):
             hard_cap = sorted_known[i + 1] if i + 1 < len(sorted_known) else hi
             walk = _walk_overlay_function(
-                data, load_addr, size, entry, hard_cap, producer_ranges)
+                data, load_addr, size, entry, hard_cap, producer_ranges,
+                allow_cross_producer_calls)
             covered |= walk['visited']
         promoted = sorted(a for a, r in included.items()
                           if r == 'DISPATCH_INTERIOR' and a not in covered
@@ -1109,7 +1140,8 @@ def classify_overlay_seeds(cap: dict, data: bytes, load_addr: int, size: int,
     for i, entry in enumerate(sorted_known):
         hard_cap = sorted_known[i + 1] if i + 1 < len(sorted_known) else hi
         walk = _walk_overlay_function(
-            data, load_addr, size, entry, hard_cap, producer_ranges)
+            data, load_addr, size, entry, hard_cap, producer_ranges,
+            allow_cross_producer_calls)
         all_branch_targets.update(walk['branch_targets'])
         all_branch_targets.update(walk['jump_table_targets'])
         all_branch_targets.update(walk['forward_branch_targets'])
@@ -2683,7 +2715,7 @@ def main():
                     fallback = optional_enrichment_fallback_capture(cap)
                     if fallback is not None:
                         print('  OPTIONAL ENRICHMENT AUDIT REJECTED; '
-                              'retrying direct-call roots only\n')
+                              'retrying conservative recipe\n')
                         _do_capture(fallback, region_coverage_cache,
                                     interior_frag_jobs, stats)
                         return
@@ -2768,7 +2800,7 @@ def main():
                     fallback = optional_enrichment_fallback_capture(cap)
                     if fallback is not None:
                         print('  OPTIONAL ENRICHMENT COMPILE REJECTED; '
-                              'retrying direct-call roots only\n')
+                              'retrying conservative recipe\n')
                         _do_capture(fallback, region_coverage_cache,
                                     interior_frag_jobs, stats)
                         return