[Constraint system] Add ConstraintSystem::getExprDepth() and use it for cleanup.

Rather than passing around or create depth maps at a few places in the
constraint solver, introduce getExprDepth() and use it consistently.

Author: DougGregor

lib/Sema/CSRanking.cpp

@@ -710,8 +710,7 @@ static Type getUnlabeledType(Type type, ASTContext &ctx) {
 
 SolutionCompareResult ConstraintSystem::compareSolutions(
     ConstraintSystem &cs, ArrayRef<Solution> solutions,
-    const SolutionDiff &diff, unsigned idx1, unsigned idx2,
-    llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &weights) {
+    const SolutionDiff &diff, unsigned idx1, unsigned idx2) {
   if (cs.getASTContext().LangOpts.DebugConstraintSolver) {
     auto &log = cs.getASTContext().TypeCheckerDebug->getStream();
     log.indent(cs.solverState->depth * 2)
@@ -743,9 +742,9 @@ SolutionCompareResult ConstraintSystem::compareSolutions(
 
   auto getWeight = [&](ConstraintLocator *locator) -> unsigned {
     if (auto *anchor = locator->getAnchor()) {
-      auto weight = weights.find(anchor);
-      if (weight != weights.end())
-        return weight->getSecond().first + 1;
+      auto weight = cs.getExprDepth(anchor);
+      if (weight)
+        return *weight + 1;
     }
 
     return 1;
@@ -1197,7 +1196,7 @@ ConstraintSystem::findBestSolution(SmallVectorImpl<Solution> &viable,
   SmallVector<bool, 16> losers(viable.size(), false);
   unsigned bestIdx = 0;
   for (unsigned i = 1, n = viable.size(); i != n; ++i) {
-    switch (compareSolutions(*this, viable, diff, i, bestIdx, ExprWeights)) {
+    switch (compareSolutions(*this, viable, diff, i, bestIdx)) {
     case SolutionCompareResult::Identical:
       // FIXME: Might want to warn about this in debug builds, so we can
       // find a way to eliminate the redundancy in the search space.
@@ -1221,7 +1220,7 @@ ConstraintSystem::findBestSolution(SmallVectorImpl<Solution> &viable,
     if (i == bestIdx)
       continue;
 
-    switch (compareSolutions(*this, viable, diff, bestIdx, i, ExprWeights)) {
+    switch (compareSolutions(*this, viable, diff, bestIdx, i)) {
     case SolutionCompareResult::Identical:
       // FIXME: Might want to warn about this in debug builds, so we can
       // find a way to eliminate the redundancy in the search space.
@@ -1273,7 +1272,7 @@ ConstraintSystem::findBestSolution(SmallVectorImpl<Solution> &viable,
       if (losers[j])
         continue;
 
-      switch (compareSolutions(*this, viable, diff, i, j, ExprWeights)) {
+      switch (compareSolutions(*this, viable, diff, i, j)) {
       case SolutionCompareResult::Identical:
         // FIXME: Dub one of these the loser arbitrarily?
         break;

lib/Sema/ConstraintSystem.cpp

@@ -2695,7 +2695,6 @@ 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 depthMap = expr->getDepthMap();
   auto indexMap = expr->getPreorderIndexMap();
 
   for (unsigned i = 0, n = diff.overloads.size(); i != n; ++i) {
@@ -2711,10 +2710,10 @@ bool ConstraintSystem::diagnoseAmbiguity(Expr *expr,
       continue;
     unsigned index = it->second;
 
-    auto e = depthMap.find(anchor);
-    if (e == depthMap.end())
+    auto optDepth = getExprDepth(anchor);
+    if (!optDepth)
       continue;
-    unsigned depth = e->second.first;
+    unsigned depth = *optDepth;
 
     // If we don't have a name to hang on to, it'll be hard to diagnose this
     // overload.

lib/Sema/ConstraintSystem.h

@@ -1991,6 +1991,18 @@ class ConstraintSystem {
     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);
+
+    return None;
+  }
+
   /// 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
@@ -3659,8 +3671,7 @@ class ConstraintSystem {
   /// \param idx2 The index of the second solution.
   static SolutionCompareResult
   compareSolutions(ConstraintSystem &cs, ArrayRef<Solution> solutions,
-                   const SolutionDiff &diff, unsigned idx1, unsigned idx2,
-                   llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &weights);
+                   const SolutionDiff &diff, unsigned idx1, unsigned idx2);
 
 public:
   /// Increase the score of the given kind for the current (partial) solution