Enable SE-0365 behavior in Swift 5 mode

CHANGELOG.md

@@ -2,13 +2,13 @@
 
 _**Note:** This is in reverse chronological order, so newer entries are added to the top._
 
-## Swift 6.0
+## Swift 5.8
 
 * [SE-0365][]:
 
   Implicit `self` is now permitted for `weak self` captures, after `self` is unwrapped.
 
-  For example, the usage of implicit `self` below is now permitted:
+  For example, the usage of implicit `self` below is permitted:
 
   ```swift
   class ViewController {
@@ -44,8 +44,6 @@ _**Note:** This is in reverse chronological order, so newer entries are added to
 
   In Swift 6, the above code will no longer compile. `weak self` captures in non-escaping closures now have the same behavior as captures in escaping closures (as described in [SE-0365][]). Code relying on the previous behavior will need to be updated to either unwrap `self` (e.g. by adding a `guard let self else return` statement), or to use a different capture method (e.g. using `[self]` or `[unowned self]` instead of `[weak self]`).
 
-## Swift 5.8
-
 * [SE-0362][]:
 
   The compiler flag `-enable-upcoming-feature X` can now be used to enable a specific feature `X` that has been accepted by the evolution process, but whose introduction into the language is waiting for the next major version (e.g., version 6). The `X` is specified by any proposal that falls into this category:

lib/AST/UnqualifiedLookup.cpp

@@ -388,6 +388,43 @@ UnqualifiedLookupFactory::ResultFinderForTypeContext::getBaseDeclForResult(
   return nominalDecl;
 }
 
+/// Whether or not the given self decl is defined in an optional
+/// 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) {
+  ASTContext &Ctx = selfDecl->getASTContext();
+
+  // Check if the implicit self decl refers to a var in a conditional stmt
+  LabeledConditionalStmt *conditionalStmt = nullptr;
+  if (auto var = dyn_cast<VarDecl>(selfDecl)) {
+    if (auto parentStmt = var->getParentPatternStmt()) {
+      conditionalStmt = dyn_cast<LabeledConditionalStmt>(parentStmt);
+    }
+  }
+
+  if (!conditionalStmt) {
+    return false;
+  }
+
+  // Find the condition that defined the self decl,
+  // and check that both its LHS and RHS are 'self'
+  for (auto cond : conditionalStmt->getCond()) {
+    if (auto pattern = cond.getPattern()) {
+      if (pattern->getBoundName() != Ctx.Id_self) {
+        continue;
+      }
+    }
+
+    if (auto selfDRE = dyn_cast<DeclRefExpr>(cond.getInitializer())) {
+      return (selfDRE->getDecl()->getName().isSimpleName(Ctx.Id_self));
+    }
+  }
+
+  return false;
+}
+
 ValueDecl *UnqualifiedLookupFactory::ResultFinderForTypeContext::lookupBaseDecl(
     const DeclContext *baseDC) const {
   // Perform an unqualified lookup for the base decl of this result. This
@@ -399,29 +436,59 @@ ValueDecl *UnqualifiedLookupFactory::ResultFinderForTypeContext::lookupBaseDecl(
   // 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`.
-  bool isInWeakSelfClosure = false;
-  if (auto closureExpr = dyn_cast<ClosureExpr>(factory->DC)) {
-    if (auto decl = closureExpr->getCapturedSelfDecl()) {
-      if (auto a = decl->getAttrs().getAttribute<ReferenceOwnershipAttr>()) {
-        isInWeakSelfClosure = a->get() == ReferenceOwnership::Weak;
-      }
+  auto closureExpr = dyn_cast<ClosureExpr>(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;
     }
   }
 
-  // We can only change the behavior of lookup in Swift 6 and later,
-  // due to a bug in Swift 5 where implicit self is always allowed
-  // for weak self captures in non-escaping closures.
-  if (!factory->Ctx.LangOpts.isSwiftVersionAtLeast(6)) {
+  if (!capturesSelfWeakly) {
     return nullptr;
   }
 
-  if (isInWeakSelfClosure) {
-    return ASTScope::lookupSingleLocalDecl(factory->DC->getParentSourceFile(),
-                                           DeclName(factory->Ctx.Id_self),
-                                           factory->Loc);
+  auto selfDecl = ASTScope::lookupSingleLocalDecl(
+      factory->DC->getParentSourceFile(), DeclName(factory->Ctx.Id_self),
+      factory->Loc);
+
+  if (!selfDecl) {
+    return nullptr;
+  }
+
+  // In Swift 5 mode, implicit self is allowed within non-escaping
+  // 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,
+  // which is not acceptable for escaping closures.
+  //
+  // Escaping closures, however, only need to permit implicit self once
+  // it has been unwrapped to be non-optional:
+  //
+  //   doVoidStuffEscaping { [weak self] in
+  //     guard let self else { return }
+  //     method()
+  //   }
+  //
+  // 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)) {
+    return nullptr;
   }
 
-  return nullptr;
+  return selfDecl;
 }
 
 // TODO (someday): Instead of adding unavailable entries to Results,