From 90e12ffba71f4d3bd617a1ffacec209c2026292a Mon Sep 17 00:00:00 2001
From: Yevgeny Rouban <yevgeny.rouban@azul.com>
Date: Mon, 17 Jun 2019 05:55:12 +0000
Subject: [PATCH] [SimplifyCFG] Fix prof branch_weights MD while removing
 unreachable switch cases

SimplifyCFG has a bug that results in inconsistent prof branch_weights metadata
if unreachable switch cases are removed. This patch fixes this bug by making use
of the newly introduced SwitchInstProfUpdateWrapper class (see patch D62122).
A new test is created.

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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@363527 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/Transforms/Utils/SimplifyCFG.cpp         |  9 +++----
 test/Transforms/SimplifyCFG/switch-profmd.ll | 35 ++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 4 deletions(-)
 create mode 100644 test/Transforms/SimplifyCFG/switch-profmd.ll

diff --git a/lib/Transforms/Utils/SimplifyCFG.cpp b/lib/Transforms/Utils/SimplifyCFG.cpp
index ec574eaa81d..a8d24378279 100644
--- a/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -4218,14 +4218,15 @@ bool SimplifyCFGOpt::SimplifyUnreachable(UnreachableInst *UI) {
         }
       }
     } else if (auto *SI = dyn_cast<SwitchInst>(TI)) {
-      for (auto i = SI->case_begin(), e = SI->case_end(); i != e;) {
+      SwitchInstProfUpdateWrapper SU(*SI);
+      for (auto i = SU->case_begin(), e = SU->case_end(); i != e;) {
         if (i->getCaseSuccessor() != BB) {
           ++i;
           continue;
         }
-        BB->removePredecessor(SI->getParent());
-        i = SI->removeCase(i);
-        e = SI->case_end();
+        BB->removePredecessor(SU->getParent());
+        i = SU.removeCase(i);
+        e = SU->case_end();
         Changed = true;
       }
     } else if (auto *II = dyn_cast<InvokeInst>(TI)) {
diff --git a/test/Transforms/SimplifyCFG/switch-profmd.ll b/test/Transforms/SimplifyCFG/switch-profmd.ll
new file mode 100644
index 00000000000..c17214e18f2
--- /dev/null
+++ b/test/Transforms/SimplifyCFG/switch-profmd.ll
@@ -0,0 +1,35 @@
+; RUN: opt -S -simplifycfg < %s | FileCheck %s
+
+declare i32 @f(i32 %val)
+
+; Check that eliminating cases with unreachable branches keeps
+; prof branch_weights metadata consistent with switch instruction.
+define i32 @test_switch_to_unreachable(i32 %val) {
+; CHECK-LABEL: test_switch_to_unreachable
+; CHECK: switch
+; CHECK-NOT: i32 0, label %on0
+  switch i32 %val, label %otherwise [
+    i32 0, label %on0
+    i32 1, label %on1
+    i32 2, label %on2
+  ], !prof !{!"branch_weights", i32 99, i32 0, i32 1, i32 2}
+; CHECK: !prof ![[MD0:[0-9]+]]
+
+otherwise:
+  %result = call i32 @f(i32 -1)
+  ret i32 %result
+
+on0:
+  unreachable
+  ret i32 125
+
+on1:
+  %result1 = call i32 @f(i32 -2)
+  ret i32 %result1
+
+on2:
+  %result2 = call i32 @f(i32 -3)
+  ret i32 %result2
+}
+
+; CHECK: ![[MD0]] = !{!"branch_weights", i32 99, i32 2, i32 1}
-- 
2.11.0