Conservatively handle jump tables in split functions

Summary:
- Allow jump table entries to point to locations inside the function and its fragments.
Reasoning behind this is that jump table identification has the logic of stopping at entry which belongs to a function different from the one originally referencing jump table. This assumption is invalid for jump tables with entries pointing to both parent function and cold fragments, leading to "unclaimed PC-relative relocations" assertion.

- Add fragment identification heuristic based on function name regex and contiguous jump table entries.
Currently, parent-to-fragment relationship is set up based on interprocedural references – direct references from the parent function. These references don't include references through jump table.
Additionally, some fragments are only reachable through jump table. In that case, in order to fully consume jump table, add parent-to-fragment relationship during `analyzeJumpTable` using the following heuristics:
  1. Fragment is identified as such based on name (contains `.cold.` part), but
  2. Parent function is not set – no direct interprocedural references to that fragment, and
  3. Fragment has the name of the form <parent>.cold(.\d+)

* For split functions with jump table entries spanning parent and fragments, mark parent and all fragments as ignored.

(cherry picked from FBD24456904)

bolt/src/BinaryContext.cpp

@@ -12,6 +12,7 @@
 #include "BinaryContext.h"
 #include "BinaryEmitter.h"
 #include "BinaryFunction.h"
+#include "NameResolver.h"
 #include "ParallelUtilities.h"
 #include "Utils.h"
 #include "llvm/ADT/Twine.h"
@@ -27,6 +28,7 @@
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSymbol.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Regex.h"
 #include <functional>
 #include <iterator>
 
@@ -518,7 +520,7 @@ BinaryContext::analyzeMemoryAt(uint64_t Address, BinaryFunction &BF) {
 
 bool BinaryContext::analyzeJumpTable(const uint64_t Address,
                                      const JumpTable::JumpTableType Type,
-                                     const BinaryFunction &BF,
+                                     BinaryFunction &BF,
                                      const uint64_t NextJTAddress,
                                      JumpTable::OffsetsType *Offsets) {
   // Is one of the targets __builtin_unreachable?
@@ -532,6 +534,29 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
       Offsets->emplace_back(Offset);
   };
 
+  auto doesBelongToFunction = [&](const uint64_t Addr,
+                                  BinaryFunction *TargetBF) -> bool {
+    if (BF.containsAddress(Addr))
+      return true;
+    if (!TargetBF || !TargetBF->isFragment())
+      return false;
+    if (!TargetBF->getParentFragment()) {
+      // TargetBF is a fragment, but parent function is not registered.
+      // This means there's no direct jump between parent and fragment.
+      // Set parent link here in jump table analysis, based on function name
+      // matching heuristic:
+      // <fragment restored name> == <base restored name>.cold(.\d+)?
+      auto TgtName = TargetBF->getOneName();
+      auto BFName = BF.getOneName();
+      auto TgtNamePrefix = NameResolver::restore(TgtName);
+      auto BFNamePrefix = Regex::escape(NameResolver::restore(BFName));
+      auto BFNameRegex = Twine(BFNamePrefix, "\\.cold(\\.\\d+)?").str();
+      if (Regex(BFNameRegex).match(TgtNamePrefix))
+        registerFragment(TargetBF, BF);
+    }
+    return TargetBF->getParentFragment() == &BF;
+  };
+
   auto Section = getSectionForAddress(Address);
   if (!Section)
     return false;
@@ -571,17 +596,28 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
       continue;
     }
 
+    // Function or one of its fragments.
+    auto TargetBF = getBinaryFunctionContainingAddress(Value);
+
     // We assume that a jump table cannot have function start as an entry.
-    if (!BF.containsAddress(Value) || Value == BF.getAddress())
+    if (!doesBelongToFunction(Value, TargetBF) || Value == BF.getAddress())
       break;
 
     // Check there's an instruction at this offset.
