Merge pull request #64463 from xedin/rework-capture-list-csgen

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

Author: xedin

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

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)) {

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()));
   }
 }