From be7d32d30c8f9e723c553fd3a90e60ab779db663 Mon Sep 17 00:00:00 2001
From: Graine25 <riengupta@gmail.com>
Date: Wed, 8 Jul 2026 15:09:39 -0700
Subject: [PATCH] fix(codegen): manual switch table configuration in block
 discovery

---
 include/rex/codegen/function_scanner.h |  8 +++---
 src/codegen/function_scanner.cpp       | 39 ++++++++++++++++++++------
 src/codegen/phase_discover.cpp         |  3 +-
 3 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/include/rex/codegen/function_scanner.h b/include/rex/codegen/function_scanner.h
index ccd10fdb0..d5a3654c8 100644
--- a/include/rex/codegen/function_scanner.h
+++ b/include/rex/codegen/function_scanner.h
@@ -253,10 +253,10 @@ struct BlockDiscoveryResult {
  * @param knownFunctions Set of known function entry points (to detect tail calls)
  * @return BlockDiscoveryResult containing blocks, branches, and jump tables
  */
-BlockDiscoveryResult discoverBlocks(DecodedBinary& decoded, uint32_t entryPoint,
-                                    const CodeRegion& containingRegion,
-                                    const std::unordered_set<uint32_t>& knownFunctions,
-                                    uint32_t pdataSize = 0);
+BlockDiscoveryResult discoverBlocks(
+    DecodedBinary& decoded, uint32_t entryPoint, const CodeRegion& containingRegion,
+    const std::unordered_set<uint32_t>& knownFunctions, uint32_t pdataSize = 0,
+    const std::unordered_map<uint32_t, JumpTable>* manualSwitchTables = nullptr);
 
 //=============================================================================
 // Jump Table Detection
diff --git a/src/codegen/function_scanner.cpp b/src/codegen/function_scanner.cpp
index 4d812d0eb..8227711e8 100644
--- a/src/codegen/function_scanner.cpp
+++ b/src/codegen/function_scanner.cpp
@@ -1812,10 +1812,10 @@ std::optional<JumpTable> detectJumpTable(DecodedBinary& decoded, uint32_t bctrAd
 // Block Discovery
 //=============================================================================
 
-BlockDiscoveryResult discoverBlocks(DecodedBinary& decoded, uint32_t entryPoint,
-                                    const CodeRegion& containingRegion,
-                                    const std::unordered_set<uint32_t>& knownFunctions,
-                                    uint32_t pdataSize) {
+BlockDiscoveryResult discoverBlocks(
+    DecodedBinary& decoded, uint32_t entryPoint, const CodeRegion& containingRegion,
+    const std::unordered_set<uint32_t>& knownFunctions, uint32_t pdataSize,
+    const std::unordered_map<uint32_t, JumpTable>* manualSwitchTables) {
   BlockDiscoveryResult result;
   std::unordered_set<uint32_t> visited;
   std::unordered_set<uint32_t> blockStarts;
@@ -1916,20 +1916,41 @@ BlockDiscoveryResult discoverBlocks(DecodedBinary& decoded, uint32_t entryPoint,
           }
           // Do not break: continue scanning the fall-through path.
         } else if (insn->opcode == rex::codegen::ppc::Opcode::bcctr) {
-          // Unconditional bctr - try to detect jump table
+          // Unconditional bctr - prefer a manually configured table, then try
+          // automatic detection. Manual tables are authoritative because they
+          // are commonly needed when the compiler emits a table without an
+          // adjacent bounds check.
           REXCODEGEN_TRACE("discoverBlocks: bctr at 0x{:08X} in func 0x{:08X}, funcEnd=0x{:08X}",
                            addr, entryPoint, funcEnd);
-          auto jt = detectJumpTable(decoded, addr, containingRegion, entryPoint, funcEnd);
+          std::optional<JumpTable> jt;
+          if (manualSwitchTables) {
+            auto manualIt = manualSwitchTables->find(addr);
+            if (manualIt != manualSwitchTables->end()) {
+              jt = manualIt->second;
+              REXCODEGEN_TRACE(
+                  "discoverBlocks: using manual jump table at bctr 0x{:08X} with {} targets", addr,
+                  jt->targets.size());
+            }
+          }
+          if (!jt) {
+            jt = detectJumpTable(decoded, addr, containingRegion, entryPoint, funcEnd);
+          }
           if (jt) {
             REXCODEGEN_TRACE("discoverBlocks: detected jump table at bctr 0x{:08X} with {} targets",
                              addr, jt->targets.size());
             result.jumpTables.push_back(*jt);
-            // Jump table targets are definitionally part of this function
-            // Extend funcEnd if any target exceeds it (within region bounds)
-            // This handles out-of-line switch case code
             for (uint32_t t : jt->targets) {
               if (t == 0)
                 continue;  // sentinel from null-padding detection
+
+              // A jump table may tail-dispatch to another known function.
+              // Do not import that function's blocks into the current function.
+              if (t != entryPoint && knownFunctions.contains(t)) {
+                continue;
+              }
+
+              // Internal jump table targets are part of this function. Extend
+              // funcEnd for out-of-line case blocks within the code region.
               if (t >= funcEnd && t < containingRegion.end) {
                 funcEnd = t + 4;  // Extend to include this target
               }
diff --git a/src/codegen/phase_discover.cpp b/src/codegen/phase_discover.cpp
index cfae8b8f3..1c735482a 100644
--- a/src/codegen/phase_discover.cpp
+++ b/src/codegen/phase_discover.cpp
@@ -95,7 +95,8 @@ void discoverFunction(CodegenContext& ctx, uint32_t funcAddr,
   }
 
   // Pass pdataSize so forward branches within function extent are correctly identified
-  auto result = discoverBlocks(decoded, funcAddr, *region, knownFunctions, pdataSize);
+  auto result = discoverBlocks(decoded, funcAddr, *region, knownFunctions, pdataSize,
+                               &ctx.Config().switchTables);
 
   if (result.blocks.empty()) {
     REXCODEGEN_WARN("Analyze: no blocks found for function 0x{:08X}", funcAddr);