[ConstraintSystem] Compute variables referenced by conjunction elements incrementally

Attempting to pre-compute a set of referenced type variables
upfront is incorrect because parameter(s) and/or result type
could be bound before conjunction is attempted. Let's compute
a set of referenced variables before each element gets attempted.

include/swift/Sema/Constraint.h

@@ -583,16 +583,14 @@ public:
       Optional<TrailingClosureMatching> trailingClosureMatching,
       ConstraintLocator *locator);
 
-  static Constraint *
+  static Constraint *createClosureBodyElement(ConstraintSystem &cs,
-  createClosureBodyElement(ConstraintSystem &cs, ASTNode node,
+                                              ASTNode node,
-                           ConstraintLocator *locator,
+                                              ConstraintLocator *locator);
-                           ArrayRef<TypeVariableType *> referencedVars = {});
 
-  static Constraint *
+  static Constraint *createClosureBodyElement(ConstraintSystem &cs,
-  createClosureBodyElement(ConstraintSystem &cs, ASTNode node,
+                                              ASTNode node,
-                           ContextualTypeInfo context,
+                                              ContextualTypeInfo context,
-                           ConstraintLocator *locator,
+                                              ConstraintLocator *locator);
-                           ArrayRef<TypeVariableType *> referencedVars = {});
 
   /// Determine the kind of constraint.
   ConstraintKind getKind() const { return Kind; }

include/swift/Sema/ConstraintSystem.h

@@ -2226,6 +2226,7 @@ public:
   friend class ComponentStep;
   friend class TypeVariableStep;
   friend class ConjunctionStep;
+  friend class ConjunctionElement;
   friend class RequirementFailure;
   friend class MissingMemberFailure;
 
@@ -5609,14 +5610,17 @@ public:
 
   bool attempt(ConstraintSystem &cs) const;
 
-  ArrayRef<TypeVariableType *> getReferencedVars() const {
-    return Element->getTypeVariables();
-  }
-
   void print(llvm::raw_ostream &Out, SourceManager *SM) const {
     Out << "conjunction element ";
     Element->print(Out, SM);
   }
