[BOLT] Make order of jump table successors deterministic

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

Reviewed By: yota9

Differential Revision: https://reviews.llvm.org/D119672

5 files changed, 108 insertions, 18 deletions

diff --git a/bolt/include/bolt/Core/BinaryBasicBlock.h b/bolt/include/bolt/Core/BinaryBasicBlock.h
index 3fbcbcb6195e..7cde2ce6fbf0 100644
--- a/bolt/include/bolt/Core/BinaryBasicBlock.h
+++ b/bolt/include/bolt/Core/BinaryBasicBlock.h
@@ -31,6 +31,7 @@ class MCCodeEmitter;
 namespace bolt {
 
 class BinaryFunction;
+class JumpTable;
 
 class BinaryBasicBlock {
 public:
@@ -623,6 +624,10 @@ public:
   /// 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);
@@ -909,7 +914,12 @@ public:
     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);
diff --git a/bolt/lib/Core/BinaryBasicBlock.cpp b/bolt/lib/Core/BinaryBasicBlock.cpp
index e76ba3c7a813..5f3a3431fcd0 100644
--- a/bolt/lib/Core/BinaryBasicBlock.cpp
+++ b/bolt/lib/Core/BinaryBasicBlock.cpp
@@ -39,10 +39,10 @@ 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) {
@@ -351,6 +351,36 @@ void BinaryBasicBlock::removeDuplicateConditionalSuccessor(MCInst *CondBranch) {
   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;
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 556e4585822e..c341ee0dcb7c 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -1875,21 +1875,7 @@ bool BinaryFunction::postProcessIndirectBranches(
     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)
diff --git a/bolt/test/X86/Inputs/jump-table-pic.s b/bolt/test/X86/Inputs/jump-table-pic.s
new file mode 100644
index 000000000000..fcc94cbb370a
--- /dev/null
+++ b/bolt/test/X86/Inputs/jump-table-pic.s
@@ -0,0 +1,52 @@
+# 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
diff --git a/bolt/test/X86/jump-table-pic-order.test b/bolt/test/X86/jump-table-pic-order.test
new file mode 100644
index 000000000000..5edf9bf24744
--- /dev/null
+++ b/bolt/test/X86/jump-table-pic-order.test
@@ -0,0 +1,12 @@
+# 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]]