[BOLT] Make order of jump table successors deterministic

[maksfb]

Summary

When a jump table is recovered in postProcessIndirectBranches(),
successors for the containing basic block are added in random order.
Make the order deterministic.

bolt/include/bolt/Core/BinaryBasicBlock.h

 #include <utility>
 
 namespace llvm {
 class MCCodeEmitter;
 
 namespace bolt {
 
 class BinaryFunction;
+class JumpTable;
 
 class BinaryBasicBlock {
 public:
   /// Profile execution information for a given edge in CFG.
   ///
   /// If MispredictedCount equals COUNT_INFERRED, then we have a profile
   /// data for a fall-through edge with a Count representing an inferred
   /// execution count, i.e. the count we calculated internally, not the one
public:
   /// from respective lists of predecessors.
   void removeAllSuccessors();
 
   /// Remove useless duplicate successors.  When the conditional
   /// successor is the same as the unconditional successor, we can
   /// remove the conditional successor and branch instruction.
   void removeDuplicateConditionalSuccessor(MCInst *CondBranch);
 
+  /// Update successors of the basic block based on the jump table instruction.
+  /// The block must end with a jump table instruction.
+  void updateJumpTableSuccessors();
+
   /// Test if BB is a predecessor of this block.
   bool isPredecessor(const BinaryBasicBlock *BB) const {
     auto Itr = std::find(Predecessors.begin(), Predecessors.end(), BB);
     return Itr != Predecessors.end();
   }
 
   /// Test if BB is a successor of this block.
   bool isSuccessor(const BinaryBasicBlock *BB) const {
public:
   uint32_t getOffset() const { return InputRange.first; }
 
   /// Get the index of this basic block.
   unsigned getIndex() const {
     assert(isValid());
     return Index;
   }
 
-  bool hasJumpTable() const;
+  /// Return jump table if the block contains a jump table instruction or
+  /// nullptr otherwise.
+  const JumpTable *getJumpTable() const;
+
+  /// Check if the block has a jump table instruction.
+  bool hasJumpTable() const { return getJumpTable() != nullptr; }
 
 private:
   void adjustNumPseudos(const MCInst &Inst, int Sign);
 
   template <typename Itr> void adjustNumPseudos(Itr Begin, Itr End, int Sign) {
     while (Begin != End)
       adjustNumPseudos(*Begin++, Sign);
   }

bolt/lib/Core/BinaryBasicBlock.cpp

 bool BinaryBasicBlock::isEntryPoint() const {
   return getParent()->isEntryPoint(*this);
 }
 
 bool BinaryBasicBlock::hasInstructions() const {
   return getParent()->hasInstructions();
 }
 
-bool BinaryBasicBlock::hasJumpTable() const {
+const JumpTable *BinaryBasicBlock::getJumpTable() const {
   const MCInst *Inst = getLastNonPseudoInstr();
   const JumpTable *JT = Inst ? Function->getJumpTable(*Inst) : nullptr;
-  return (JT != nullptr);
+  return JT;
 }
 
 void BinaryBasicBlock::adjustNumPseudos(const MCInst &Inst, int Sign) {
   BinaryContext &BC = Function->getBinaryContext();
   if (BC.MIB->isPseudo(Inst))
     NumPseudos += Sign;
 }
 
void BinaryBasicBlock::removeDuplicateConditionalSuccessor(MCInst *CondBranch) {
   Successors.push_back(Succ);
 
   uint64_t Count = COUNT_NO_PROFILE;
   if (CondBI.Count != COUNT_NO_PROFILE && UncondBI.Count != COUNT_NO_PROFILE)
     Count = CondBI.Count + UncondBI.Count;
   BranchInfo.push_back({Count, 0});
 }
 
+void BinaryBasicBlock::updateJumpTableSuccessors() {
+  const JumpTable *JT = getJumpTable();
+  assert(JT && "Expected jump table instruction.");
+
+  // Clear existing successors.
+  removeAllSuccessors();
+
+  // Generate the list of successors in deterministic order without duplicates.
+  SmallVector<BinaryBasicBlock *, 16> SuccessorBBs;
+  for (const MCSymbol *Label : JT->Entries) {
+    BinaryBasicBlock *BB = getFunction()->getBasicBlockForLabel(Label);
+    // Ignore __builtin_unreachable()
+    if (!BB) {
+      assert(Label == getFunction()->getFunctionEndLabel() &&
+             "JT label should match a block or end of function.");
+      continue;
+    }
+    SuccessorBBs.emplace_back(BB);
+  }
+  llvm::sort(SuccessorBBs,
+             [](const BinaryBasicBlock *BB1, const BinaryBasicBlock *BB2) {
+               return BB1->getInputOffset() < BB2->getInputOffset();
+             });
+  SuccessorBBs.erase(std::unique(SuccessorBBs.begin(), SuccessorBBs.end()),
+                     SuccessorBBs.end());
+
+  for (BinaryBasicBlock *BB : SuccessorBBs)
+    addSuccessor(BB);
+}
+
 void BinaryBasicBlock::adjustExecutionCount(double Ratio) {
   auto adjustedCount = [&](uint64_t Count) -> uint64_t {
     double NewCount = Count * Ratio;
     if (!NewCount && Count && (Ratio > 0.0))
       NewCount = 1;
     return NewCount;
   };
 

bolt/lib/Core/BinaryFunction.cpp

bool BinaryFunction::postProcessIndirectBranches(
   // If there's only one jump table, and one indirect jump, and no other
   // references, then we should be able to derive the jump table even if we
   // fail to match the pattern.
   if (HasUnknownControlFlow && NumIndirectJumps == 1 &&
       JumpTables.size() == 1 && LastIndirectJump) {
     BC.MIB->setJumpTable(*LastIndirectJump, LastJT, LastJTIndexReg, AllocId);
     HasUnknownControlFlow = false;
 
-    // re-populate successors based on the jump table.
-    std::set<const MCSymbol *> JTLabels;
-    LastIndirectJumpBB->removeAllSuccessors();
-    const JumpTable *JT = getJumpTableContainingAddress(LastJT);
-    for (const MCSymbol *Label : JT->Entries)
-      JTLabels.emplace(Label);
-    for (const MCSymbol *Label : JTLabels) {
-      BinaryBasicBlock *BB = getBasicBlockForLabel(Label);
-      // Ignore __builtin_unreachable()
-      if (!BB) {
-        assert(Label == getFunctionEndLabel() && "if no BB found, must be end");
-        continue;
-      }
-      LastIndirectJumpBB->addSuccessor(BB);
-    }
+    LastIndirectJumpBB->updateJumpTableSuccessors();
   }
 
   if (HasFixedIndirectBranch)
     return false;
 
   if (HasUnknownControlFlow && !BC.HasRelocations)
     return false;
 

bolt/test/X86/Inputs/jump-table-pic.s

+# Test case with a simple PIC-style jump table where the register containing
+# the jump table address is defined outside of the basic block containing the
+# jump on register.
+#
+# One of the destinations of the jump table points past the end of the function
+# similar to the code generated for __builtin_unreachable().
+
+  .text
+  .globl  main
+  .type main, @function
+main:
+  .cfi_startproc
+  pushq %rbx
+  .cfi_def_cfa_offset 16
+  .cfi_offset 3, -16
+  leaq  .LJUMPTABLE(%rip), %rcx
+.L13:
+  cmpl  $3, %edi
+  ja  .L2
+
+  movslq  (%rcx,%rdi,4), %rax
+  addq  %rcx, %rax
+  jmp *%rax
+
+.L12:
+  movq  %rdi, %rax
+  popq  %rbx
+  .cfi_remember_state
+  .cfi_def_cfa_offset 8
+  ret
+
+.L5:
+  .cfi_restore_state
+  addq  $9, %rdi
+  jmp .L12
+.L6:
+  addq  $8, %rdi
+  jmp .L12
+.L2:
+  addq  $2, %rdi
+  jmp .L12
+.LUNREACHABLE:
+  .cfi_endproc
+  .size main, .-main
+
+  .section  .rodata
+  .align 4
+.LJUMPTABLE:
+  .long .L2-.LJUMPTABLE
+  .long .L6-.LJUMPTABLE
+  .long .L5-.LJUMPTABLE
+  .long .LUNREACHABLE-.LJUMPTABLE

bolt/test/X86/jump-table-pic-order.test

+# Check that successors of a basic block with jump table are generated
+# in the same order as they appear in the input code.
+
+RUN: %clang %cflags %S/Inputs/jump-table-pic.s -o %t.exe -Wl,-q
+RUN: llvm-bolt %t.exe -strict -print-cfg -print-only=main -o /dev/null \
+RUN:   | FileCheck %s
+
+CHECK:      BB Layout   : {{.*, .*, .*,}} [[BB4to6:.*, .*, .*]]
+
+# Check that successors appear in the order matching the input layout.
+CHECK:      jmpq    *%rax # JUMPTABLE
+CHECK-NEXT: Successors: [[BB4to6]]