[ConstraintSystem] Track AST depth information directly

Instead of storing information about expression depths in the
solver state (which gets recomputed for salvage) let's track
it directly in constraint system, which also gives solver
access to it when needed e.g. for fixes.

Author: xedin

lib/Sema/CSRanking.cpp

@@ -1210,10 +1210,9 @@ SolutionCompareResult ConstraintSystem::compareSolutions(
                   : SolutionCompareResult::Incomparable;
 }
 
-Optional<unsigned> ConstraintSystem::findBestSolution(
-    SmallVectorImpl<Solution> &viable,
-    llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &weights,
-    bool minimize) {
+Optional<unsigned>
+ConstraintSystem::findBestSolution(SmallVectorImpl<Solution> &viable,
+                                   bool minimize) {
   if (viable.empty())
     return None;
   if (viable.size() == 1)
@@ -1236,7 +1235,7 @@ Optional<unsigned> ConstraintSystem::findBestSolution(
   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, weights)) {
+    switch (compareSolutions(*this, viable, diff, i, bestIdx, ExprWeights)) {
     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.
@@ -1260,7 +1259,7 @@ Optional<unsigned> ConstraintSystem::findBestSolution(
     if (i == bestIdx)
       continue;
 
-    switch (compareSolutions(*this, viable, diff, bestIdx, i, weights)) {
+    switch (compareSolutions(*this, viable, diff, bestIdx, i, ExprWeights)) {
     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.
@@ -1312,7 +1311,7 @@ Optional<unsigned> ConstraintSystem::findBestSolution(
       if (losers[j])
         continue;
 
-      switch (compareSolutions(*this, viable, diff, i, j, weights)) {
+      switch (compareSolutions(*this, viable, diff, i, j, ExprWeights)) {
       case SolutionCompareResult::Identical:
         // FIXME: Dub one of these the loser arbitrarily?
         break;

lib/Sema/CSSolver.cpp

@@ -327,16 +327,12 @@ void truncate(llvm::SmallSetVector<T, N> &vec, unsigned newSize) {
 } // end anonymous namespace
 
 ConstraintSystem::SolverState::SolverState(
-    Expr *const expr, ConstraintSystem &cs,
-    FreeTypeVariableBinding allowFreeTypeVariables)
+    ConstraintSystem &cs, FreeTypeVariableBinding allowFreeTypeVariables)
     : CS(cs), AllowFreeTypeVariables(allowFreeTypeVariables) {
   assert(!CS.solverState &&
          "Constraint system should not already have solver state!");
   CS.solverState = this;
 
-  if (expr)
-    ExprWeights = expr->getDepthMap();
-
   ++NumSolutionAttempts;
   SolutionAttempt = NumSolutionAttempts;
 
@@ -498,12 +494,12 @@ Optional<Solution>
 ConstraintSystem::solveSingle(FreeTypeVariableBinding allowFreeTypeVariables,
                               bool allowFixes) {
 
-  SolverState state(nullptr, *this, allowFreeTypeVariables);
+  SolverState state(*this, allowFreeTypeVariables);
   state.recordFixes = allowFixes;
 
   SmallVector<Solution, 4> solutions;
   solve(solutions);
-  filterSolutions(solutions, state.ExprWeights);
+  filterSolutions(solutions);
 
   if (solutions.size() != 1)
     return Optional<Solution>();
@@ -538,7 +534,7 @@ bool ConstraintSystem::Candidate::solve(
   };
 
   // Allocate new constraint system for sub-expression.
-  ConstraintSystem cs(TC, DC, None);
+  ConstraintSystem cs(TC, DC, None, E);
   cs.baseCS = &BaseCS;
 
   // Set up expression type checker timer for the candidate.
@@ -589,7 +585,7 @@ bool ConstraintSystem::Candidate::solve(
   // Try to solve the system and record all available solutions.
   llvm::SmallVector<Solution, 2> solutions;
   {
-    SolverState state(E, cs, FreeTypeVariableBinding::Allow);
+    SolverState state(cs, FreeTypeVariableBinding::Allow);
 
     // Use solve which doesn't try to filter solution list.
     // Because we want the whole set of possible domain choices.
@@ -1179,7 +1175,7 @@ bool ConstraintSystem::solve(Expr *const expr,
                              SmallVectorImpl<Solution> &solutions,
                              FreeTypeVariableBinding allowFreeTypeVariables) {
   // Set up solver state.
-  SolverState state(expr, *this, allowFreeTypeVariables);
+  SolverState state(*this, allowFreeTypeVariables);
 
   // Solve the system.
   solve(solutions);
@@ -1200,7 +1196,7 @@ bool ConstraintSystem::solve(Expr *const expr,
   // a single best solution to use, if not explicitly disabled
   // by constraint system options.
   if (!retainAllSolutions())
-    filterSolutions(solutions, state.ExprWeights);
+    filterSolutions(solutions);
 
   // We fail if there is no solution or the expression was too complex.
   return solutions.empty() || getExpressionTooComplex(solutions);

lib/Sema/CSStep.h

@@ -227,7 +227,7 @@ class SolverStep {
 
   void filterSolutions(SmallVectorImpl<Solution> &solutions, bool minimize) {
     if (!CS.retainAllSolutions())
-      CS.filterSolutions(solutions, CS.solverState->ExprWeights, minimize);
+      CS.filterSolutions(solutions, minimize);
   }
 
   /// Check whether constraint solver is running in "debug" mode,

lib/Sema/ConstraintSystem.cpp

@@ -70,11 +70,15 @@ ExpressionTimer::~ExpressionTimer() {
 }
 
 ConstraintSystem::ConstraintSystem(TypeChecker &tc, DeclContext *dc,
-                                   ConstraintSystemOptions options)
+                                   ConstraintSystemOptions options,
+                                   Expr *expr)
   : TC(tc), DC(dc), Options(options),
     Arena(tc.Context, Allocator),
     CG(*new ConstraintGraph(*this))
 {
+  if (expr)
+    ExprWeights = expr->getDepthMap();
+
   assert(DC && "context required");
 }
 
@@ -2063,16 +2067,15 @@ bool ConstraintSystem::salvage(SmallVectorImpl<Solution> &viable, Expr *expr) {
 
   {
     // Set up solver state.
-    SolverState state(expr, *this, FreeTypeVariableBinding::Disallow);
+    SolverState state(*this, FreeTypeVariableBinding::Disallow);
     state.recordFixes = true;
 
     // Solve the system.
     solve(viable);
 
     // Check whether we have a best solution; this can happen if we found
     // a series of fixes that worked.
-    if (auto best = findBestSolution(viable, state.ExprWeights,
-                                     /*minimize=*/true)) {
+    if (auto best = findBestSolution(viable, /*minimize=*/true)) {
       if (*best != 0)
         viable[0] = std::move(viable[*best]);
       viable.erase(viable.begin() + 1, viable.end());

lib/Sema/ConstraintSystem.h

@@ -962,6 +962,8 @@ class ConstraintSystem {
 
 private:
 
+  llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> ExprWeights;
+
   /// Allocator used for all of the related constraint systems.
   llvm::BumpPtrAllocator Allocator;
 
@@ -1163,12 +1165,10 @@ class ConstraintSystem {
 
   /// Describes the current solver state.
   struct SolverState {
-    SolverState(Expr *const expr, ConstraintSystem &cs,
+    SolverState(ConstraintSystem &cs,
                 FreeTypeVariableBinding allowFreeTypeVariables);
     ~SolverState();
 
-    llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> ExprWeights;
-
     /// The constraint system.
     ConstraintSystem &CS;
 
@@ -1315,11 +1315,6 @@ 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
@@ -1518,7 +1513,8 @@ class ConstraintSystem {
   };
 
   ConstraintSystem(TypeChecker &tc, DeclContext *dc,
-                   ConstraintSystemOptions options);
+                   ConstraintSystemOptions options,
+                   Expr *expr = nullptr);
   ~ConstraintSystem();
 
   /// Retrieve the type checker associated with this constraint system.
@@ -1570,12 +1566,11 @@ class ConstraintSystem {
   /// more details.
   void
   filterSolutions(SmallVectorImpl<Solution> &solutions,
-                  llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &weights,
                   bool minimize = false) {
     if (solutions.size() < 2)
       return;
 
-    if (auto best = findBestSolution(solutions, weights, minimize)) {
+    if (auto best = findBestSolution(solutions, minimize)) {
       if (*best != 0)
         solutions[0] = std::move(solutions[*best]);
       solutions.erase(solutions.begin() + 1, solutions.end());
@@ -3150,7 +3145,6 @@ class ConstraintSystem {
   /// \param diff The differences among the solutions.
   /// \param idx1 The index of the first solution.
   /// \param idx2 The index of the second solution.
-  /// \param weights The weights of the sub-expressions used for ranking.
   static SolutionCompareResult
   compareSolutions(ConstraintSystem &cs, ArrayRef<Solution> solutions,
                    const SolutionDiff &diff, unsigned idx1, unsigned idx2,
@@ -3169,7 +3163,6 @@ class ConstraintSystem {
   /// solution.
   ///
   /// \param solutions The set of viable solutions to consider.
-  /// \param weights The weights of the sub-expressions used for ranking.
   ///
   /// \param minimize If true, then in the case where there is no single
   /// best solution, minimize the set of solutions by removing any solutions
@@ -3180,7 +3173,6 @@ class ConstraintSystem {
   /// best solution.
   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

lib/Sema/TypeCheckConstraints.cpp

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