diff --git a/llvm/lib/Analysis/ConstantFolding.cpp b/llvm/lib/Analysis/ConstantFolding.cpp --- a/llvm/lib/Analysis/ConstantFolding.cpp +++ b/llvm/lib/Analysis/ConstantFolding.cpp @@ -1121,36 +1121,58 @@ ConstantFoldConstantImpl(const Constant *C, const DataLayout &DL, const TargetLibraryInfo *TLI, SmallDenseMap &FoldedOps) { - if (!isa(C) && !isa(C)) - return const_cast(C); + SmallVector Worklist = {const_cast(C)}; + + // "Recursively" iterate over all users of the constant. Try to fold constant + // if all its operand were processed before. + while (!Worklist.empty()) { + Constant *CurrentC = Worklist.back(); + if (FoldedOps.count(CurrentC)) { + Worklist.pop_back(); + continue; + } - SmallVector Ops; - for (const Use &OldU : C->operands()) { - Constant *OldC = cast(&OldU); - Constant *NewC = OldC; - // Recursively fold the ConstantExpr's operands. If we have already folded - // a ConstantExpr, we don't have to process it again. - if (isa(OldC) || isa(OldC)) { - auto It = FoldedOps.find(OldC); - if (It == FoldedOps.end()) { - NewC = ConstantFoldConstantImpl(OldC, DL, TLI, FoldedOps); - FoldedOps.insert({OldC, NewC}); - } else { - NewC = It->second; + if (!isa(CurrentC) && !isa(CurrentC)) { + Worklist.pop_back(); + FoldedOps.try_emplace(CurrentC, CurrentC); + continue; + } + + SmallVector Ops; + for (const Use &U : CurrentC->operands()) { + Constant *OldC = cast(&U); + Constant *NewC = OldC; + if (isa(OldC) || isa(OldC)) { + auto It = FoldedOps.find(OldC); + if (It == FoldedOps.end()) { + NewC = nullptr; + Worklist.push_back(OldC); + } else { + NewC = It->second; + } } + if (NewC) + Ops.push_back(NewC); } - Ops.push_back(NewC); - } - if (auto *CE = dyn_cast(C)) { - if (Constant *Res = - ConstantFoldInstOperandsImpl(CE, CE->getOpcode(), Ops, DL, TLI)) - return Res; - return const_cast(C); - } + // If all operands of the constant were processed, try to fold them + if (Ops.size() == CurrentC->getNumOperands()) { + assert(Worklist.back() == CurrentC && + "Some operand of the Constant was added to the Worklist."); + Worklist.pop_back(); - assert(isa(C)); - return ConstantVector::get(Ops); + if (auto *CE = dyn_cast(CurrentC)) { + if (Constant *Res = + ConstantFoldInstOperandsImpl(CE, CE->getOpcode(), Ops, DL, TLI)) + FoldedOps.try_emplace(CurrentC, Res); + else + FoldedOps.try_emplace(CurrentC, CurrentC); + } else { + FoldedOps.try_emplace(CurrentC, ConstantVector::get(Ops)); + } + } + } + return FoldedOps.find(C)->second; } } // end anonymous namespace