[InstCombine] fix bug when offsetting case values of a switch (PR31260)

We could truncate the condition and then try to fold the add into the
original condition value causing wrong case constants to be used.

Move the offset transform ahead of the truncate transform and return
after each transform, so there's no chance of getting confused values.

Fix for:
https://llvm.org/bugs/show_bug.cgi?id=31260

llvm-svn: 289442

2 files changed, 20 insertions, 29 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index f96766db6df..0d29c1dceea 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2232,6 +2232,20 @@ Instruction *InstCombiner::visitBranchInst(BranchInst &BI) {
 
 Instruction *InstCombiner::visitSwitchInst(SwitchInst &SI) {
   Value *Cond = SI.getCondition();
+  Value *Op0;
+  ConstantInt *AddRHS;
+  if (match(Cond, m_Add(m_Value(Op0), m_ConstantInt(AddRHS)))) {
+    // Change 'switch (X+4) case 1:' into 'switch (X) case -3'.
+    for (SwitchInst::CaseIt CaseIter : SI.cases()) {
+      Constant *NewCase = ConstantExpr::getSub(CaseIter.getCaseValue(), AddRHS);
+      assert(isa<ConstantInt>(NewCase) &&
+             "Result of expression should be constant");
+      CaseIter.setValue(cast<ConstantInt>(NewCase));
+    }
+    SI.setCondition(Op0);
+    return &SI;
+  }
+
   unsigned BitWidth = cast<IntegerType>(Cond->getType())->getBitWidth();
   APInt KnownZero(BitWidth, 0), KnownOne(BitWidth, 0);
   computeKnownBits(Cond, KnownZero, KnownOne, 0, &SI);
@@ -2252,9 +2266,7 @@ Instruction *InstCombiner::visitSwitchInst(SwitchInst &SI) {
   // Shrink the condition operand if the new type is smaller than the old type.
   // This may produce a non-standard type for the switch, but that's ok because
   // the backend should extend back to a legal type for the target.
-  bool TruncCond = false;
   if (NewWidth > 0 && NewWidth < BitWidth) {
-    TruncCond = true;
     IntegerType *Ty = IntegerType::get(SI.getContext(), NewWidth);
     Builder->SetInsertPoint(&SI);
     Value *NewCond = Builder->CreateTrunc(Cond, Ty, "trunc");
@@ -2264,32 +2276,10 @@ Instruction *InstCombiner::visitSwitchInst(SwitchInst &SI) {
       APInt TruncatedCase = CaseIter.getCaseValue()->getValue().trunc(NewWidth);
       CaseIter.setValue(ConstantInt::get(SI.getContext(), TruncatedCase));
     }
-  }
-
-  Value *Op0 = nullptr;
-  ConstantInt *AddRHS = nullptr;
-  if (match(Cond, m_Add(m_Value(Op0), m_ConstantInt(AddRHS)))) {
-    // Change 'switch (X+4) case 1:' into 'switch (X) case -3'.
-    for (SwitchInst::CaseIt CaseIter : SI.cases()) {
-      ConstantInt *CaseVal = CaseIter.getCaseValue();
-      Constant *LHS = CaseVal;
-      if (TruncCond) {
-        LHS = LeadingKnownZeros
-                  ? ConstantExpr::getZExt(CaseVal, Cond->getType())
-                  : ConstantExpr::getSExt(CaseVal, Cond->getType());
-      }
-      Constant *NewCaseVal = ConstantExpr::getSub(LHS, AddRHS);
-      assert(isa<ConstantInt>(NewCaseVal) &&
-             "Result of expression should be constant");
-      CaseIter.setValue(cast<ConstantInt>(NewCaseVal));
-    }
-    SI.setCondition(Op0);
-    if (auto *CondI = dyn_cast<Instruction>(Cond))
-      Worklist.Add(CondI);
     return &SI;
   }
 
-  return TruncCond ? &SI : nullptr;
+  return nullptr;
 }
 
 Instruction *InstCombiner::visitExtractValueInst(ExtractValueInst &EV) {
diff --git a/llvm/test/Transforms/InstCombine/narrow-switch.ll b/llvm/test/Transforms/InstCombine/narrow-switch.ll
index e58b4956568..8f4c24b9deb 100644
--- a/llvm/test/Transforms/InstCombine/narrow-switch.ll
+++ b/llvm/test/Transforms/InstCombine/narrow-switch.ll
@@ -135,10 +135,11 @@ sw.default:
 define i8 @PR31260(i8 %x) {
 ; ALL-LABEL: @PR31260(
 ; ALL-NEXT:  entry:
-; ALL-NEXT:    [[T4:%.*]] = and i8 %x, 2
-; ALL-NEXT:    switch i8 [[T4]], label %exit [
-; ALL-NEXT:    i8 -128, label %case126
-; ALL-NEXT:    i8 -126, label %case124
+; ALL-NEXT:    [[TMP0:%.*]] = trunc i8 %x to i2
+; ALL-NEXT:    [[TRUNC:%.*]] = and i2 [[TMP0]], -2
+; ALL-NEXT:    switch i2 [[TRUNC]], label %exit [
+; ALL-NEXT:    i2 0, label %case126
+; ALL-NEXT:    i2 -2, label %case124
 ; ALL-NEXT:    ]
 ; ALL:       exit:
 ; ALL-NEXT:    ret i8 1