Sema: Compute and store captures for default argument expressions

Author: slavapestov

include/swift/AST/Decl.h

@@ -5164,6 +5164,7 @@ class ParamDecl : public VarDecl {
     PointerUnion<Expr *, VarDecl *> DefaultArg;
     Initializer *InitContext = nullptr;
     StringRef StringRepresentation;
+    CaptureInfo Captures;
   };
 
   enum class Flags : uint8_t {
@@ -5249,6 +5250,13 @@ class ParamDecl : public VarDecl {
 
   void setDefaultArgumentInitContext(Initializer *initContext);
 
+  const CaptureInfo &getDefaultArgumentCaptureInfo() const {
+    assert(DefaultValueAndFlags.getPointer());
+    return DefaultValueAndFlags.getPointer()->Captures;
+  }
+
+  void setDefaultArgumentCaptureInfo(const CaptureInfo &captures);
+
   /// Extracts the text of the default argument attached to the provided
   /// ParamDecl, removing all inactive #if clauses and providing only the
   /// text of active #if clauses.

lib/AST/Decl.cpp

@@ -5866,6 +5866,11 @@ void ParamDecl::setDefaultArgumentInitContext(Initializer *initContext) {
   DefaultValueAndFlags.getPointer()->InitContext = initContext;
 }
 
+void ParamDecl::setDefaultArgumentCaptureInfo(const CaptureInfo &captures) {
+  assert(DefaultValueAndFlags.getPointer());
+  DefaultValueAndFlags.getPointer()->Captures = captures;
+}
+
 /// Return nullptr if there is no property wrapper
 Expr *swift::findOriginalPropertyWrapperInitialValue(VarDecl *var,
                                                      Expr *init) {

lib/Sema/TypeCheckCaptures.cpp

@@ -325,11 +325,8 @@ class FindCapturedVars : public ASTWalker {
     return { false, DRE };
   }
 
-  void propagateCaptures(AnyFunctionRef innerClosure, SourceLoc captureLoc) {
-    TypeChecker::computeCaptures(innerClosure);
-
-    auto &captureInfo = innerClosure.getCaptureInfo();
-
+  void propagateCaptures(const CaptureInfo &captureInfo,
+                         SourceLoc loc) {
     for (auto capture : captureInfo.getCaptures()) {
       // If the decl was captured from us, it isn't captured *by* us.
       if (capture.getDecl()->getDeclContext() == CurDC)
@@ -347,18 +344,19 @@ class FindCapturedVars : public ASTWalker {
       if (!NoEscape)
         Flags &= ~CapturedValue::IsNoEscape;
 
-      addCapture(CapturedValue(capture.getDecl(), Flags, captureLoc));
+      addCapture(CapturedValue(capture.getDecl(), Flags, capture.getLoc()));
     }
 
     if (GenericParamCaptureLoc.isInvalid())
       if (captureInfo.hasGenericParamCaptures())
-        GenericParamCaptureLoc = innerClosure.getLoc();
+        GenericParamCaptureLoc = loc;
 
-    if (DynamicSelfCaptureLoc.isInvalid())
+    if (DynamicSelfCaptureLoc.isInvalid()) {
       if (captureInfo.hasDynamicSelfCapture()) {
-        DynamicSelfCaptureLoc = innerClosure.getLoc();
+        DynamicSelfCaptureLoc = loc;
         DynamicSelf = captureInfo.getDynamicSelfType();
       }
+    }
 
     if (!OpaqueValue) {
       if (captureInfo.hasOpaqueValueCapture())
@@ -368,13 +366,13 @@ class FindCapturedVars : public ASTWalker {
 
   bool walkToDeclPre(Decl *D) override {
     if (auto *AFD = dyn_cast<AbstractFunctionDecl>(D)) {
-      propagateCaptures(AFD, AFD->getLoc());
+      TypeChecker::computeCaptures(AFD);
+      propagateCaptures(AFD->getCaptureInfo(), AFD->getLoc());
 
-      // Can default parameter initializers capture state?  That seems like
-      // a really bad idea.
-      for (auto *param : *AFD->getParameters())
-        if (auto E = param->getDefaultValue())
-          E->walk(*this);
+      for (auto *P : *AFD->getParameters())
+        if (P->getDefaultValue())
+          propagateCaptures(P->getDefaultArgumentCaptureInfo(),
+                            P->getLoc());
 
       return false;
     }
@@ -569,7 +567,8 @@ class FindCapturedVars : public ASTWalker {
     // list computed; we just propagate it, filtering out stuff that they
     // capture from us.
     if (auto *SubCE = dyn_cast<AbstractClosureExpr>(E)) {
-      propagateCaptures(SubCE, SubCE->getStartLoc());
+      TypeChecker::computeCaptures(SubCE);
+      propagateCaptures(SubCE->getCaptureInfo(), SubCE->getLoc());
       return { false, E };
     }
 
@@ -609,17 +608,36 @@ void TypeChecker::computeCaptures(AnyFunctionRef AFR) {
     finder.checkType(AFR.getType(), AFR.getLoc());
   }
 
-  auto captures = finder.getCaptureInfo();
-
   // A generic function always captures outer generic parameters.
+  bool isGeneric = false;
   auto *AFD = AFR.getAbstractFunctionDecl();
-  if (AFD) {
-    if (AFD->getGenericParams())
-      captures.setGenericParamCaptures(true);
-  }
+  if (AFD)
+    isGeneric = AFD->isGeneric();
 
+  auto captures = finder.getCaptureInfo();
+  if (isGeneric)
+    captures.setGenericParamCaptures(true);
   AFR.setCaptureInfo(captures);
 
+  // Compute captures for default argument expressions.
+  if (auto *AFD = AFR.getAbstractFunctionDecl()) {
+    for (auto *P : *AFD->getParameters()) {
+      if (auto E = P->getDefaultValue()) {
+        FindCapturedVars finder(Context,
+                                E->getLoc(),
+                                AFD,
+                                /*isNoEscape=*/false,
+                                /*isObjC=*/false);
+        E->walk(finder);
+
+        auto captures = finder.getCaptureInfo();
+        if (isGeneric)
+          captures.setGenericParamCaptures(true);
+        P->setDefaultArgumentCaptureInfo(captures);
+      }
+    }
+  }
+
   // Extensions of generic ObjC functions can't use generic parameters from
   // their context.
   if (AFD && finder.getGenericParamCaptureLoc().isValid()) {

test/SILGen/capture_order.swift

@@ -192,10 +192,10 @@ class rdar40600800 {
     }
 
     func innerFunction() {
-      let closure = { // expected-note {{captured here}}
+      let closure = {
       // FIXME: Bogus warning!
       // expected-warning@-2 {{initialization of immutable value 'closure' was never used; consider replacing with assignment to '_' or removing it}}
-        callback()
+        callback() // expected-note {{captured here}}
       }
     }
   }