-    if (BF.getState() == BinaryFunction::State::Disassembled &&
-        !BF.getInstructionAtOffset(Value - BF.getAddress()))
+    if (TargetBF->getState() == BinaryFunction::State::Disassembled &&
+        !TargetBF->getInstructionAtOffset(Value - TargetBF->getAddress())) {
       break;
+    }
 
-    addOffset(Value - BF.getAddress());
     ++NumRealEntries;
+
+    if (TargetBF == &BF) {
+      // Address inside the function.
+      addOffset(Value - TargetBF->getAddress());
+    } else {
+      // Address in split fragment.
+      BF.setHasSplitJumpTable(true);
+    }
   }
 
   // It's a jump table if the number of real entries is more than 1, or there's
@@ -591,6 +627,7 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
 }
 
 void BinaryContext::populateJumpTables() {
+  std::vector<BinaryFunction *> FuncsToSkip;
   for (auto JTI = JumpTables.begin(), JTE = JumpTables.end(); JTI != JTE;
        ++JTI) {
     auto *JT = JTI->second;
@@ -638,11 +675,21 @@ void BinaryContext::populateJumpTables() {
         DataPCRelocations.erase(DataPCRelocations.find(Address));
       }
     }
+
+    // Mark to skip the function and all its fragments.
+    if (BF.hasSplitJumpTable())
+      FuncsToSkip.push_back(&BF);
   }
 
   assert((!opts::StrictMode || !DataPCRelocations.size()) &&
          "unclaimed PC-relative relocations left in data\n");
   clearList(DataPCRelocations);
