1 files changed, 51 insertions, 23 deletions

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 2f0347771ba..aa5bfad5356 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -4568,42 +4568,70 @@ ConstantRange ScalarEvolution::getRangeViaFactoring(const SCEV *Start,
     Start = SA->getOperand(1);
   }
 
-  if (!isa<SCEVUnknown>(Start) || !isa<SCEVUnknown>(Step))
-    // We don't have anything new to contribute in this case.
-    return ConstantRange(BitWidth, /* isFullSet = */ true);
-
   //    RangeOf({C?A:B,+,C?P:Q}) == RangeOf(C?{A,+,P}:{B,+,Q})
   // == RangeOf({A,+,P}) union RangeOf({B,+,Q})
 
   struct SelectPattern {
     Value *Condition = nullptr;
-    const APInt *TrueValue = nullptr;
-    const APInt *FalseValue = nullptr;
+    APInt TrueValue;
+    APInt FalseValue;
+
+    explicit SelectPattern(ScalarEvolution &SE, unsigned BitWidth,
+                           const SCEV *S) {
+      Optional<unsigned> CastOp;
+
+      assert(SE.getTypeSizeInBits(S->getType()) == BitWidth &&
+             "Should be!");
+
+      // Peel off a cast operation
+      if (auto *SCast = dyn_cast<SCEVCastExpr>(S)) {
+        CastOp = SCast->getSCEVType();
+        S = SCast->getOperand();
+      }
 
-    explicit SelectPattern(const SCEVUnknown *SU) {
       using namespace llvm::PatternMatch;
 
-      if (!match(SU->getValue(),
-                 m_Select(m_Value(Condition), m_APInt(TrueValue),
-                          m_APInt(FalseValue)))) {
+      auto *SU = dyn_cast<SCEVUnknown>(S);
+      const APInt *TrueVal, *FalseVal;
+      if (!SU ||
+          !match(SU->getValue(), m_Select(m_Value(Condition), m_APInt(TrueVal),
+                                          m_APInt(FalseVal)))) {
         Condition = nullptr;
-        TrueValue = FalseValue = nullptr;
+        return;
       }
-    }
 
-    bool isRecognized() {
-      assert(((Condition && TrueValue && FalseValue) ||
-              (!Condition && !TrueValue && !FalseValue)) &&
-             "Invariant: either all three are non-null or all three are null");
-      return TrueValue != nullptr;
+      TrueValue = *TrueVal;
+      FalseValue = *FalseVal;
+
+      // Re-apply the cast we peeled off earlier
+      if (CastOp.hasValue())
+        switch (*CastOp) {
+        default:
+          llvm_unreachable("Unknown SCEV cast type!");
+
+        case scTruncate:
+          TrueValue = TrueValue.trunc(BitWidth);
+          FalseValue = FalseValue.trunc(BitWidth);
+          break;
+        case scZeroExtend:
+          TrueValue = TrueValue.zext(BitWidth);
+          FalseValue = FalseValue.zext(BitWidth);
+          break;
+        case scSignExtend:
+          TrueValue = TrueValue.sext(BitWidth);
+          FalseValue = FalseValue.sext(BitWidth);
+          break;
+        }
     }
+
+    bool isRecognized() { return Condition != nullptr; }
   };
 
-  SelectPattern StartPattern(cast<SCEVUnknown>(Start));
+  SelectPattern StartPattern(*this, BitWidth, Start);
   if (!StartPattern.isRecognized())
     return ConstantRange(BitWidth, /* isFullSet = */ true);
 
-  SelectPattern StepPattern(cast<SCEVUnknown>(Step));
+  SelectPattern StepPattern(*this, BitWidth, Step);
   if (!StepPattern.isRecognized())
     return ConstantRange(BitWidth, /* isFullSet = */ true);
 
@@ -4622,10 +4650,10 @@ ConstantRange ScalarEvolution::getRangeViaFactoring(const SCEV *Start,
   // FIXME: without the explicit `this` receiver below, MSVC errors out with
   // C2352 and C2512 (otherwise it isn't needed).
 
-  const SCEV *TrueStart = this->getConstant(*StartPattern.TrueValue + Offset);
-  const SCEV *TrueStep = this->getConstant(*StepPattern.TrueValue);
-  const SCEV *FalseStart = this->getConstant(*StartPattern.FalseValue + Offset);
-  const SCEV *FalseStep = this->getConstant(*StepPattern.FalseValue);
+  const SCEV *TrueStart = this->getConstant(StartPattern.TrueValue + Offset);
+  const SCEV *TrueStep = this->getConstant(StepPattern.TrueValue);
+  const SCEV *FalseStart = this->getConstant(StartPattern.FalseValue + Offset);
+  const SCEV *FalseStep = this->getConstant(StepPattern.FalseValue);
 
   ConstantRange TrueRange =
       this->getRangeForAffineAR(TrueStart, TrueStep, MaxBECount, BitWidth);