Merge pull request #33116 from DougGregor/captures-in-constraint-system

[Constraint solver] Type check captures as part of the constraint system

Author: DougGregor

include/swift/AST/ASTWalker.h

@@ -229,6 +229,11 @@ class ASTWalker {
   /// TapExpr.
   virtual bool shouldWalkIntoTapExpression() { return true; }
 
+  /// This method configures whether the walker should visit the capture
+  /// initializer expressions within a capture list directly, rather than
+  /// walking the declarations.
+  virtual bool shouldWalkCaptureInitializerExpressions() { return false; }
+
   /// This method configures whether the walker should exhibit the legacy
   /// behavior where accessors appear as peers of their storage, rather
   /// than children nested inside of it.

lib/AST/ASTWalker.cpp

@@ -803,8 +803,16 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
 
   Expr *visitCaptureListExpr(CaptureListExpr *expr) {
     for (auto c : expr->getCaptureList()) {
-      if (doIt(c.Var) || doIt(c.Init))
+      if (Walker.shouldWalkCaptureInitializerExpressions()) {
+        for (auto entryIdx : range(c.Init->getNumPatternEntries())) {
+          if (auto newInit = doIt(c.Init->getInit(entryIdx)))
+            c.Init->setInit(entryIdx, newInit);
+          else
+            return nullptr;
+        }
+      } else if (doIt(c.Var) || doIt(c.Init)) {
         return nullptr;
+      }
     }
 
     ClosureExpr *body = expr->getClosureBody();

lib/AST/Expr.cpp

@@ -1968,8 +1968,15 @@ bool ClosureExpr::hasEmptyBody() const {
 }
 
 bool ClosureExpr::capturesSelfEnablingImplictSelf() const {
-  if (auto *VD = getCapturedSelfDecl())
-    return VD->isSelfParamCapture() && !VD->getType()->is<WeakStorageType>();
+  if (auto *VD = getCapturedSelfDecl()) {
+    if (!VD->isSelfParamCapture())
+      return false;
+
+    if (auto *attr = VD->getAttrs().getAttribute<ReferenceOwnershipAttr>())
+      return attr->get() != ReferenceOwnership::Weak;
+
+    return true;
+  }
   return false;
 }
 

lib/Sema/CSApply.cpp

@@ -7649,6 +7649,13 @@ namespace {
         TapsToTypeCheck.push_back(std::make_pair(tap, Rewriter.dc));
       }
 
+      if (auto captureList = dyn_cast<CaptureListExpr>(expr)) {
+        // Rewrite captures.
+        for (const auto &capture : captureList->getCaptureList()) {
+          (void)rewriteTarget(SolutionApplicationTarget(capture.Init));
+        }
+      }
+
       Rewriter.walkToExprPre(expr);
       return { true, expr };
     }

lib/Sema/CSGen.cpp

@@ -2783,6 +2783,8 @@ namespace {
 
         CollectVarRefs(ConstraintSystem &cs) : cs(cs) { }
 
+        bool shouldWalkCaptureInitializerExpressions() override { return true; }
+
         std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
           // If there are any error expressions in this closure
           // it wouldn't be possible to infer its type.
@@ -3729,6 +3731,16 @@ namespace {
         }
       }
 
+      // Generate constraints for each of the entries in the capture list.
+      if (auto captureList = dyn_cast<CaptureListExpr>(expr)) {
+        auto &CS = CG.getConstraintSystem();
+        for (const auto &capture : captureList->getCaptureList()) {
+          SolutionApplicationTarget target(capture.Init);
+          if (CS.generateConstraints(target, FreeTypeVariableBinding::Disallow))
+            return {false, nullptr};
+        }
+      }
+
       // Both multi- and single-statement closures now behave the same way
       // when it comes to constraint generation.
       if (auto closure = dyn_cast<ClosureExpr>(expr)) {

lib/Sema/MiscDiagnostics.cpp

@@ -96,6 +96,8 @@ static void diagSyntacticUseRestrictions(const Expr *E, const DeclContext *DC,
       return false;
     }
 
+    bool shouldWalkCaptureInitializerExpressions() override { return true; }
+
     bool shouldWalkIntoTapExpression() override { return false; }
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
@@ -1332,6 +1334,8 @@ static void diagRecursivePropertyAccess(const Expr *E, const DeclContext *DC) {
       return false;
     }
 
+    bool shouldWalkCaptureInitializerExpressions() override { return true; }
+
     bool shouldWalkIntoTapExpression() override { return false; }
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
@@ -1505,6 +1509,8 @@ static void diagnoseImplicitSelfUseInClosure(const Expr *E,
       return false;
     }
 
+    bool shouldWalkCaptureInitializerExpressions() override { return true; }
+
     bool shouldWalkIntoTapExpression() override { return false; }
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
@@ -1515,7 +1521,6 @@ static void diagnoseImplicitSelfUseInClosure(const Expr *E,
           Closures.push_back(CE);
       }
 
-
       // If we aren't in a closure, no diagnostics will be produced.
       if (Closures.size() == 0)
         return { true, E };
@@ -3332,6 +3337,8 @@ static void checkStmtConditionTrailingClosure(ASTContext &ctx, const Expr *E) {
       return false;
     }
 
+    bool shouldWalkCaptureInitializerExpressions() override { return true; }
+
     bool shouldWalkIntoTapExpression() override { return false; }
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
@@ -3482,6 +3489,8 @@ class ObjCSelectorWalker : public ASTWalker {
     return false;
   }
 
+  bool shouldWalkCaptureInitializerExpressions() override { return true; }
+
   bool shouldWalkIntoTapExpression() override { return false; }
 
   std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
@@ -4222,6 +4231,8 @@ static void diagnoseUnintendedOptionalBehavior(const Expr *E,
       return false;
     }
 
+    bool shouldWalkCaptureInitializerExpressions() override { return true; }
+
     bool shouldWalkIntoTapExpression() override { return false; }
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
@@ -4298,6 +4309,8 @@ static void diagnoseDeprecatedWritableKeyPath(const Expr *E,
       return false;
     }
 
+    bool shouldWalkCaptureInitializerExpressions() override { return true; }
+
     bool shouldWalkIntoTapExpression() override { return false; }
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {

lib/Sema/TypeCheckConstraints.cpp

@@ -1027,6 +1027,8 @@ namespace {
 
     bool walkToClosureExprPre(ClosureExpr *expr);
 
+    bool shouldWalkCaptureInitializerExpressions() override { return true; }
+
     std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
       // If this is a call, record the argument expression.
       if (auto call = dyn_cast<ApplyExpr>(expr)) {
@@ -1076,21 +1078,6 @@ namespace {
         return std::make_pair(recursive, expr);
       };
 
-      // For capture lists, we typecheck the decls they contain.
-      if (auto captureList = dyn_cast<CaptureListExpr>(expr)) {
-        // Validate the capture list.
-        for (auto capture : captureList->getCaptureList()) {
-          TypeChecker::typeCheckDecl(capture.Init);
-          TypeChecker::typeCheckDecl(capture.Var);
-        }
-
-        // Since closure expression is contained by capture list
-        // let's handle it directly to avoid walking into capture
-        // list itself.
-        captureList->getClosureBody()->walk(*this);
-        return finish(false, expr);
-      }
-
       // For closures, type-check the patterns and result type as written,
       // but do not walk into the body. That will be type-checked after
       // we've determine the complete function type.

test/expr/closure/closures.swift

@@ -504,3 +504,22 @@ var qux1: (() -> Void)? = {}
 
 SR9839(qux1!)
 SR9839(forceUnwrap(qux1))
+
+// rdar://problem/65155671 - crash referencing parameter of outer closure
+func rdar65155671(x: Int) {
+    { a in
+      _ = { [a] in a }
+    }(x)
+}
+
+func sr3186<T, U>(_ f: (@escaping (@escaping (T) -> U) -> ((T) -> U))) -> ((T) -> U) {
+    return { x in return f(sr3186(f))(x) }
+}
+
+class SR3186 {
+  init() {
+    // expected-warning@+1{{capture 'self' was never used}}
+    let v = sr3186 { f in { [unowned self, f] x in x != 1000 ? f(x + 1) : "success" } }(0)
+    print("\(v)")
+  }
+}