XenonAnalyse: optional switch-aware block walker for forward-jumping switch tables

XenonAnalyse/function.cpp

@@ -38,7 +38,8 @@ size_t Function::SearchBlock(size_t address) const
     return -1;
 }
 
-Function Function::Analyze(const void* code, size_t size, size_t base)
+Function Function::Analyze(const void* code, size_t size, size_t base,
+                           const AnalyzerSwitchTableMap* switchMap)
 {
     Function fn{ base, 0 };
 
@@ -59,6 +60,11 @@ Function Function::Analyze(const void* code, size_t size, size_t base)
     blockStack.reserve(32);
     blockStack.emplace_back();
 
+    // Set when the walker pushes switch-label successor blocks via the
+    // switchMap. When set, the end-of-Analyze discontinuity erase is
+    // skipped — see the comment on the discontinuity pass below for why.
+    bool switchAwareTermination = false;
+
     #define RESTORE_DATA() if (!blockStack.empty()) data = (dataStart + ((blocks[blockStack.back()].base + blocks[blockStack.back()].size) / sizeof(*data))) - 1; // continue adds one
 
     // TODO: Branch fallthrough
@@ -182,6 +188,66 @@ Function Function::Analyze(const void* code, size_t size, size_t base)
                 {
                     // 5th bit of BO tells cpu to ignore the counter, which is a blr/bctr otherwise it's conditional
                     const bool conditional = !(PPC_BO(instruction) & 0x10);
+
+                    // Switch-aware branch: when this is an unconditional
+                    // bctr (xop 528) at a known-switch site, push every
+                    // label (plus the default) as a successor block
+                    // instead of terminating without successors.
+                    //
+                    // Without this branch, an unconditional bctr pops
+                    // the current block and adds no new blocks, so the
+                    // walker never reaches the switch-dispatched code.
+                    // The discontinuity pass at the end of this function
+                    // then erases every block past the first address
+                    // gap — which for a switch-dispatched function is
+                    // immediately after bctr+4, sweeping all label
+                    // blocks away. Downstream consumers (e.g., a
+                    // recompiler's "label in fn.base..fn.base+fn.size"
+                    // boundary check) then mis-flag the labels as out
+                    // of function range.
+                    //
+                    // Two guards mirror the existing pre-base-branch
+                    // handling for unconditional `b` (see the "Branches
+                    // before base are just tail calls" note above):
+                    //
+                    //   label < base:         skip. Labels that point
+                    //     before the function base are either malformed
+                    //     TOML entries or tail-call-style jumps; they
+                    //     do not extend `fn.size`.
+                    //
+                    //   label >= base + size: skip. Labels outside the
+                    //     caller's analysis window cannot be walked
+                    //     safely; emplacing them here would compute an
+                    //     out-of-buffer `data` pointer at the next
+                    //     RESTORE_DATA.
+                    if (!conditional && xop == 528 && switchMap)
+                    {
+                        auto it = switchMap->find(addr);
+                        if (it != switchMap->end())
+                        {
+                            switchAwareTermination = true;
+                            auto pushLabel = [&](uint32_t label)
+                            {
+                                if (label < base) return;
+                                if (label >= base + size) return;
+                                if (fn.SearchBlock(label) == -1)
+                                {
+                                    const size_t lBase = label - base;
+                                    blocks.emplace_back(lBase, 0);
+                                    DEBUG(blocks.back().parent = blockBase);
+                                    blockStack.emplace_back(blocks.size() - 1);
+                                }
+                            };
+                            for (uint32_t label : it->second.labels)
+                            {
+                                pushLabel(label);
+                            }
+                            pushLabel(it->second.defaultLabel);
+                            RESTORE_DATA();
+                            continue;
+                        }
+                    }
+
                     if (conditional)
                     {
                         // right block's just going to return
@@ -210,29 +276,42 @@ Function Function::Analyze(const void* code, size_t size, size_t base)
         }
     }
 
-    // Sort and invalidate discontinuous blocks
+    // Sort and invalidate discontinuous blocks.
+    //
+    // When switchAwareTermination is set, the walker pushed successor
+    // blocks for every label of at least one switch dispatch. Those
+    // label blocks are separated from the pre-bctr block by the jump-
+    // table bytes themselves (data region living in the code section),
+    // which creates a legitimate address gap. The generic discontinuity
+    // heuristic would erase every block past that gap, which is exactly
+    // the set of blocks we just worked to make reachable. Skip the
+    // erase in that case; sort still runs because fn.size below picks
+    // max(block.base + block.size) and the sort is cheap.
     if (blocks.size() > 1)
     {
         std::sort(blocks.begin(), blocks.end(), [](const Block& a, const Block& b)
         {
             return a.base < b.base;
         });
 
-        size_t discontinuity = -1;
-        for (size_t i = 0; i < blocks.size() - 1; i++)
+        if (!switchAwareTermination)
         {
-            if (blocks[i].base + blocks[i].size >= blocks[i + 1].base)
+            size_t discontinuity = -1;
+            for (size_t i = 0; i < blocks.size() - 1; i++)
             {
-                continue;
-            }
+                if (blocks[i].base + blocks[i].size >= blocks[i + 1].base)
+                {
+                    continue;
+                }
 
-            discontinuity = i + 1;
-            break;
-        }
+                discontinuity = i + 1;
+                break;
+            }
 
-        if (discontinuity != -1)
-        {
-            blocks.erase(blocks.begin() + discontinuity, blocks.end());
+            if (discontinuity != -1)
+            {
+                blocks.erase(blocks.begin() + discontinuity, blocks.end());
+            }
         }
     }
 