+
+private:
+  /// Find type variables referenced by this conjunction element.
+  /// If this is a closure body element, it would look inside \c ASTNode.
+  void
+  findReferencedVariables(ConstraintSystem &cs,
+                          SmallPtrSetImpl<TypeVariableType *> &typeVars) const;
 };
 
 class TypeVariableBinding {

lib/Sema/CSClosure.cpp

@@ -29,46 +29,18 @@ class TypeVariableRefFinder : public ASTWalker {
   ConstraintSystem &CS;
   ASTNode Parent;
 
-  llvm::SmallSetVector<TypeVariableType *, 2> referencedVars;
+  llvm::SmallPtrSetImpl<TypeVariableType *> &ReferencedVars;
 
 public:
-  TypeVariableRefFinder(ConstraintSystem &cs, ASTNode parent)
+  TypeVariableRefFinder(
-      : CS(cs), Parent(parent) {}
+      ConstraintSystem &cs, ASTNode parent,
+      llvm::SmallPtrSetImpl<TypeVariableType *> &referencedVars)
+      : CS(cs), Parent(parent), ReferencedVars(referencedVars) {}
 
   std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
     if (auto *DRE = dyn_cast<DeclRefExpr>(expr)) {
-      if (auto type = CS.getTypeIfAvailable(DRE->getDecl())) {
+      if (auto type = CS.getTypeIfAvailable(DRE->getDecl()))
-        // The logic below is handling not-yet resolved parameter
+        inferVariables(type);
-        // types referenced in the body e.g. `$0` or `x`.
-        if (auto *typeVar = type->getAs<TypeVariableType>()) {
-          referencedVars.insert(typeVar);
-
-          // It is possible that contextual type of a parameter
-          // has been assigned to an anonymous of named argument
-          // early, to facilitate closure type checking. Such a
-          // type can have type variables inside e.g.
-          //
-          // func test<T>(_: (UnsafePointer<T>) -> Void) {}
-          //
-          // test { ptr in
-          //  ...
-          // }
-          //
-          // Type variable representing `ptr` in the body of
-          // this closure would be bound to `UnsafePointer<$T>`
-          // in this case, where `$T` is a type variable for a
-          // generic parameter `T`.
-          auto simplifiedTy =
-              CS.getFixedTypeRecursive(typeVar, /*wantRValue=*/false);
-
-          if (!simplifiedTy->isEqual(typeVar) &&
-              simplifiedTy->hasTypeVariable()) {
-            SmallPtrSet<TypeVariableType *, 4> typeVars;
-            simplifiedTy->getTypeVariables(typeVars);
-            referencedVars.insert(typeVars.begin(), typeVars.end());
-          }
-        }
-      }
     }
 
     return {true, expr};
@@ -80,17 +52,45 @@ public:
     // explicitly.
     if (isa<ReturnStmt>(stmt)) {
       if (auto *closure = getAsExpr<ClosureExpr>(Parent)) {
-        auto resultTy = CS.getClosureType(closure)->getResult();
+        inferVariables(CS.getClosureType(closure)->getResult());
-        if (auto *typeVar = resultTy->getAs<TypeVariableType>())
-          referencedVars.insert(typeVar);
       }
     }
 
     return {true, stmt};
   }
 
-  ArrayRef<TypeVariableType *> getReferencedVars() const {
+private:
-    return referencedVars.getArrayRef();
+  void inferVariables(Type type) {
+    auto *typeVar = type->getWithoutSpecifierType()->getAs<TypeVariableType>();
+    if (!typeVar)
+      return;
+
+    // Record the type variable itself because it has to
+    // be in scope even when already bound.
+    ReferencedVars.insert(typeVar);
+
+    // It is possible that contextual type of a parameter/result
+    // has been assigned to e.g. an anonymous or named argument
+    // early, to facilitate closure type checking. Such a
+    // type can have type variables inside e.g.
+    //
+    // func test<T>(_: (UnsafePointer<T>) -> Void) {}
+    //
+    // test { ptr in
+    //  ...
+    // }
+    //
+    // Type variable representing `ptr` in the body of
+    // this closure would be bound to `UnsafePointer<$T>`
+    // in this case, where `$T` is a type variable for a
+    // generic parameter `T`.
+    auto simplifiedTy = CS.getFixedTypeRecursive(typeVar, /*wantRValue=*/false);
+
+    if (!simplifiedTy->isEqual(typeVar) && simplifiedTy->hasTypeVariable()) {
+      SmallPtrSet<TypeVariableType *, 4> typeVars;
+      simplifiedTy->getTypeVariables(typeVars);
+      ReferencedVars.insert(typeVars.begin(), typeVars.end());
+    }
   }
 };
 
@@ -151,16 +151,8 @@ static void createConjunction(ConstraintSystem &cs,
     if (!isViableElement(element))
       continue;
 
-    TypeVariableRefFinder refFinder(cs, locator->getAnchor());
+    constraints.push_back(
-
+        Constraint::createClosureBodyElement(cs, element, context, elementLoc));
-    // Solvable elements have to bring any of the referenced
-    // outer context type variables into scope.
-    if (element.is<Decl *>() || element.is<StmtCondition *>() ||
-        element.is<Expr *>() || element.isStmt(StmtKind::Return))
-      element.walk(refFinder);
-
-    constraints.push_back(Constraint::createClosureBodyElement(
-        cs, element, context, elementLoc, refFinder.getReferencedVars()));
   }
 
   cs.addUnsolvedConstraint(Constraint::createConjunction(
@@ -1334,3 +1326,21 @@ bool ConstraintSystem::applySolutionToBody(Solution &solution,
   closure->setBodyState(ClosureExpr::BodyState::TypeCheckedWithSignature);
   return false;
 }
+
+void ConjunctionElement::findReferencedVariables(
+    ConstraintSystem &cs, SmallPtrSetImpl<TypeVariableType *> &typeVars) const {
+  auto referencedVars = Element->getTypeVariables();
+  typeVars.insert(referencedVars.begin(), referencedVars.end());
+
+  if (Element->getKind() != ConstraintKind::ClosureBodyElement)
+    return;
+
+  ASTNode element = Element->getClosureElement();
+  auto *locator = Element->getLocator();
+
+  TypeVariableRefFinder refFinder(cs, locator->getAnchor(), typeVars);
+
+  if (element.is<Decl *>() || element.is<StmtCondition *>() ||
+      element.is<Expr *>() || element.isStmt(StmtKind::Return))
+    element.walk(refFinder);
+}

lib/Sema/CSSolver.cpp

@@ -2313,6 +2313,15 @@ void DisjunctionChoice::propagateConversionInfo(ConstraintSystem &cs) const {
 }
 
 bool ConjunctionElement::attempt(ConstraintSystem &cs) const {
+  // First, let's bring all referenced variables into scope.
+  {
+    llvm::SmallPtrSet<TypeVariableType *, 4> referencedVars;
+    findReferencedVariables(cs, referencedVars);
+
+    for (auto *typeVar : referencedVars)
+      cs.addTypeVariable(typeVar);
+  }
+
   auto result = cs.simplifyConstraint(*Element);
   return result != ConstraintSystem::SolutionKind::Error;
 }

lib/Sema/CSStep.cpp

@@ -799,10 +799,6 @@ bool ConjunctionStep::attempt(const ConjunctionElement &element) {
     CS.applySolution(Solutions.pop_back_val());
   }
 
-  // Bring all of the referenced variables into scope.
-  for (auto *typeVar : element.getReferencedVars())
-    CS.addTypeVariable(typeVar);
-
   // Make sure that element is solved in isolation
   // by dropping all scoring information.
   CS.CurrentScore = Score();

lib/Sema/Constraint.cpp

@@ -346,8 +346,7 @@ Constraint *Constraint::clone(ConstraintSystem &cs) const {
                   getLocator());
 
   case ConstraintKind::ClosureBodyElement:
-    return createClosureBodyElement(cs, getClosureElement(), getLocator(),
+    return createClosureBodyElement(cs, getClosureElement(), getLocator());
-                                    getTypeVariables());
   }
 
   llvm_unreachable("Unhandled ConstraintKind in switch.");
@@ -1015,19 +1014,17 @@ Constraint *Constraint::createApplicableFunction(
   return constraint;
 }
 
-Constraint *Constraint::createClosureBodyElement(
+Constraint *Constraint::createClosureBodyElement(ConstraintSystem &cs,
-    ConstraintSystem &cs, ASTNode node, ConstraintLocator *locator,
+                                                 ASTNode node,
-    ArrayRef<TypeVariableType *> referencedVars) {
+                                                 ConstraintLocator *locator) {
-  return createClosureBodyElement(cs, node, ContextualTypeInfo(), locator,
+  return createClosureBodyElement(cs, node, ContextualTypeInfo(), locator);
-                                  referencedVars);
 }
 
-Constraint *Constraint::createClosureBodyElement(
+Constraint *Constraint::createClosureBodyElement(ConstraintSystem &cs,
-    ConstraintSystem &cs, ASTNode node, ContextualTypeInfo context,
+                                                 ASTNode node,
-    ConstraintLocator *locator, ArrayRef<TypeVariableType *> referencedVars) {
+                                                 ContextualTypeInfo context,
+                                                 ConstraintLocator *locator) {
   SmallPtrSet<TypeVariableType *, 4> typeVars;
-  typeVars.insert(referencedVars.begin(), referencedVars.end());
-
   unsigned size = totalSizeToAlloc<TypeVariableType *>(typeVars.size());
   void *mem = cs.getAllocator().Allocate(size, alignof(Constraint));
   return new (mem) Constraint(node, context, locator, typeVars);