[IR] move helper function to replace undef constant (elements) with fixed constants

This is the NFC part of D69519.
We had this functionality locally in instcombine, but it can be used
elsewhere, so hoisting it to Constant class.

3 files changed, 31 insertions, 24 deletions

diff --git a/llvm/include/llvm/IR/Constant.h b/llvm/include/llvm/IR/Constant.h
index 2b6a6e4141b..b91100a0e89 100644
--- a/llvm/include/llvm/IR/Constant.h
+++ b/llvm/include/llvm/IR/Constant.h
@@ -188,6 +188,10 @@ public:
     return const_cast<Constant*>(
                       static_cast<const Constant *>(this)->stripPointerCasts());
   }
+
+  /// Try to replace undefined constant C or undefined elements in C with
+  /// Replacement. If no changes are made, the constant C is returned.
+  static Constant *replaceUndefsWith(Constant *C, Constant *Replacement);
 };
 
 } // end namespace llvm
diff --git a/llvm/lib/IR/Constants.cpp b/llvm/lib/IR/Constants.cpp
index f792f01efc1..ab4e478f700 100644
--- a/llvm/lib/IR/Constants.cpp
+++ b/llvm/lib/IR/Constants.cpp
@@ -31,6 +31,7 @@
 #include <algorithm>
 
 using namespace llvm;
+using namespace PatternMatch;
 
 //===----------------------------------------------------------------------===//
 //                              Constant Class
@@ -259,10 +260,9 @@ bool Constant::isElementWiseEqual(Value *Y) const {
   auto *Cy = dyn_cast<Constant>(Y);
   if (!Cy)
     return false;
-  return PatternMatch::match(ConstantExpr::getICmp(ICmpInst::Predicate::ICMP_EQ,
-                                                   const_cast<Constant *>(this),
-                                                   Cy),
-                             PatternMatch::m_One());
+  return match(ConstantExpr::getICmp(ICmpInst::Predicate::ICMP_EQ,
+                                     const_cast<Constant *>(this), Cy),
+               m_One());
 }
 
 bool Constant::containsUndefElement() const {
@@ -595,6 +595,27 @@ void Constant::removeDeadConstantUsers() const {
   }
 }
 
+Constant *Constant::replaceUndefsWith(Constant *C, Constant *Replacement) {
+  Type *Ty = C->getType();
+  if (C && match(C, m_Undef())) {
+    assert(Ty == Replacement->getType() && "Expected matching types");
+    return Replacement;
+  }
+
+  // Don't know how to deal with this constant.
+  if (!Ty->isVectorTy())
+    return C;
+
+  unsigned NumElts = Ty->getVectorNumElements();
+  SmallVector<Constant *, 32> NewC(NumElts);
+  for (unsigned i = 0; i != NumElts; ++i) {
+    Constant *EltC = C->getAggregateElement(i);
+    assert(EltC->getType() == Replacement->getType() &&
+           "Expected matching types");
+    NewC[i] = EltC && match(EltC, m_Undef()) ? Replacement : EltC;
+  }
+  return ConstantVector::get(NewC);
+}
 
 
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
index 64294838644..d28601ff9e2 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
@@ -138,24 +138,6 @@ Value *InstCombiner::reassociateShiftAmtsOfTwoSameDirectionShifts(
   return Ret;
 }
 
-// Try to replace `undef` constants in C with Replacement.
-static Constant *replaceUndefsWith(Constant *C, Constant *Replacement) {
-  if (C && match(C, m_Undef()))
-    return Replacement;
-
-  if (auto *CV = dyn_cast<ConstantVector>(C)) {
-    llvm::SmallVector<Constant *, 32> NewOps(CV->getNumOperands());
-    for (unsigned i = 0, NumElts = NewOps.size(); i != NumElts; ++i) {
-      Constant *EltC = CV->getOperand(i);
-      NewOps[i] = EltC && match(EltC, m_Undef()) ? Replacement : EltC;
-    }
-    return ConstantVector::get(NewOps);
-  }
-
-  // Don't know how to deal with this constant.
-  return C;
-}
-
 // If we have some pattern that leaves only some low bits set, and then performs
 // left-shift of those bits, if none of the bits that are left after the final
 // shift are modified by the mask, we can omit the mask.
@@ -216,7 +198,7 @@ dropRedundantMaskingOfLeftShiftInput(BinaryOperator *OuterShift,
     // completely unknown. Replace the the `undef` shift amounts with final
     // shift bitwidth to ensure that the value remains undef when creating the
     // subsequent shift op.
-    SumOfShAmts = replaceUndefsWith(
+    SumOfShAmts = Constant::replaceUndefsWith(
         SumOfShAmts, ConstantInt::get(SumOfShAmts->getType()->getScalarType(),
                                       ExtendedTy->getScalarSizeInBits()));
     auto *ExtendedSumOfShAmts = ConstantExpr::getZExt(SumOfShAmts, ExtendedTy);
@@ -241,7 +223,7 @@ dropRedundantMaskingOfLeftShiftInput(BinaryOperator *OuterShift,
     // bitwidth of innermost shift to ensure that the value remains undef when
     // creating the subsequent shift op.
     unsigned WidestTyBitWidth = WidestTy->getScalarSizeInBits();
-    ShAmtsDiff = replaceUndefsWith(
+    ShAmtsDiff = Constant::replaceUndefsWith(
         ShAmtsDiff, ConstantInt::get(ShAmtsDiff->getType()->getScalarType(),
                                      -WidestTyBitWidth));
     auto *ExtendedNumHighBitsToClear = ConstantExpr::getZExt(