[AST] Augment getDepthMap with information about parent expressions

Which is very useful for the solver because otherwise it'd have to
compute and store this information twice.

Author: xedin

include/swift/AST/Expr.h

@@ -524,10 +524,10 @@ class alignas(8) Expr {
   /// the parent map.
   llvm::DenseMap<Expr *, Expr *> getParentMap();
 
-  /// Produce a mapping from each subexpression to its depth in the root
-  /// expression. The root expression has depth 0, its children have depth
-  /// 1, etc.
-  llvm::DenseMap<Expr *, unsigned> getDepthMap();
+  /// 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

lib/AST/Expr.cpp

@@ -698,17 +698,18 @@ llvm::DenseMap<Expr *, Expr *> Expr::getParentMap() {
   return parentMap;
 }
 
-llvm::DenseMap<Expr *, unsigned> Expr::getDepthMap() {
+llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> Expr::getDepthMap() {
   class RecordingTraversal : public ASTWalker {
   public:
-    llvm::DenseMap<Expr *, unsigned> &DepthMap;
+    llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &DepthMap;
     unsigned Depth = 0;
 
-    explicit RecordingTraversal(llvm::DenseMap<Expr *, unsigned> &depthMap)
-      : DepthMap(depthMap) { }
+    explicit RecordingTraversal(
+        llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &depthMap)
+        : DepthMap(depthMap) {}
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
-      DepthMap[E] = Depth;
+      DepthMap[E] = {Depth, Parent.getAsExpr()};
       Depth++;
       return { true, E };
     }
@@ -719,7 +720,7 @@ llvm::DenseMap<Expr *, unsigned> Expr::getDepthMap() {
     }
   };
 
-  llvm::DenseMap<Expr *, unsigned> depthMap;
+  llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> depthMap;
   RecordingTraversal traversal(depthMap);
   walk(traversal);
   return depthMap;

lib/Sema/CSRanking.cpp

@@ -772,7 +772,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 *, unsigned> &weights) {
+    llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &weights) {
   if (cs.TC.getLangOpts().DebugConstraintSolver) {
     auto &log = cs.getASTContext().TypeCheckerDebug->getStream();
     log.indent(cs.solverState->depth * 2)
@@ -806,7 +806,7 @@ SolutionCompareResult ConstraintSystem::compareSolutions(
     if (auto *anchor = locator->getAnchor()) {
       auto weight = weights.find(anchor);
       if (weight != weights.end())
-        return weight->getSecond() + 1;
+        return weight->getSecond().first + 1;
     }
 
     return 1;
@@ -1210,10 +1210,10 @@ SolutionCompareResult ConstraintSystem::compareSolutions(
                   : SolutionCompareResult::Incomparable;
 }
 
-Optional<unsigned>
-ConstraintSystem::findBestSolution(SmallVectorImpl<Solution> &viable,
-                                   llvm::DenseMap<Expr *, unsigned> &weights,
-                                   bool minimize) {
+Optional<unsigned> ConstraintSystem::findBestSolution(
+    SmallVectorImpl<Solution> &viable,
+    llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &weights,
+    bool minimize) {
   if (viable.empty())
     return None;
   if (viable.size() == 1)

lib/Sema/ConstraintSystem.cpp

@@ -2289,10 +2289,11 @@ bool ConstraintSystem::diagnoseAmbiguity(Expr *expr,
     if (it == indexMap.end())
       continue;
     unsigned index = it->second;
-    it = depthMap.find(anchor);
-    if (it == depthMap.end())
+
+    auto e = depthMap.find(anchor);
+    if (e == depthMap.end())
       continue;
-    unsigned depth = it->second;
+    unsigned depth = e->second.first;
 
     // If we don't have a name to hang on to, it'll be hard to diagnose this
     // overload.

lib/Sema/ConstraintSystem.h

@@ -1167,7 +1167,7 @@ class ConstraintSystem {
                 FreeTypeVariableBinding allowFreeTypeVariables);
     ~SolverState();
 
-    llvm::DenseMap<Expr *, unsigned> ExprWeights;
+    llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> ExprWeights;
 
     /// The constraint system.
     ConstraintSystem &CS;
@@ -1315,6 +1315,11 @@ class ConstraintSystem {
       return AllowFreeTypeVariables != FreeTypeVariableBinding::Disallow;
     }
 
+    Expr *getParentExpr(Expr *expr) {
+      const auto &e = ExprWeights.find(expr);
+      return (e != ExprWeights.end()) ? e->second.second : nullptr;
+    }
+
   private:
     /// The list of constraints that have been retired along the
     /// current path, this list is used in LIFO fashion when constraints
@@ -1563,9 +1568,10 @@ class ConstraintSystem {
   /// set of solutions should be filtered even if there is
   /// no single best solution, see `findBestSolution` for
   /// more details.
-  void filterSolutions(SmallVectorImpl<Solution> &solutions,
-                       llvm::DenseMap<Expr *, unsigned> &weights,
-                       bool minimize = false) {
+  void
+  filterSolutions(SmallVectorImpl<Solution> &solutions,
+                  llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &weights,
+                  bool minimize = false) {
     if (solutions.size() < 2)
       return;
 
@@ -3148,7 +3154,7 @@ class ConstraintSystem {
   static SolutionCompareResult
   compareSolutions(ConstraintSystem &cs, ArrayRef<Solution> solutions,
                    const SolutionDiff &diff, unsigned idx1, unsigned idx2,
-                   llvm::DenseMap<Expr *, unsigned> &weights);
+                   llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &weights);
 
 public:
   /// Increase the score of the given kind for the current (partial) solution
@@ -3172,9 +3178,10 @@ class ConstraintSystem {
   ///
   /// \returns The index of the best solution, or nothing if there was no
   /// best solution.
-  Optional<unsigned> findBestSolution(SmallVectorImpl<Solution> &solutions,
-                                      llvm::DenseMap<Expr *, unsigned> &weights,
-                                      bool minimize);
+  Optional<unsigned>
+  findBestSolution(SmallVectorImpl<Solution> &solutions,
+                   llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &weights,
+                   bool minimize);
 
   /// Apply a given solution to the expression, producing a fully
   /// type-checked expression.