[ConstraintSolver] NFC: refactor candidate solver to use constraint system arena

While shrinking we have to allocate containers for the reduced domains
for some of the candidates, it's currently done using permanent arena
of the `ASTContext` allocator. This patch changes candidate solver to
use arena associated with the parent constraint system, which significantly
limits lifetime of domain containers.

Author: xedin

lib/Sema/CSSolver.cpp

@@ -765,7 +765,7 @@ ConstraintSystem::solveSingle(FreeTypeVariableBinding allowFreeTypeVariables) {
 }
 
 bool ConstraintSystem::Candidate::solve(
-    llvm::SmallDenseSet<Expr *> &shrunkExprs) {
+    llvm::SmallDenseSet<OverloadSetRefExpr *> &shrunkExprs) {
   // Don't attempt to solve candidate if there is closure
   // expression involved, because it's handled specially
   // by parent constraint system (e.g. parameter lists).
@@ -879,7 +879,7 @@ bool ConstraintSystem::Candidate::solve(
 
 void ConstraintSystem::Candidate::applySolutions(
     llvm::SmallVectorImpl<Solution> &solutions,
-    llvm::SmallDenseSet<Expr *> &shrunkExprs) const {
+    llvm::SmallDenseSet<OverloadSetRefExpr *> &shrunkExprs) const {
   // A collection of OSRs with their newly reduced domains,
   // it's domains are sets because multiple solutions can have the same
   // choice for one of the type variables, and we want no duplication.
@@ -922,8 +922,10 @@ void ConstraintSystem::Candidate::applySolutions(
     if (OSR->getDecls().size() == choices.size()) continue;
 
     // Update the expression with the reduced domain.
-    MutableArrayRef<ValueDecl *> decls
-      = TC.Context.AllocateUninitialized<ValueDecl *>(choices.size());
+    MutableArrayRef<ValueDecl *> decls(
+                                Allocator.Allocate<ValueDecl *>(choices.size()),
+                                choices.size());
+
     std::uninitialized_copy(choices.begin(), choices.end(), decls.begin());
     OSR->setDecls(decls);
 
@@ -1238,7 +1240,7 @@ void ConstraintSystem::shrink(Expr *expr) {
   // so we can start solving them separately.
   expr->walk(collector);
 
-  llvm::SmallDenseSet<Expr *> shrunkExprs;
+  llvm::SmallDenseSet<OverloadSetRefExpr *> shrunkExprs;
   for (auto &candidate : collector.Candidates) {
     // If there are no results, let's forget everything we know about the
     // system so far. This actually is ok, because some of the expressions
@@ -1261,6 +1263,17 @@ void ConstraintSystem::shrink(Expr *expr) {
       });
     }
   }
+
+  // Once "shrinking" is done let's re-allocate final version of
+  // the candidate list to the permanent arena, so it could
+  // survive even after primary constraint system is destroyed.
+  for (auto &OSR : shrunkExprs) {
+    auto choices = OSR->getDecls();
+    auto decls = TC.Context.AllocateUninitialized<ValueDecl *>(choices.size());
+
+    std::uninitialized_copy(choices.begin(), choices.end(), decls.begin());
+    OSR->setDecls(decls);
+  }
 }
 
 ConstraintSystem::SolutionKind

lib/Sema/ConstraintSystem.h

@@ -1018,6 +1018,7 @@ class ConstraintSystem {
     Expr *E;
     TypeChecker &TC;
     DeclContext *DC;
+    llvm::BumpPtrAllocator &Allocator;
 
     // Contextual Information.
     Type CT;
@@ -1026,7 +1027,8 @@ class ConstraintSystem {
   public:
     Candidate(ConstraintSystem &cs, Expr *expr, Type ct = Type(),
               ContextualTypePurpose ctp = ContextualTypePurpose::CTP_Unused)
-        : E(expr), TC(cs.TC), DC(cs.DC), CT(ct), CTP(ctp) {}
+        : E(expr), TC(cs.TC), DC(cs.DC), Allocator(cs.Allocator), CT(ct),
+          CTP(ctp) {}
 
     /// \brief Return underlying expression.
     Expr *getExpr() const { return E; }
@@ -1038,7 +1040,7 @@ class ConstraintSystem {
     /// domains have been successfully shrunk so far.
     ///
     /// \returns true on solver failure, false otherwise.
-    bool solve(llvm::SmallDenseSet<Expr *> &shrunkExprs);
+    bool solve(llvm::SmallDenseSet<OverloadSetRefExpr *> &shrunkExprs);
 
     /// \brief Apply solutions found by solver as reduced OSR sets for
     /// for current and all of it's sub-expressions.
@@ -1048,14 +1050,16 @@ class ConstraintSystem {
     ///
     /// \param shrunkExprs The set of expressions which
     /// domains have been successfully shrunk so far.
-    void applySolutions(llvm::SmallVectorImpl<Solution> &solutions,
-                        llvm::SmallDenseSet<Expr *> &shrunkExprs) const;
+    void applySolutions(
+        llvm::SmallVectorImpl<Solution> &solutions,
+        llvm::SmallDenseSet<OverloadSetRefExpr *> &shrunkExprs) const;
 
     /// Check if attempt at solving of the candidate makes sense given
     /// the current conditions - number of shrunk domains which is related
     /// to the given candidate over the total number of disjunctions present.
-    static bool isTooComplexGiven(ConstraintSystem *const cs,
-                                  llvm::SmallDenseSet<Expr *> &shrunkExprs) {
+    static bool
+    isTooComplexGiven(ConstraintSystem *const cs,
+                      llvm::SmallDenseSet<OverloadSetRefExpr *> &shrunkExprs) {
       SmallVector<Constraint *, 8> disjunctions;
       cs->collectDisjunctions(disjunctions);
 
@@ -1066,7 +1070,11 @@ class ConstraintSystem {
           continue;
 
         if (auto *anchor = locator->getAnchor()) {
-          if (shrunkExprs.count(anchor) > 0)
+          auto *OSR = dyn_cast<OverloadSetRefExpr>(anchor);
+          if (!OSR)
+            continue;
+
+          if (shrunkExprs.count(OSR) > 0)
             --unsolvedDisjunctions;
         }
       }