[Constraint solver] Lazily populate the expression depth/index map.

Author: DougGregor

lib/Sema/ConstraintSystem.cpp

@@ -73,35 +73,6 @@ ExpressionTimer::~ExpressionTimer() {
       .highlight(E->getSourceRange());
 }
 
-/// Extend the given depth map by adding depths for all of the subexpressions
-/// of the given expression.
-static void extendDepthMap(
-   Expr *expr,
-   llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &depthMap) {
-  class RecordingTraversal : public ASTWalker {
-  public:
-    llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &DepthMap;
-    unsigned Depth = 0;
-
-    explicit RecordingTraversal(
-        llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &depthMap)
-        : DepthMap(depthMap) {}
-
-    std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
-      DepthMap[E] = {Depth, Parent.getAsExpr()};
-      Depth++;
-      return { true, E };
-    }
-
-    Expr *walkToExprPost(Expr *E) override {
-      Depth--;
-      return E;
-    }
-  };
-
-  RecordingTraversal traversal(depthMap);
-  expr->walk(traversal);
-}
 
 ConstraintSystem::ConstraintSystem(DeclContext *dc,
                                    ConstraintSystemOptions options,
@@ -111,7 +82,7 @@ ConstraintSystem::ConstraintSystem(DeclContext *dc,
     CG(*new ConstraintGraph(*this))
 {
   if (expr) {
-    extendDepthMap(expr, ExprWeights);
+    InputExprs.insert(expr);
   }
 
   assert(DC && "context required");
@@ -530,6 +501,54 @@ ConstraintSystem::getCalleeLocator(ConstraintLocator *locator,
   return getConstraintLocator(anchor);
 }
 
+/// Extend the given depth map by adding depths for all of the subexpressions
+/// of the given expression.
+static void extendDepthMap(
+   Expr *expr,
+   llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &depthMap) {
+  class RecordingTraversal : public ASTWalker {
+  public:
+    llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &DepthMap;
+    unsigned Depth = 0;
+
+    explicit RecordingTraversal(
+        llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &depthMap)
+        : DepthMap(depthMap) {}
+
+    std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
+      DepthMap[E] = {Depth, Parent.getAsExpr()};
+      Depth++;
+      return { true, E };
+    }
+
+    Expr *walkToExprPost(Expr *E) override {
+      Depth--;
+      return E;
+    }
+  };
+
+  RecordingTraversal traversal(depthMap);
+  expr->walk(traversal);
+}
+
+Optional<std::pair<unsigned, Expr *>> ConstraintSystem::getExprDepthAndParent(
+    Expr *expr) {
+  // Bring the set of expression weights up to date.
+  while (NumInputExprsInWeights < InputExprs.size()) {
+    extendDepthMap(InputExprs[NumInputExprsInWeights], ExprWeights);
+    ++NumInputExprsInWeights;
+  }
+
+  auto e = ExprWeights.find(expr);
+  if (e != ExprWeights.end())
+    return e->second;
+
+  if (baseCS && baseCS != this)
+    return baseCS->getExprDepthAndParent(expr);
+
+  return None;
+}
+
 Type ConstraintSystem::openUnboundGenericType(UnboundGenericType *unbound,
                                               ConstraintLocatorBuilder locator,
                                               OpenedTypeMap &replacements) {

lib/Sema/ConstraintSystem.h

@@ -1024,6 +1024,12 @@ class ConstraintSystem {
   unsigned CountDisjunctions = 0;
 
 private:
+  /// The set of expressions for which we have generated constraints.
+  llvm::SetVector<Expr *> InputExprs;
+
+  /// The number of input expressions whose parents and depths have
+  /// been entered into \c ExprWeights.
+  unsigned NumInputExprsInWeights = 0;
 
   llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> ExprWeights;
 
@@ -1982,29 +1988,22 @@ class ConstraintSystem {
   getConstraintLocator(const ConstraintLocatorBuilder &builder);
 
   /// Lookup and return parent associated with given expression.
-  Expr *getParentExpr(Expr *expr) const {
-    auto e = ExprWeights.find(expr);
-    if (e != ExprWeights.end())
-      return e->second.second;
-
-    if (baseCS && baseCS != this)
-      return baseCS->getParentExpr(expr);
-
+  Expr *getParentExpr(Expr *expr) {
+    if (auto result = getExprDepthAndParent(expr))
+      return result->second;
     return nullptr;
   }
 
   /// Retrieve the depth of the given expression.
-  Optional<unsigned> getExprDepth(Expr *expr) const {
-    auto e = ExprWeights.find(expr);
-    if (e != ExprWeights.end())
-      return e->second.first;
-
-    if (baseCS && baseCS != this)
-      return baseCS->getExprDepth(expr);
-
+  Optional<unsigned> getExprDepth(Expr *expr) {
+    if (auto result = getExprDepthAndParent(expr))
+      return result->first;
     return None;
   }
 
+  /// Retrieve the depth and parent expression of the given expression.
+  Optional<std::pair<unsigned, Expr *>> getExprDepthAndParent(Expr *expr);
+
   /// Returns a locator describing the callee for the anchor of a given locator.
   ///
   /// - For an unresolved dot/member anchor, this will be a locator describing