+  // Functions containing split jump tables need to be skipped with all
+  // fragments.
+  for (auto BF : FuncsToSkip) {
+    BF->setIgnored();
+    BF->ignoreFragments();
+  }
 }
 
 MCSymbol *BinaryContext::getOrCreateGlobalSymbol(uint64_t Address,

bolt/src/BinaryContext.h

@@ -433,7 +433,7 @@ class BinaryContext {
   /// could be partially populated if the jump table detection fails.
   bool analyzeJumpTable(const uint64_t Address,
                         const JumpTable::JumpTableType Type,
-                        const BinaryFunction &BF,
+                        BinaryFunction &BF,
                         const uint64_t NextJTAddress = 0,
                         JumpTable::OffsetsType *Offsets = nullptr);
 

bolt/src/BinaryFunction.h

@@ -313,6 +313,10 @@ class BinaryFunction {
   /// True if the original entry point was patched.
   bool IsPatched{false};
 
+  /// True if the function contains jump table with entries pointing to
+  /// locations in fragments.
+  bool HasSplitJumpTable{false};
+
   /// The address for the code for this function in codegen memory.
   /// Used for functions that are emitted in a dedicated section with a fixed
   /// address. E.g. for functions that are overwritten in-place.
@@ -1400,6 +1404,12 @@ class BinaryFunction {
     return IsPseudo;
   }
 
+  /// Return true if the function contains a jump table with entries pointing
+  /// to split fragments.
+  bool hasSplitJumpTable() const {
+    return HasSplitJumpTable;
+  }
+
   /// Return true if the original function code has all necessary relocations
   /// to track addresses of functions emitted to new locations.
   bool hasExternalRefRelocations() const {
@@ -1876,6 +1886,10 @@ class BinaryFunction {
     IsPatched = V;
   }
 
+  void setHasSplitJumpTable(bool V) {
+    HasSplitJumpTable = V;
+  }
+
   void setFolded(BinaryFunction *BF) {
     FoldedIntoFunction = BF;
   }
@@ -2480,6 +2494,12 @@ class BinaryFunction {
   FragmentInfo &cold() { return ColdFragment; }
 
   const FragmentInfo &cold() const { return ColdFragment; }
+
+  /// Mark child fragments as ignored.
+  void ignoreFragments() {
+    for (auto *Fragment : Fragments)
+      Fragment->setIgnored();
+  }
 };
 
 inline raw_ostream &operator<<(raw_ostream &OS,

bolt/test/X86/split_func_jump_table_fragment.s

@@ -0,0 +1,57 @@
+# This reproduces a bug with jump table identification where jump table has
+# entries pointing to code in function and its cold fragment.
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# RUN: strip --strip-unneeded %t.o
+# RUN: %host_cc %t.o -o %t.exe -Wl,-q
+# RUN: llvm-bolt %t.exe -o %t.out -lite=0 -v=1 2>&1 | FileCheck %s
+
+# CHECK-NOT: unclaimed PC-relative relocations left in data
+# CHECK: BOLT-INFO: marking main.cold.1 as a fragment of main
+  .text
+  .globl main
+  .type main, %function
+  .p2align 2
+main:
+LBB0:
+  andl $0xf, %ecx
+  cmpb $0x4, %cl
+  # exit through abort in main.cold.1, registers cold fragment the regular way
+  ja main.cold.1
+
+# jump table dispatch, jumping to label indexed by val in %ecx
+LBB1:
+  leaq JUMP_TABLE(%rip), %r8
+  movzbl %cl, %ecx
+  movq (%r8,%rcx,8), %rax
+  addq %r8, %rax
+  jmpq *%rax
+
+LBB2:
+  xorq %rax, %rax
+LBB3:
+  addq $0x8, %rsp
+  ret
+.size main, .-main
+
+  .globl main.cold.1
+  .type main.cold.1, %function
+  .p2align 2
+main.cold.1:
+  # load bearing nop: pad LBB4 so that it can't be treated
+  # as __builtin_unreachable by analyzeJumpTable
+  nop
+LBB4:
+  callq abort
+.size main.cold.1, .-main.cold.1
+
+  .rodata
+# jmp table, entries must be R_X86_64_PC32 relocs
+  .globl JUMP_TABLE
+JUMP_TABLE:
+  .quad LBB2-JUMP_TABLE
+  .quad LBB3-JUMP_TABLE
+  .quad LBB4-JUMP_TABLE
+  .quad LBB3-JUMP_TABLE

bolt/test/X86/split_func_jump_table_fragment_noparent.s

@@ -0,0 +1,59 @@
+# This reproduces a bug with jump table identification where jump table has
+# entries pointing to code in function and its cold fragment.
+# The fragment is only reachable through jump table.
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# RUN: strip --strip-unneeded %t.o
+# RUN: %host_cc %t.o -o %t.exe -Wl,-q
+# RUN: llvm-bolt %t.exe -o %t.out -lite=0 -v=1 2>&1 | FileCheck %s
+
+# CHECK-NOT: unclaimed PC-relative relocations left in data
+# CHECK: BOLT-INFO: marking main.cold.1 as a fragment of main
+  .text
+  .globl main
+  .type main, %function
+  .p2align 2
+main:
+LBB0:
+  andl $0xf, %ecx
+  cmpb $0x4, %cl
+  # exit through ret
+  ja LBB3
+
+# jump table dispatch, jumping to label indexed by val in %ecx
+LBB1:
+  leaq JUMP_TABLE(%rip), %r8
+  movzbl %cl, %ecx
+  movslq (%r8,%rcx,4), %rax
+  addq %rax, %r8
+  jmpq *%r8
+
+LBB2:
+  xorq %rax, %rax
+LBB3:
+  addq $0x8, %rsp
+  ret
+.size main, .-main
+
+# cold fragment is only reachable through jump table
+  .globl main.cold.1
+  .type main.cold.1, %function
+  .p2align 2
+main.cold.1:
+  # load bearing nop: pad LBB4 so that it can't be treated
+  # as __builtin_unreachable by analyzeJumpTable
+  nop
+LBB4:
+  callq abort
+.size main.cold.1, .-main.cold.1
+
+  .rodata
+# jmp table, entries must be R_X86_64_PC32 relocs
+  .globl JUMP_TABLE
+JUMP_TABLE:
+  .long LBB2-JUMP_TABLE
+  .long LBB3-JUMP_TABLE
+  .long LBB4-JUMP_TABLE
+  .long LBB3-JUMP_TABLE