XenonAnalyse/function.h

@@ -1,6 +1,8 @@
 #pragma once
 
 #include <cstddef>
+#include <cstdint>
+#include <unordered_map>
 #include <vector>
 
 #ifdef _DEBUG
@@ -9,6 +11,29 @@
 #define DEBUG(X)
 #endif
 
+// Optional switch-table map consumed by the block walker in
+// Function::Analyze when it reaches an unconditional bctr at a known
+// switch-dispatch site. When the caller hands in a populated map, the
+// walker pushes every label (and the default) as a successor block
+// instead of terminating, which grows `fn.size` to include switch-
+// dispatched code that the legacy walker would discard via the end-
+// of-Analyze discontinuity pass.
+//
+// The map is keyed by the BCTR guest VA — NOT by the first-instruction-
+// of-pattern VA that XenonAnalyse's own switch-table TOML emits. The
+// caller (typically XenonRecomp) performs the switch-base → bctr-VA
+// offset math at map-construction time and hands us a pre-keyed map.
+//
+// Passing nullptr (the default) preserves the legacy walker behavior
+// byte-for-byte. No existing callers need to change.
+struct AnalyzerSwitchTable
+{
+    uint32_t defaultLabel{};
+    std::vector<uint32_t> labels{};
+};
+
+using AnalyzerSwitchTableMap = std::unordered_map<uint32_t, AnalyzerSwitchTable>;
+
 struct Function
 {
     struct Block
@@ -18,16 +43,16 @@ struct Function
         size_t projectedSize{ static_cast<size_t>(-1) }; // scratch
         DEBUG(size_t parent{});
 
-        Block() 
+        Block()
         {
         }
 
         Block(size_t base, size_t size)
-            : base(base), size(size) 
+            : base(base), size(size)
         {
         }
 
-        Block(size_t base, size_t size, size_t projectedSize) 
+        Block(size_t base, size_t size, size_t projectedSize)
             : base(base), size(size), projectedSize(projectedSize)
         {
         }
@@ -45,7 +70,8 @@ struct Function
         : base(base), size(size)
     {
     }
-    
+
     size_t SearchBlock(size_t address) const;
-    static Function Analyze(const void* code, size_t size, size_t base);
+    static Function Analyze(const void* code, size_t size, size_t base,
+                            const AnalyzerSwitchTableMap* switchMap = nullptr);
 };

XenonAnalyse/tests/README.md

@@ -0,0 +1,68 @@
+# XenonAnalyse tests
+
+Synthetic hand-crafted PPC bytecode fixtures for `Function::Analyze`'s
+switch-aware block walker.
+
+## Fixtures
+
+`switch_aware_walker_test.cpp` exercises three scenarios against a
+synthetic 4-case absolute-form switch function built in-source from
+hand-emitted big-endian instruction words:
+
+1. **Happy path** — `switchMap` populated. Walker pushes all 4 labels
+   + default as successor blocks, `fn.size` extends to cover through
+   the default block.
+2. **Null-map safety** — `switchMap = nullptr`. Walker uses legacy
+   pre-patch behavior; discontinuity pass erases the label blocks;
+   `fn.size` covers only the pre-bctr head.
+3. **Wrong-map miss** — `switchMap` populated, but with an entry for
+   an unrelated bctr VA. `switchMap->find(our_bctr)` returns `end()`;
+   walker falls through to legacy behavior.
+
+All three assert on `fn.size` landing in the expected range. The
+tests use `assert`-style `fprintf(stderr, "[FAIL ...]" ...) + return 1`
+rather than a framework, to keep the test self-contained (no new
+dependencies).
+
+## Running
+
+Compile against a built `LibXenonAnalyse` + `XenonUtils` + `disasm` +
+`fmt`. Example (from a VS 2022 Developer Command Prompt with LLVM in
+PATH, after building the rest of the tree):
+
+```
+cd XenonAnalyse/tests
+clang-cl /std:c++20 /EHsc /nologo ^
+  /I"..\" /I"..\..\XenonUtils" ^
+  switch_aware_walker_test.cpp ^
+  "..\..\build\XenonAnalyse\LibXenonAnalyse.lib" ^
+  "..\..\build\XenonUtils\XenonUtils.lib" ^
+  "..\..\build\thirdparty\disasm\disasm.lib" ^
+  "..\..\build\thirdparty\fmt\fmt.lib" ^
+  /link /OUT:walker_test.exe
+walker_test.exe
+```
+
+Expected:
+
+```
+[ok   happy-path] fn.base=0x10000, fn.size=0x38, blocks=7
+[ok   null-map ] fn.base=0x10000, fn.size=0x1C (pre-patch truncation preserved)
+[ok   wrong-map] fn.base=0x10000, fn.size=0x1C (unrelated map entry correctly ignored)
+
+All 3 tests PASSED
+```
+
+## CMake integration
+
+Not included in this commit. If the repository adopts a test-runner
+pattern, a minimal CMake shape would be:
+
+```cmake
+add_executable(xenonanalyse_switch_walker_test switch_aware_walker_test.cpp)
+target_link_libraries(xenonanalyse_switch_walker_test PRIVATE LibXenonAnalyse)
+add_test(NAME switch_walker COMMAND xenonanalyse_switch_walker_test)
+```
+
+Or as an opt-in build behind `-DXENONANALYSE_TESTS=ON`. Whatever
+convention the repository prefers.