Merge pull request #28276 from DougGregor/constraint-solver-expr-ctor

[Constraint solver] Remove expression from the constructor.

Author: DougGregor

include/swift/AST/Expr.h

@@ -530,17 +530,6 @@ class alignas(8) Expr {
   /// the parent map.
   llvm::DenseMap<Expr *, Expr *> getParentMap();
 
-  /// Produce a mapping from each subexpression to its depth and parent,
-  /// in the root expression. The root expression has depth 0, its children have
-  /// depth 1, etc.
-  llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> getDepthMap();
-
-  /// Produce a mapping from each expression to its index according to a
-  /// preorder traversal of the expressions. The parent has index 0, its first
-  /// child has index 1, its second child has index 2 if the first child is a
-  /// leaf node, etc.
-  llvm::DenseMap<Expr *, unsigned> getPreorderIndexMap();
-
   SWIFT_DEBUG_DUMP;
   void dump(raw_ostream &OS, unsigned Indent = 0) const;
   void dump(raw_ostream &OS, llvm::function_ref<Type(const Expr *)> getType,

lib/AST/Expr.cpp

@@ -721,56 +721,6 @@ llvm::DenseMap<Expr *, Expr *> Expr::getParentMap() {
   return parentMap;
 }
 
-llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> Expr::getDepthMap() {
-  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;
-    }
-  };
-
-  llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> depthMap;
-  RecordingTraversal traversal(depthMap);
-  walk(traversal);
-  return depthMap;
-}
-
-llvm::DenseMap<Expr *, unsigned> Expr::getPreorderIndexMap() {
-  class RecordingTraversal : public ASTWalker {
-  public:
-    llvm::DenseMap<Expr *, unsigned> &IndexMap;
-    unsigned Index = 0;
-
-    explicit RecordingTraversal(llvm::DenseMap<Expr *, unsigned> &indexMap)
-      : IndexMap(indexMap) { }
-
-    std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
-      IndexMap[E] = Index;
-      Index++;
-      return { true, E };
-    }
-  };
-
-  llvm::DenseMap<Expr *, unsigned> indexMap;
-  RecordingTraversal traversal(indexMap);
-  walk(traversal);
-  return indexMap;
-}
-
 //===----------------------------------------------------------------------===//
 // Support methods for Exprs.
 //===----------------------------------------------------------------------===//

lib/Sema/CSGen.cpp

@@ -3754,6 +3754,8 @@ namespace {
 } // end anonymous namespace
 
 Expr *ConstraintSystem::generateConstraints(Expr *expr, DeclContext *dc) {
+  InputExprs.insert(expr);
+
   // Remove implicit conversions from the expression.
   expr = expr->walk(SanitizeExpr(*this));
 

lib/Sema/CSSolver.cpp

@@ -570,7 +570,7 @@ bool ConstraintSystem::Candidate::solve(
   };
 
   // Allocate new constraint system for sub-expression.
-  ConstraintSystem cs(DC, None, E);
+  ConstraintSystem cs(DC, None);
   cs.baseCS = &BaseCS;
 
   // Set up expression type checker timer for the candidate.

lib/Sema/ConstraintSystem.cpp

@@ -73,16 +73,13 @@ ExpressionTimer::~ExpressionTimer() {
       .highlight(E->getSourceRange());
 }
 
+
 ConstraintSystem::ConstraintSystem(DeclContext *dc,
-                                   ConstraintSystemOptions options,
-                                   Expr *expr)
+                                   ConstraintSystemOptions options)
   : Context(dc->getASTContext()), DC(dc), Options(options),
     Arena(dc->getASTContext(), Allocator),
     CG(*new ConstraintGraph(*this))
 {
-  if (expr)
-    ExprWeights = expr->getDepthMap();
-
   assert(DC && "context required");
 }
 
@@ -499,6 +496,57 @@ 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) {
+  // Check whether the parent has this information.
+  if (baseCS && baseCS != this) {
+    if (auto known = baseCS->getExprDepthAndParent(expr))
+      return *known;
+  }
+
+  // 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;
+
+  return None;
+}
+
 Type ConstraintSystem::openUnboundGenericType(UnboundGenericType *unbound,
                                               ConstraintLocatorBuilder locator,
                                               OpenedTypeMap &replacements) {
@@ -2683,6 +2731,29 @@ static DeclName getOverloadChoiceName(ArrayRef<OverloadChoice> choices) {
   return name;
 }
 
+/// Extend the given index map with all of the subexpressions in the given
+/// expression.
+static void extendPreorderIndexMap(
+    Expr *expr, llvm::DenseMap<Expr *, unsigned> &indexMap) {
+  class RecordingTraversal : public ASTWalker {
+  public:
+    llvm::DenseMap<Expr *, unsigned> &IndexMap;
+    unsigned Index = 0;
+
+    explicit RecordingTraversal(llvm::DenseMap<Expr *, unsigned> &indexMap)
+      : IndexMap(indexMap) { }
+
+    std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
+      IndexMap[E] = Index;
+      Index++;
+      return { true, E };
+    }
+  };
+
+  RecordingTraversal traversal(indexMap);
+  expr->walk(traversal);
+}
+
 bool ConstraintSystem::diagnoseAmbiguity(Expr *expr,
                                          ArrayRef<Solution> solutions) {
   // Produce a diff of the solutions.
@@ -2703,7 +2774,10 @@ bool ConstraintSystem::diagnoseAmbiguity(Expr *expr,
   // Heuristically, all other things being equal, we should complain about the
   // ambiguous expression that (1) has the most overloads, (2) is deepest, or
   // (3) comes earliest in the expression.
-  auto indexMap = expr->getPreorderIndexMap();
+  llvm::DenseMap<Expr *, unsigned> indexMap;
+  for (auto expr : InputExprs) {
+    extendPreorderIndexMap(expr, indexMap);
+  }
 
   for (unsigned i = 0, n = diff.overloads.size(); i != n; ++i) {
     auto &overload = diff.overloads[i];

lib/Sema/ConstraintSystem.h

@@ -1038,6 +1038,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;
 
@@ -1662,8 +1668,7 @@ class ConstraintSystem {
   };
 
   ConstraintSystem(DeclContext *dc,
-                   ConstraintSystemOptions options,
-                   Expr *expr = nullptr);
+                   ConstraintSystemOptions options);
   ~ConstraintSystem();
 
   /// Retrieve the type checker associated with this constraint system.
@@ -1986,29 +1991,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

lib/Sema/TypeCheckConstraints.cpp

@@ -2197,7 +2197,7 @@ Type TypeChecker::typeCheckExpressionImpl(Expr *&expr, DeclContext *dc,
   if (options.contains(TypeCheckExprFlags::SubExpressionDiagnostics))
     csOptions |= ConstraintSystemFlags::SubExpressionDiagnostics;
 
-  ConstraintSystem cs(dc, csOptions, expr);
+  ConstraintSystem cs(dc, csOptions);
   cs.baseCS = baseCS;
 
   // Verify that a purpose was specified if a convertType was.  Note that it is