1 files changed, 23 insertions, 19 deletions

diff --git a/llvm/lib/Analysis/ConstantFolding.cpp b/llvm/lib/Analysis/ConstantFolding.cpp
index 8f67492971d..1a6a54b9a76 100644
--- a/llvm/lib/Analysis/ConstantFolding.cpp
+++ b/llvm/lib/Analysis/ConstantFolding.cpp
@@ -776,7 +776,9 @@ Constant *SymbolicallyEvaluateGEP(const GEPOperator *GEP,
           Res = ConstantExpr::getSub(Res, CE->getOperand(1));
           Res = ConstantExpr::getIntToPtr(Res, ResTy);
           if (auto *ResCE = dyn_cast<ConstantExpr>(Res))
-            Res = ConstantFoldConstantExpression(ResCE, DL, TLI);
+            if (auto *FoldedRes =
+                    ConstantFoldConstantExpression(ResCE, DL, TLI))
+              Res = FoldedRes;
           return Res;
         }
       }
@@ -985,7 +987,8 @@ Constant *llvm::ConstantFoldInstruction(Instruction *I, const DataLayout &DL,
         return nullptr;
       // Fold the PHI's operands.
       if (auto *NewC = dyn_cast<ConstantExpr>(C))
-        C = ConstantFoldConstantExpression(NewC, DL, TLI);
+        if (auto *FoldedC = ConstantFoldConstantExpression(NewC, DL, TLI))
+          C = FoldedC;
       // If the incoming value is a different constant to
       // the one we saw previously, then give up.
       if (CommonValue && C != CommonValue)
@@ -1007,12 +1010,9 @@ Constant *llvm::ConstantFoldInstruction(Instruction *I, const DataLayout &DL,
   for (const Use &OpU : I->operands()) {
     auto *Op = cast<Constant>(&OpU);
     // Fold the Instruction's operands.
-    if (auto *NewCE = dyn_cast<ConstantExpr>(Op)) {
-      auto *FoldedOp = ConstantFoldConstantExpression(NewCE, DL, TLI);
-      if (!FoldedOp)
-        return nullptr;
-      Op = FoldedOp;
-    }
+    if (auto *NewCE = dyn_cast<ConstantExpr>(Op))
+      if (auto *FoldedOp = ConstantFoldConstantExpression(NewCE, DL, TLI))
+        Op = FoldedOp;
 
     Ops.push_back(Op);
   }
@@ -1042,22 +1042,26 @@ Constant *llvm::ConstantFoldInstruction(Instruction *I, const DataLayout &DL,
 
 namespace {
 
-Constant *
-ConstantFoldConstantExpressionImpl(const ConstantExpr *CE, const DataLayout &DL,
-                                   const TargetLibraryInfo *TLI,
-                                   SmallPtrSetImpl<ConstantExpr *> &FoldedOps) {
+Constant *ConstantFoldConstantExpressionImpl(
+    const ConstantExpr *CE, const DataLayout &DL, const TargetLibraryInfo *TLI,
+    SmallDenseMap<ConstantExpr *, Constant *> &FoldedOps) {
   SmallVector<Constant *, 8> Ops;
   for (const Use &NewU : CE->operands()) {
     auto *NewC = cast<Constant>(&NewU);
     // Recursively fold the ConstantExpr's operands. If we have already folded
     // a ConstantExpr, we don't have to process it again.
     if (auto *NewCE = dyn_cast<ConstantExpr>(NewC)) {
-      if (FoldedOps.insert(NewCE).second){
-        auto *FoldedC =
-            ConstantFoldConstantExpressionImpl(NewCE, DL, TLI, FoldedOps);
-        if (!FoldedC)
-          return nullptr;
-        NewC = FoldedC;
+      auto It = FoldedOps.find(NewCE);
+      if (It == FoldedOps.end()) {
+        if (auto *FoldedC =
+                ConstantFoldConstantExpressionImpl(NewCE, DL, TLI, FoldedOps)) {
+          NewC = FoldedC;
+          FoldedOps.insert({NewCE, FoldedC});
+        } else {
+          FoldedOps.insert({NewCE, NewCE});
+        }
+      } else {
+        NewC = It->second;
       }
     }
     Ops.push_back(NewC);
@@ -1076,7 +1080,7 @@ ConstantFoldConstantExpressionImpl(const ConstantExpr *CE, const DataLayout &DL,
 Constant *llvm::ConstantFoldConstantExpression(const ConstantExpr *CE,
                                                const DataLayout &DL,
                                                const TargetLibraryInfo *TLI) {
-  SmallPtrSet<ConstantExpr *, 4> FoldedOps;
+  SmallDenseMap<ConstantExpr *, Constant *> FoldedOps;
   return ConstantFoldConstantExpressionImpl(CE, DL, TLI, FoldedOps);
 }