Improve diagnostics for weak captures and captures under different names

Author: Jumhyn

include/swift/AST/DiagnosticsSema.def

@@ -3033,6 +3033,11 @@ NOTE(note_capture_self_explicitly,none,
       "capture 'self' explicitly to enable implicit 'self' in this closure", ())
 NOTE(note_reference_self_explicitly,none,
       "reference 'self.' explicitly", ())
+NOTE(note_other_self_capture,none,
+      "variable other than 'self' captured here under the name 'self' does not "
+      "enable implicit 'self'", ())
+NOTE(note_self_captured_weakly,none,
+      "weak capture of 'self' here does not enable implicit 'self'", ())
 ERROR(implicit_use_of_self_in_closure,none,
       "implicit use of 'self' in closure; use 'self.' to make"
       " capture semantics explicit", ())

include/swift/AST/Expr.h

@@ -3688,6 +3688,7 @@ class ClosureExpr : public AbstractClosureExpr {
   /// Is this a completely empty closure?
   bool hasEmptyBody() const;
 
+  /// VarDecl captured by this closure under the literal name \c self , if any.
   VarDecl *getCapturedSelfDecl() const { return CapturedSelfDecl; }
 
   static bool classof(const Expr *E) {
@@ -3794,7 +3795,7 @@ class CaptureListExpr final : public Expr,
   const ClosureExpr *getClosureBody() const { return closureBody; }
 
   void setClosureBody(ClosureExpr *body) { closureBody = body; }
-      
+
   /// This is a bit weird, but the capture list is lexically contained within
   /// the closure, so the ClosureExpr has the full source range.
   SWIFT_FORWARD_SOURCE_LOCS_TO(closureBody)

lib/AST/Decl.cpp

@@ -1532,7 +1532,7 @@ VarDecl *PatternBindingDecl::getSingleInitializerVar() const {
     return dyn_cast<VarDecl>(dyn_cast<DeclRefExpr>(getInit(0))->getDecl());
   return nullptr;
 }
-
+    
 bool VarDecl::isInitExposedToClients() const {
   auto parent = dyn_cast<NominalTypeDecl>(getDeclContext());
   if (!parent) return false;

lib/AST/Expr.cpp

@@ -1233,10 +1233,8 @@ bool CaptureListEntry::isSimpleSelfCapture() const {
     return false;
   if (auto *DRE = dyn_cast<DeclRefExpr>(Init->getInit(0)))
     if (auto *VD = dyn_cast<VarDecl>(DRE->getDecl())) {
-      auto *ownership = Var->getAttrs().getAttribute<ReferenceOwnershipAttr>();
       return (VD->isSelfParameter() || VD->isSelfParamCapture())
-             && VD->getName() == Var->getName()
-             && (!ownership || ownership->get() != ReferenceOwnership::Weak);
+             && VD->getName() == Var->getName();
     }
   return false;
 }

lib/AST/UnqualifiedLookup.cpp

@@ -700,8 +700,7 @@ void UnqualifiedLookupFactory::lookupNamesIntroducedByMemberFunction(
         // for lookup.
         DeclContext *const BaseDC =
             isOutsideBodyOfFunction(AFD) ? fnDeclContext
-            : capturedSelfContext == nullptr ?  AFD
-            : capturedSelfContext;
+            : capturedSelfContext ?:  AFD;
         // If we are inside of a method, check to see if there are any ivars in
         // scope, and if so, whether this is a reference to one of them.
         // FIXME: We should persist this information between lookups.
@@ -733,8 +732,12 @@ void UnqualifiedLookupFactory::lookupNamesIntroducedByClosure(
     AbstractClosureExpr *ACE, Optional<bool> isCascadingUse) {
   if (auto *CE = dyn_cast<ClosureExpr>(ACE)) {
     lookForLocalVariablesIn(CE);
-      if (capturedSelfContext == nullptr && CE->getCapturedSelfDecl())
-        capturedSelfContext = CE;
+    // If we don't already have a captured self context, and this closure
+    // captures the self param, remember that.
+    if (capturedSelfContext == nullptr)
+      if (auto *VD = CE->getCapturedSelfDecl())
+        if (VD->isSelfParamCapture() && !VD->getType()->is<WeakStorageType>())
+          capturedSelfContext = CE;
   }
   ifNotDoneYet([&] {
     // clang-format off

lib/Sema/MiscDiagnostics.cpp

@@ -1386,7 +1386,7 @@ static void diagnoseImplicitSelfUseInClosure(TypeChecker &TC, const Expr *E,
                       diag::property_use_in_closure_without_explicit_self,
                       MRE->getMember().getDecl()->getBaseName().getIdentifier());
         }
-      
+
       // Handle method calls with a specific diagnostic + fixit.
       if (auto *DSCE = dyn_cast<DotSyntaxCallExpr>(E))
         if (isImplicitSelfParamUse(DSCE->getBase()) &&
@@ -1403,10 +1403,22 @@ static void diagnoseImplicitSelfUseInClosure(TypeChecker &TC, const Expr *E,
       if (memberLoc.isValid()) {
         if (auto *CE = dyn_cast<ClosureExpr>(ACE)) {
           // If we've already captured something with the name "self" other than
-          // the actual self param, don't bother with any fix-its.
-          auto *VD = CE->getCapturedSelfDecl();
-          if (VD && !VD->isSelfParamCapture())
+          // the actual self param, offer special diagnostics.
+          if (auto *VD = CE->getCapturedSelfDecl()) {
+            // Either this is a weak capture of self...
+            if (VD->getType()->is<WeakStorageType>()) {
+              TC.diagnose(VD->getLoc(), diag::note_self_captured_weakly);
+            // ...or something completely different.
+            } else {
+              TC.diagnose(VD->getLoc(), diag::note_other_self_capture);
+            }
+            // Bail on the rest of the diagnostics. Offering the option to
+            // capture 'self' explicitly will result in an error, and using
+            // 'self.' explicitly will be accessing something other than the
+            // self param.
+            // FIXME: We could offer a special fixit in the [weak self] case to insert 'self?.'...
             return { false, E };
+          }
 
           auto diag = TC.diagnose(CE->getLoc(),
                                   diag::note_capture_self_explicitly);