Teach SimplifyCFG that (switch (select cond, X, Y)) is better expressed as a branch.

Based on a patch by Alistair Lynn.


git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@126647 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 164 insertions, 1 deletions

diff --git a/lib/Transforms/Utils/SimplifyCFG.cpp b/lib/Transforms/Utils/SimplifyCFG.cpp
index c6708857cb..72fa468db6 100644
--- a/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -1800,6 +1800,26 @@ static bool SimplifyTerminatorOnSelect(TerminatorInst *OldTerm, Value *Cond,
   return true;
 }
 
+// SimplifySwitchOnSelect - Replaces
+//   (switch (select cond, X, Y)) on constant X, Y
+// with a branch - conditional if X and Y lead to distinct BBs,
+// unconditional otherwise.
+static bool SimplifySwitchOnSelect(SwitchInst *SI, SelectInst *Select) {
+  // Check for constant integer values in the select.
+  ConstantInt *TrueVal = dyn_cast<ConstantInt>(Select->getTrueValue());
+  ConstantInt *FalseVal = dyn_cast<ConstantInt>(Select->getFalseValue());
+  if (!TrueVal || !FalseVal)
+    return false;
+
+  // Find the relevant condition and destinations.
+  Value *Condition = Select->getCondition();
+  BasicBlock *TrueBB = SI->getSuccessor(SI->findCaseValue(TrueVal));
+  BasicBlock *FalseBB = SI->getSuccessor(SI->findCaseValue(FalseVal));
+
+  // Perform the actual simplification.
+  return SimplifyTerminatorOnSelect(SI, Condition, TrueBB, FalseBB);
+}
+
 // SimplifyIndirectBrOnSelect - Replaces
 //   (indirectbr (select cond, blockaddress(@fn, BlockA),
 //                             blockaddress(@fn, BlockB)))
@@ -2309,7 +2329,12 @@ bool SimplifyCFGOpt::SimplifySwitch(SwitchInst *SI) {
   if (BasicBlock *OnlyPred = BB->getSinglePredecessor())
     if (SimplifyEqualityComparisonWithOnlyPredecessor(SI, OnlyPred))
       return SimplifyCFG(BB) | true;
-  
+
+  Value *Cond = SI->getCondition();
+  if (SelectInst *Select = dyn_cast<SelectInst>(Cond))
+    if (SimplifySwitchOnSelect(SI, Select))
+      return SimplifyCFG(BB) | true;
+
   // If the block only contains the switch, see if we can fold the block
   // away into any preds.
   BasicBlock::iterator BBI = BB->begin();
diff --git a/test/Transforms/SimplifyCFG/switch-on-const-select.ll b/test/Transforms/SimplifyCFG/switch-on-const-select.ll
new file mode 100644
index 0000000000..5494a651d4
--- /dev/null
+++ b/test/Transforms/SimplifyCFG/switch-on-const-select.ll
@@ -0,0 +1,138 @@
+; RUN: opt < %s -simplifycfg -S | FileCheck %s
+
+; Test basic folding to a conditional branch.
+define i32 @foo(i64 %x, i64 %y) nounwind {
+; CHECK: @foo
+entry:
+    %eq = icmp eq i64 %x, %y
+    br i1 %eq, label %b, label %switch
+switch:
+    %lt = icmp slt i64 %x, %y
+; CHECK: br i1 %lt, label %a, label %b
+    %qux = select i1 %lt, i32 0, i32 2
+    switch i32 %qux, label %bees [
+        i32 0, label %a
+        i32 1, label %b
+        i32 2, label %b
+    ]
+a:
+    tail call void @bees.a() nounwind
+    ret i32 1
+; CHECK: b:
+; CHECK-NEXT: %retval = phi i32 [ 0, %switch ], [ 2, %entry ]
+b:
+    %retval = phi i32 [0, %switch], [0, %switch], [2, %entry]
+    tail call void @bees.b() nounwind
+    ret i32 %retval
+; CHECK-NOT: bees:
+bees:
+    tail call void @llvm.trap() nounwind
+    unreachable
+}
+
+; Test basic folding to an unconditional branch.
+define i32 @bar(i64 %x, i64 %y) nounwind {
+; CHECK: @bar
+entry:
+; CHECK-NEXT: entry:
+; CHECK-NEXT: tail call void @bees.a() nounwind
+; CHECK-NEXT: ret i32 0
+    %lt = icmp slt i64 %x, %y
+    %qux = select i1 %lt, i32 0, i32 2
+    switch i32 %qux, label %bees [
+        i32 0, label %a
+        i32 1, label %b
+        i32 2, label %a
+    ]
+a:
+    %retval = phi i32 [0, %entry], [0, %entry], [1, %b]
+    tail call void @bees.a() nounwind
+    ret i32 0
+b:
+    tail call void @bees.b() nounwind
+    br label %a
+bees:
+    tail call void @llvm.trap() nounwind
+    unreachable
+}
+
+; Test the edge case where both values from the select are the default case.
+define void @bazz(i64 %x, i64 %y) nounwind {
+; CHECK: @bazz
+entry:
+; CHECK-NEXT: entry:
+; CHECK-NEXT: tail call void @bees.b() nounwind
+; CHECK-NEXT: ret void
+    %lt = icmp slt i64 %x, %y
+    %qux = select i1 %lt, i32 10, i32 12
+    switch i32 %qux, label %b [
+        i32 0, label %a
+        i32 1, label %bees
+        i32 2, label %bees
+    ]
+a:
+    tail call void @bees.a() nounwind
+    ret void
+b:
+    tail call void @bees.b() nounwind
+    ret void
+bees:
+    tail call void @llvm.trap()
+    unreachable
+}
+
+; Test the edge case where both values from the select are equal.
+define void @quux(i64 %x, i64 %y) nounwind {
+; CHECK: @quux
+entry:
+; CHECK-NEXT: entry:
+; CHECK-NEXT: tail call void @bees.a() nounwind
+; CHECK-NEXT: ret void
+    %lt = icmp slt i64 %x, %y
+    %qux = select i1 %lt, i32 0, i32 0
+    switch i32 %qux, label %b [
+        i32 0, label %a
+        i32 1, label %bees
+        i32 2, label %bees
+    ]
+a:
+    tail call void @bees.a() nounwind
+    ret void
+b:
+    tail call void @bees.b() nounwind
+    ret void
+bees:
+    tail call void @llvm.trap()
+    unreachable
+}
+
+; A final test, for phi node munging.
+define i32 @xyzzy(i64 %x, i64 %y) {
+; CHECK: @xyzzy
+entry:
+    %eq = icmp eq i64 %x, %y
+    br i1 %eq, label %r, label %cont
+cont:
+; CHECK: %lt = icmp slt i64 %x, %y
+    %lt = icmp slt i64 %x, %y
+; CHECK-NEXT: br i1 %lt, label %a, label %r
+    %qux = select i1 %lt, i32 0, i32 2
+    switch i32 %qux, label %bees [
+        i32 0, label %a
+        i32 1, label %r
+        i32 2, label %r
+    ]
+r:
+    %val = phi i32 [0, %entry], [1, %cont], [1, %cont]
+    ret i32 %val
+a:
+    ret i32 -1
+; CHECK-NOT: bees:
+bees:
+    tail call void @llvm.trap()
+    unreachable
+}
+
+declare void @llvm.trap() nounwind noreturn
+declare void @bees.a() nounwind
+declare void @bees.b() nounwind