[CSGen] Type-check capture list together with closure body

Delay constraint generation for capture list until body of
the associated closure is resolved. This means that we can
unify capture checking with that of regular pattern bindings
for multi-statement closures.

Author: xedin

Diff for: lib/Sema/CSGen.cpp

@@ -2985,7 +2985,14 @@ namespace {
         }
       } collectVarRefs(CS);
 
-      closure->walk(collectVarRefs);
+      // Walk the capture list if this closure has one,  because it could
+      // reference declarations from the outer closure.
+      if (auto *captureList =
+              getAsExpr<CaptureListExpr>(CS.getParentExpr(closure))) {
+        captureList->walk(collectVarRefs);
+      } else {
+        closure->walk(collectVarRefs);
+      }
 
       auto inferredType = inferClosureType(closure);
       if (!inferredType || inferredType->hasError())
@@ -2994,15 +3001,6 @@ namespace {
       SmallVector<TypeVariableType *, 4> referencedVars{
           collectVarRefs.varRefs.begin(), collectVarRefs.varRefs.end()};
 
-      if (auto *captureList =
-              getAsExpr<CaptureListExpr>(CS.getParentExpr(closure))) {
-        for (const auto &capture : captureList->getCaptureList()) {
-          if (auto *typeVar =
-                  CS.getType(capture.getVar())->getAs<TypeVariableType>())
-            referencedVars.push_back(typeVar);
-        }
-      }
-
       CS.addUnsolvedConstraint(Constraint::create(
           CS, ConstraintKind::DefaultClosureType, closureType, inferredType,
           locator, referencedVars));
@@ -4114,13 +4112,6 @@ namespace {
       // Generate constraints for each of the entries in the capture list.
       if (auto captureList = dyn_cast<CaptureListExpr>(expr)) {
         TypeChecker::diagnoseDuplicateCaptureVars(captureList);
-
-        auto &CS = CG.getConstraintSystem();
-        for (const auto &capture : captureList->getCaptureList()) {
-          SyntacticElementTarget target(capture.PBD);
-          if (CS.generateConstraints(target))
-            return Action::Stop();
-        }
       }
 
       // Both multi- and single-statement closures now behave the same way

Diff for: lib/Sema/CSSyntacticElement.cpp

@@ -1012,7 +1012,30 @@ class SyntacticElementConstraintGenerator
   void visitBraceStmt(BraceStmt *braceStmt) {
     auto &ctx = cs.getASTContext();
 
+    CaptureListExpr *captureList = nullptr;
+    {
+      if (locator->directlyAt<ClosureExpr>()) {
+        auto *closure = castToExpr<ClosureExpr>(locator->getAnchor());
+        captureList = getAsExpr<CaptureListExpr>(cs.getParentExpr(closure));
+      }
+    }
+
     if (context.isSingleExpressionClosure(cs)) {
+      // Generate constraints for the capture list first.
+      //
+      // TODO: This should be a conjunction connected to
+      // the closure body to make sure that each capture
+      // is solved in isolation.
+      if (captureList) {
+        for (const auto &capture : captureList->getCaptureList()) {
+          SyntacticElementTarget target(capture.PBD);
+          if (cs.generateConstraints(target)) {
+            hadError = true;
+            return;
+          }
+        }
+      }
+
       for (auto node : braceStmt->getElements()) {
         if (auto expr = node.dyn_cast<Expr *>()) {
           auto generatedExpr = cs.generateConstraints(
@@ -1029,6 +1052,8 @@ class SyntacticElementConstraintGenerator
       return;
     }
 
+    SmallVector<ElementInfo, 4> elements;
+
     // If this brace statement represents a body of an empty or
     // multi-statement closure.
     if (locator->directlyAt<ClosureExpr>()) {
@@ -1054,10 +1079,24 @@ class SyntacticElementConstraintGenerator
             cs.getConstraintLocator(closure, ConstraintLocator::ClosureResult));
       }
 
+      // If this multi-statement closure has captures, let's solve
+      // them first.
+      if (captureList) {
+        for (const auto &capture : captureList->getCaptureList())
+          visitPatternBinding(capture.PBD, elements);
+      }
+
       // Let's not walk into the body if empty or multi-statement closure
       // doesn't participate in inference.
-      if (!cs.participatesInInference(closure))
+      if (!cs.participatesInInference(closure)) {
+        // Although the body doesn't participate in inference we still
+        // want to type-check captures to make sure that the context
+        // is valid.
+        if (captureList)
+          createConjunction(cs, elements, locator);
+
         return;
+      }
     }
 
     if (isChildOf(StmtKind::Case)) {
@@ -1070,7 +1109,6 @@ class SyntacticElementConstraintGenerator
       }
     }
 
-    SmallVector<ElementInfo, 4> elements;
     for (auto element : braceStmt->getElements()) {
       if (cs.isForCodeCompletion() &&
           !cs.containsIDEInspectionTarget(element)) {

Diff for: unittests/Sema/ConstraintGenerationTests.cpp

@@ -116,7 +116,7 @@ TEST_F(SemaTest, TestImplicitConditionalCastConstraintGeneration) {
               ->isEqual(getStdlibType("Int")));
 }
 
-TEST_F(SemaTest, TestCaptureListIsConnectedToTheClosure) {
+TEST_F(SemaTest, TestCaptureListIsNotOpenedEarly) {
   ConstraintSystem cs(DC, ConstraintSystemOptions());
 
   DeclAttributes attrs;
@@ -160,19 +160,19 @@ TEST_F(SemaTest, TestCaptureListIsConnectedToTheClosure) {
   auto *processed = cs.generateConstraints(assign, DC);
   ASSERT_NE(processed, nullptr);
 
-  auto *closureType = cs.getType(closure)->castTo<TypeVariableType>();
-  auto &CG = cs.getConstraintGraph();
+  for (const auto &capture : captureList->getCaptureList()) {
+    ASSERT_FALSE(cs.hasType(capture.getVar()));
+  }
 
-  for (auto *constraint : CG[closureType].getConstraints()) {
-    if (constraint->getKind() != ConstraintKind::DefaultClosureType)
-      continue;
+  auto *closureType = cs.getType(closure)->castTo<TypeVariableType>();
 
-    for (const auto &capture : captureList->getCaptureList()) {
-      auto *capturedVar =
-          cs.getType(capture.getVar())->castTo<TypeVariableType>();
+  ASTExtInfo extInfo;
+  ASSERT_TRUE(cs.resolveClosure(
+      closureType,
+      FunctionType::get(/*params*/ {}, Context.TheEmptyTupleType, extInfo),
+      cs.getConstraintLocator(closure)));
 
-      ASSERT_TRUE(
-          llvm::is_contained(constraint->getTypeVariables(), capturedVar));
-    }
+  for (const auto &capture : captureList->getCaptureList()) {
+    ASSERT_TRUE(cs.hasType(capture.getVar()));
   }
 }