[SCEV] Slightly generalize getRangeViaFactoring

Building on the previous change, this generalizes
ScalarEvolution::getRangeViaFactoring to work with
{Ext(C?A:B)+k0,+,Ext(C?A:B)+k1} where Ext can be a zero extend, sign
extend or truncate operation, and k0 and k1 are constants.

llvm-svn: 262979

1 files changed, 18 insertions, 13 deletions

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index aa5bfad5356..b8bf0f1cb9e 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -4557,17 +4557,6 @@ ConstantRange ScalarEvolution::getRangeViaFactoring(const SCEV *Start,
                                                     const SCEV *Step,
                                                     const SCEV *MaxBECount,
                                                     unsigned BitWidth) {
-  APInt Offset(BitWidth, 0);
-
-  if (auto *SA = dyn_cast<SCEVAddExpr>(Start)) {
-    // Peel off a constant offset, if possible.  In the future we could consider
-    // being smarter here and handle {Start+Step,+,Step} too.
-    if (SA->getNumOperands() != 2 || !isa<SCEVConstant>(SA->getOperand(0)))
-      return ConstantRange(BitWidth, /* isFullSet = */ true);
-    Offset = cast<SCEVConstant>(SA->getOperand(0))->getAPInt();
-    Start = SA->getOperand(1);
-  }
-
   //    RangeOf({C?A:B,+,C?P:Q}) == RangeOf(C?{A,+,P}:{B,+,Q})
   // == RangeOf({A,+,P}) union RangeOf({B,+,Q})
 
@@ -4579,10 +4568,22 @@ ConstantRange ScalarEvolution::getRangeViaFactoring(const SCEV *Start,
     explicit SelectPattern(ScalarEvolution &SE, unsigned BitWidth,
                            const SCEV *S) {
       Optional<unsigned> CastOp;
+      APInt Offset(BitWidth, 0);
 
       assert(SE.getTypeSizeInBits(S->getType()) == BitWidth &&
              "Should be!");
 
+      // Peel off a constant offset:
+      if (auto *SA = dyn_cast<SCEVAddExpr>(S)) {
+        // In the future we could consider being smarter here and handle
+        // {Start+Step,+,Step} too.
+        if (SA->getNumOperands() != 2 || !isa<SCEVConstant>(SA->getOperand(0)))
+          return;
+
+        Offset = cast<SCEVConstant>(SA->getOperand(0))->getAPInt();
+        S = SA->getOperand(1);
+      }
+
       // Peel off a cast operation
       if (auto *SCast = dyn_cast<SCEVCastExpr>(S)) {
         CastOp = SCast->getSCEVType();
@@ -4622,6 +4623,10 @@ ConstantRange ScalarEvolution::getRangeViaFactoring(const SCEV *Start,
           FalseValue = FalseValue.sext(BitWidth);
           break;
         }
+
+      // Re-apply the constant offset we peeled off earlier
+      TrueValue += Offset;
+      FalseValue += Offset;
     }
 
     bool isRecognized() { return Condition != nullptr; }
@@ -4650,9 +4655,9 @@ 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 *TrueStart = this->getConstant(StartPattern.TrueValue);
   const SCEV *TrueStep = this->getConstant(StepPattern.TrueValue);
-  const SCEV *FalseStart = this->getConstant(StartPattern.FalseValue + Offset);
+  const SCEV *FalseStart = this->getConstant(StartPattern.FalseValue);
   const SCEV *FalseStep = this->getConstant(StepPattern.FalseValue);
 
   ConstantRange TrueRange =