Merge pull request #63196 from xedin/connect-captures-to-closure

[CSGen] Captures should always be connected to their closure

Author: xedin

lib/Sema/CSGen.cpp

@@ -2970,6 +2970,15 @@ 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));

unittests/Sema/ConstraintGenerationTests.cpp

@@ -115,3 +115,64 @@ TEST_F(SemaTest, TestImplicitConditionalCastConstraintGeneration) {
   ASSERT_TRUE(getTypeOfCoercedExpr(cast<ConditionalCheckedCastExpr>(resultExpr))
               ->isEqual(getStdlibType("Int")));
 }
+
+TEST_F(SemaTest, TestCaptureListIsConnectedToTheClosure) {
+  ConstraintSystem cs(DC, ConstraintSystemOptions());
+
+  DeclAttributes attrs;
+  auto *closure = new (Context) ClosureExpr(attrs,
+                                            /*bracetRange=*/SourceRange(),
+                                            /*capturedSelfDecl=*/nullptr,
+                                            ParameterList::createEmpty(Context),
+                                            /*asyncLoc=*/SourceLoc(),
+                                            /*throwsLoc=*/SourceLoc(),
+                                            /*arrowLoc=*/SourceLoc(),
+                                            /*inLoc=*/SourceLoc(),
+                                            /*explicitResultType=*/nullptr, DC);
+  closure->setImplicit();
+  closure->setBody(BraceStmt::createImplicit(Context, /*elements=*/{}),
+                   /*isSingleExpression=*/true);
+
+  SmallVector<CaptureListEntry> captures;
+  {
+    // The capture variable.
+    auto *xVar = new (Context)
+        VarDecl(/*isStatic=*/false, VarDecl::Introducer::Var,
+                /*loc=*/SourceLoc(), Context.getIdentifier("x"), DC);
+    xVar->setImplicit();
+
+    auto *PBD = PatternBindingDecl::createImplicit(
+        Context, StaticSpellingKind::None,
+        /*pattern=*/NamedPattern::createImplicit(Context, xVar),
+        IntegerLiteralExpr::createFromUnsigned(Context, 42, SourceLoc()), DC);
+
+    captures.push_back(CaptureListEntry(PBD));
+  }
+
+  auto *captureList = CaptureListExpr::create(Context, captures, closure);
+
+  // _ = { [x = 42] in }
+  Expr *assign = new (Context)
+      AssignExpr(new (Context) DiscardAssignmentExpr(/*loc=*/SourceLoc(),
+                                                     /*Implicit=*/true),
+                 /*EqualLoc=*/SourceLoc(), captureList, /*Implicit=*/true);
+
+  auto *processed = cs.generateConstraints(assign, DC);
+  ASSERT_NE(processed, nullptr);
+
+  auto *closureType = cs.getType(closure)->castTo<TypeVariableType>();
+  auto &CG = cs.getConstraintGraph();
+
+  for (auto *constraint : CG[closureType].getConstraints()) {
+    if (constraint->getKind() != ConstraintKind::DefaultClosureType)
+      continue;
+
+    for (const auto &capture : captureList->getCaptureList()) {
+      auto *capturedVar =
+          cs.getType(capture.getVar())->castTo<TypeVariableType>();
+
+      ASSERT_TRUE(
+          llvm::is_contained(constraint->getTypeVariables(), capturedVar));
+    }
+  }
+}