1 files changed, 20 insertions, 1 deletions
diff --git a/llvm/include/llvm/Analysis/ScalarEvolutionExpressions.h b/llvm/include/llvm/Analysis/ScalarEvolutionExpressions.h
index ff24cafbe68..9414c7bb574 100644
--- a/llvm/include/llvm/Analysis/ScalarEvolutionExpressions.h
+++ b/llvm/include/llvm/Analysis/ScalarEvolutionExpressions.h
@@ -537,14 +537,33 @@ namespace llvm {
     T.visitAll(Root);
   }
 
-  /// Recursively visits a SCEV expression and re-writes it.
+  /// This visitor recursively visits a SCEV expression and re-writes it.
+  /// The result from each visit is cached, so it will return the same
+  /// SCEV for the same input.
   template<typename SC>
   class SCEVRewriteVisitor : public SCEVVisitor<SC, const SCEV *> {
   protected:
     ScalarEvolution &SE;
+    // Memoize the result of each visit so that we only compute once for
+    // the same input SCEV. This is to avoid redundant computations when
+    // a SCEV is referenced by multiple SCEVs. Without memoization, this
+    // visit algorithm would have exponential time complexity in the worst
+    // case, causing the compiler to hang on certain tests.
+    DenseMap<const SCEV *, const SCEV *> RewriteResults;
+
   public:
     SCEVRewriteVisitor(ScalarEvolution &SE) : SE(SE) {}
 
+    const SCEV *visit(const SCEV *S) {
+      auto It = RewriteResults.find(S);
+      if (It != RewriteResults.end())
+        return It->second;
+      auto *Result = SCEVVisitor<SC, const SCEV *>::visit(S);
+      assert(RewriteResults.insert({S, Result}).second &&
+             "Should insert a new entry");
+      return Result;
+    }
+
     const SCEV *visitConstant(const SCEVConstant *Constant) {
       return Constant;
     }