From 954fc5a62524eb16ddee565a3e861c126a55fad7 Mon Sep 17 00:00:00 2001
From: "Brian M. Rzycki" <brzycki@gmail.com>
Date: Mon, 8 Apr 2019 18:20:35 +0000
Subject: [PATCH] [JumpThreading] Fix incorrect fold conditional after
 indirectbr/callbr

Fixes bug 40992: https://bugs.llvm.org/show_bug.cgi?id=40992

There is potential for miscompiled code emitted from JumpThreading when
analyzing a block with one or more indirectbr or callbr predecessors. The
ProcessThreadableEdges() function incorrectly folds conditional branches
into an unconditional branch.

This patch prevents incorrect branch folding without fully pessimizing
other potential threading opportunities through the same basic block.

This IR shape was manually fed in via opt and is unclear if clang and the
full pass pipeline will ever emit similar code shapes.

Thanks to Matthias Liedtke for the bug report and simplified IR example.

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


git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@357930 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/Transforms/Scalar/JumpThreading.cpp       |  6 +--
 .../pr40992-indirectbr-folding.ll             | 44 +++++++++++++++++++
 2 files changed, 45 insertions(+), 5 deletions(-)
 create mode 100644 test/Transforms/JumpThreading/pr40992-indirectbr-folding.ll

diff --git a/lib/Transforms/Scalar/JumpThreading.cpp b/lib/Transforms/Scalar/JumpThreading.cpp
index 264ea3aa22a..123c8b9630c 100644
--- a/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/lib/Transforms/Scalar/JumpThreading.cpp
@@ -1606,7 +1606,6 @@ bool JumpThreadingPass::ProcessThreadableEdges(Value *Cond, BasicBlock *BB,
   Constant *OnlyVal = nullptr;
   Constant *MultipleVal = (Constant *)(intptr_t)~0ULL;
 
-  unsigned PredWithKnownDest = 0;
   for (const auto &PredValue : PredValues) {
     BasicBlock *Pred = PredValue.second;
     if (!SeenPreds.insert(Pred).second)
@@ -1643,9 +1642,6 @@ bool JumpThreadingPass::ProcessThreadableEdges(Value *Cond, BasicBlock *BB,
         OnlyVal = MultipleVal;
     }
 
-    // We know where this predecessor is going.
-    ++PredWithKnownDest;
-
     // If the predecessor ends with an indirect goto, we can't change its
     // destination. Same for CallBr.
     if (isa<IndirectBrInst>(Pred->getTerminator()) ||
@@ -1663,7 +1659,7 @@ bool JumpThreadingPass::ProcessThreadableEdges(Value *Cond, BasicBlock *BB,
   // not thread. By doing so, we do not need to duplicate the current block and
   // also miss potential opportunities in case we dont/cant duplicate.
   if (OnlyDest && OnlyDest != MultipleDestSentinel) {
-    if (PredWithKnownDest == (size_t)pred_size(BB)) {
+    if (BB->hasNPredecessors(PredToDestList.size())) {
       bool SeenFirstBranchToOnlyDest = false;
       std::vector <DominatorTree::UpdateType> Updates;
       Updates.reserve(BB->getTerminator()->getNumSuccessors() - 1);
diff --git a/test/Transforms/JumpThreading/pr40992-indirectbr-folding.ll b/test/Transforms/JumpThreading/pr40992-indirectbr-folding.ll
new file mode 100644
index 00000000000..b94d4c1b701
--- /dev/null
+++ b/test/Transforms/JumpThreading/pr40992-indirectbr-folding.ll
@@ -0,0 +1,44 @@
+; RUN: opt -S < %s -jump-threading | FileCheck %s
+
+; PR40992: Do not incorrectly fold %bb5 into an unconditional br to %bb7.
+;          Also verify we correctly thread %bb1 -> %bb7 when %c is false.
+
+define i32 @jtbr(i1 %v1, i1 %v2, i1 %v3) {
+; CHECK: bb0:
+bb0:
+  br label %bb1
+
+; CHECK: bb1:
+; CHECK-NEXT: and
+; CHECK-NEXT: br i1 %c, label %bb2, label %bb7
+bb1:
+  %c = and i1 %v1, %v2
+  br i1 %c, label %bb2, label %bb5
+
+; CHECK: bb2:
+; CHECK-NEXT: select
+; CHECK-NEXT: indirectbr i8* %ba, [label %bb3, label %bb5]
+bb2:
+  %ba = select i1 %v3, i8* blockaddress(@jtbr, %bb3), i8* blockaddress(@jtbr, %bb4)
+  indirectbr i8* %ba, [label %bb3, label %bb4]
+
+; CHECK: bb3:
+bb3:
+  br label %bb1
+
+; CHECK-NOT: bb4:
+bb4:
+  br label %bb5
+
+; CHECK: bb5:
+bb5:
+  br i1 %c, label %bb6, label %bb7
+
+; CHECK: bb6:
+bb6:
+  ret i32 0
+
+; CHECK: bb7:
+bb7:
+  ret i32 1
+}
-- 
2.50.1