[SE-0365] Fix issue where implicit self was unexpectedly disallowed in nested closures

Author: calda

lib/AST/UnqualifiedLookup.cpp

@@ -393,8 +393,7 @@ UnqualifiedLookupFactory::ResultFinderForTypeContext::getBaseDeclForResult(
 /// unwrapping condition (e.g. `guard let self else { return }`).
 /// If this is true, then we know any implicit self reference in the
 /// following scope is guaranteed to be non-optional.
-bool implicitSelfReferenceIsUnwrapped(const ValueDecl *selfDecl,
-                                      const AbstractClosureExpr *inClosure) {
+bool implicitSelfReferenceIsUnwrapped(const ValueDecl *selfDecl) {
   ASTContext &Ctx = selfDecl->getASTContext();
 
   // Check if the implicit self decl refers to a var in a conditional stmt
@@ -438,29 +437,12 @@ ValueDecl *UnqualifiedLookupFactory::ResultFinderForTypeContext::lookupBaseDecl(
     const DeclContext *baseDC) const {
   // Perform an unqualified lookup for the base decl of this result. This
   // handles cases where self was rebound (e.g. `guard let self = self`)
-  // earlier in the scope.
-  //
-  // Only do this in closures that capture self weakly, since implicit self
-  // isn't allowed to be rebound in other contexts. In other contexts, implicit
-  // self _always_ refers to the context's self `ParamDecl`, even if there
-  // is another local decl with the name `self` that would be found by
-  // `lookupSingleLocalDecl`.
+  // earlier in this closure or some outer closure.
   auto closureExpr = closestParentClosure(factory->DC);
   if (!closureExpr) {
     return nullptr;
   }
 
-  bool capturesSelfWeakly = false;
-  if (auto decl = closureExpr->getCapturedSelfDecl()) {
-    if (auto a = decl->getAttrs().getAttribute<ReferenceOwnershipAttr>()) {
-      capturesSelfWeakly = a->get() == ReferenceOwnership::Weak;
-    }
-  }
-
-  if (!capturesSelfWeakly) {
-    return nullptr;
-  }
-
   auto selfDecl = ASTScope::lookupSingleLocalDecl(
       factory->DC->getParentSourceFile(), DeclName(factory->Ctx.Id_self),
       factory->Loc);
@@ -469,16 +451,24 @@ ValueDecl *UnqualifiedLookupFactory::ResultFinderForTypeContext::lookupBaseDecl(
     return nullptr;
   }
 
+  bool capturesSelfWeakly = false;
+  if (auto decl = closureExpr->getCapturedSelfDecl()) {
+    if (auto a = decl->getAttrs().getAttribute<ReferenceOwnershipAttr>()) {
+      capturesSelfWeakly = a->get() == ReferenceOwnership::Weak;
+    }
+  }
+
   // In Swift 5 mode, implicit self is allowed within non-escaping
-  // closures even before self is unwrapped. For example, this is allowed:
+  // `weak self` closures even before self is unwrapped.
+  // For example, this is allowed:
   //
   //   doVoidStuffNonEscaping { [weak self] in
   //     method() // implicitly `self.method()`
   //   }
   //
   // To support this, we have to preserve the lookup behavior from
   // Swift 5.7 and earlier where implicit self defaults to the closure's
-  // `ParamDecl`. This causes the closure to capture self strongly, however,
+  // `ParamDecl`. This causes the closure to capture self strongly,
   // which is not acceptable for escaping closures.
   //
   // Escaping closures, however, only need to permit implicit self once
@@ -492,11 +482,29 @@ ValueDecl *UnqualifiedLookupFactory::ResultFinderForTypeContext::lookupBaseDecl(
   // In these cases, using the Swift 6 lookup behavior doesn't affect
   // how the body is type-checked, so it can be used in Swift 5 mode
   // without breaking source compatibility for non-escaping closures.
-  if (!factory->Ctx.LangOpts.isSwiftVersionAtLeast(6) &&
-      !implicitSelfReferenceIsUnwrapped(selfDecl, closureExpr)) {
+  if (capturesSelfWeakly && !factory->Ctx.LangOpts.isSwiftVersionAtLeast(6) &&
+      !implicitSelfReferenceIsUnwrapped(selfDecl)) {
     return nullptr;
   }
 
+  // Closures are only allowed to rebind self in specific circumstances:
+  //  1. In a capture list with an explicit capture.
+  //  2. In a `guard let self = self` / `if let self = self` condition
+  //     in a closure that captures self weakly.
+  //
+  // These rebindings can be done by any parent closure, and apply
+  // to all nested closures. We only need to check these structural
+  // requirements loosely here -- more extensive validation happens
+  // later in MiscDiagnostics::diagnoseImplicitSelfUseInClosure.
+  //
+  // Other types of rebindings, like an arbitrary "let `self` = foo",
+  // are never allowed to rebind self.
+  if (auto var = dyn_cast<VarDecl>(selfDecl)) {
+    if (!(var->isCaptureList() || var->getParentPatternStmt())) {
+      return nullptr;
+    }
+  }
+
   return selfDecl;
 }