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

include/swift/AST/DeclContext.h

@@ -82,6 +82,7 @@ namespace swift {
   class SerializedTopLevelCodeDecl;
   class StructDecl;
   class AccessorDecl;
+  class ClosureExpr;
 
   template <typename T>
   struct AvailableDuringLoweringDeclFilter;
@@ -504,6 +505,14 @@ class alignas(1 << DeclContextAlignInBits) DeclContext
         const_cast<DeclContext *>(this)->getInnermostSkippedFunctionContext();
   }
 
+  /// Returns the innermost context that is a ClosureExpr, which defines how
+  /// self behaves, unless within a type context that redefines self.
+  LLVM_READONLY
+  ClosureExpr *getInnermostClosureForSelfCapture();
+  const ClosureExpr *getInnermostClosureForSelfCapture() const {
+    return const_cast<DeclContext *>(this)->getInnermostClosureForSelfCapture();
+  }
+
   /// Returns the semantic parent of this context.  A context has a
   /// parent if and only if it is not a module context.
   DeclContext *getParent() const {

include/swift/AST/Expr.h

@@ -4393,7 +4393,7 @@ struct CaptureListEntry {
   explicit CaptureListEntry(PatternBindingDecl *PBD);
 
   VarDecl *getVar() const;
-  bool isSimpleSelfCapture() const;
+  bool isSimpleSelfCapture(bool excludeWeakCaptures = true) const;
 };
 
 /// CaptureListExpr - This expression represents the capture list on an explicit

include/swift/AST/Stmt.h

@@ -709,9 +709,13 @@ class alignas(1 << PatternAlignInBits) StmtConditionElement {
   }
 
   /// Whether or not this conditional stmt rebinds self with a `let self`
-  /// or `let self = self` condition. If `requireLoadExpr` is `true`,
-  /// additionally requires that the RHS of the self condition is a `LoadExpr`.
-  bool rebindsSelf(ASTContext &Ctx, bool requireLoadExpr = false) const;
+  /// or `let self = self` condition.
+  ///  - If `requiresCaptureListRef` is `true`, additionally requires that the
+  ///    RHS of the self condition references a var defined in a capture list.
+  ///  - If `requireLoadExpr` is `true`, additionally requires that the RHS of
+  ///    the self condition is a `LoadExpr`.
+  bool rebindsSelf(ASTContext &Ctx, bool requiresCaptureListRef = false,
+                   bool requireLoadExpr = false) const;
 
   SourceLoc getStartLoc() const;
   SourceLoc getEndLoc() const;
@@ -817,9 +821,13 @@ class LabeledConditionalStmt : public LabeledStmt {
   StmtCondition *getCondPointer() { return &Cond; }
 
   /// Whether or not this conditional stmt rebinds self with a `let self`
-  /// or `let self = self` condition. If `requireLoadExpr` is `true`,
-  /// additionally requires that the RHS of the self condition is a `LoadExpr`.
-  bool rebindsSelf(ASTContext &Ctx, bool requireLoadExpr = false) const;
+  /// or `let self = self` condition.
+  ///  - If `requiresCaptureListRef` is `true`, additionally requires that the
+  ///    RHS of the self condition references a var defined in a capture list.
+  ///  - If `requireLoadExpr` is `true`, additionally requires that the RHS of
+  ///    the self condition is a `LoadExpr`.
+  bool rebindsSelf(ASTContext &Ctx, bool requiresCaptureListRef = false,
+                   bool requireLoadExpr = false) const;
 
   static bool classof(const Stmt *S) {
     return S->getKind() >= StmtKind::First_LabeledConditionalStmt &&

lib/AST/DeclContext.cpp

@@ -263,6 +263,25 @@ DeclContext *DeclContext::getInnermostSkippedFunctionContext() {
   return nullptr;
 }
 
+ClosureExpr *DeclContext::getInnermostClosureForSelfCapture() {
+  auto dc = this;
+  if (auto closure = dyn_cast<ClosureExpr>(dc)) {
+    return closure;
+  }
+
+  // Stop searching if we find a type decl, since types always
+  // redefine what 'self' means, even when nested inside a closure.
+  if (dc->isTypeContext()) {
+    return nullptr;
+  }
+
+  if (auto parent = dc->getParent()) {
+    return parent->getInnermostClosureForSelfCapture();
+  }
+
+  return nullptr;
+}
+
 DeclContext *DeclContext::getParentForLookup() const {
   if (isa<ExtensionDecl>(this)) {
     // If we are inside an extension, skip directly

lib/AST/Expr.cpp

@@ -1353,22 +1353,29 @@ VarDecl *CaptureListEntry::getVar() const {
   return PBD->getSingleVar();
 }
 
-bool CaptureListEntry::isSimpleSelfCapture() const {
+bool CaptureListEntry::isSimpleSelfCapture(bool excludeWeakCaptures) const {
   auto *Var = getVar();
   auto &ctx = Var->getASTContext();
 
   if (Var->getName() != ctx.Id_self)
     return false;
 
-  if (auto *attr = Var->getAttrs().getAttribute<ReferenceOwnershipAttr>())
-    if (attr->get() == ReferenceOwnership::Weak)
-      return false;
-
   if (PBD->getPatternList().size() != 1)
     return false;
 
   auto *expr = PBD->getInit(0);
 
+  if (auto *attr = Var->getAttrs().getAttribute<ReferenceOwnershipAttr>()) {
+    if (attr->get() == ReferenceOwnership::Weak && excludeWeakCaptures) {
+      return false;
+    }
+  }
+
+  // Look through any ImplicitConversionExpr that may contain the DRE
+  if (auto conversion = dyn_cast_or_null<ImplicitConversionExpr>(expr)) {
+    expr = conversion->getSubExpr();
+  }
+
   if (auto *DRE = dyn_cast<DeclRefExpr>(expr)) {
     if (auto *VD = dyn_cast<VarDecl>(DRE->getDecl())) {
       return VD->getName() == ctx.Id_self;

lib/AST/Stmt.cpp

@@ -526,19 +526,28 @@ void LabeledConditionalStmt::setCond(StmtCondition e) {
 }
 
 /// Whether or not this conditional stmt rebinds self with a `let self`
-/// or `let self = self` condition. If `requireLoadExpr` is `true`,
-/// additionally requires that the RHS of the self condition is a `LoadExpr`.
+/// or `let self = self` condition.
+///  - If `requiresCaptureListRef` is `true`, additionally requires that the
+///    RHS of the self condition references a var defined in a capture list.
+///  - If `requireLoadExpr` is `true`, additionally requires that the RHS of
+///    the self condition is a `LoadExpr`.
 bool LabeledConditionalStmt::rebindsSelf(ASTContext &Ctx,
+                                         bool requiresCaptureListRef,
                                          bool requireLoadExpr) const {
-  return llvm::any_of(getCond(), [&Ctx, requireLoadExpr](const auto &cond) {
-    return cond.rebindsSelf(Ctx, requireLoadExpr);
+  return llvm::any_of(getCond(), [&Ctx, requiresCaptureListRef,
+                                  requireLoadExpr](const auto &cond) {
+    return cond.rebindsSelf(Ctx, requiresCaptureListRef, requireLoadExpr);
   });
 }
 
 /// Whether or not this conditional stmt rebinds self with a `let self`
-/// or `let self = self` condition. If `requireLoadExpr` is `true`,
-/// additionally requires that the RHS of the self condition is a `LoadExpr`.
+/// or `let self = self` condition.
+///  - If `requiresCaptureListRef` is `true`, additionally requires that the
+///    RHS of the self condition references a var defined in a capture list.
+///  - If `requireLoadExpr` is `true`, additionally requires that the RHS of
+///    the self condition is a `LoadExpr`.
 bool StmtConditionElement::rebindsSelf(ASTContext &Ctx,
+                                       bool requiresCaptureListRef,
                                        bool requireLoadExpr) const {
   auto pattern = getPatternOrNull();
   if (!pattern) {
@@ -578,8 +587,22 @@ bool StmtConditionElement::rebindsSelf(ASTContext &Ctx,
   if (auto *DRE = dyn_cast<DeclRefExpr>(
           exprToCheckForDRE->getSemanticsProvidingExpr())) {
     auto *decl = DRE->getDecl();
-    return decl && decl->hasName() ? decl->getName().isSimpleName(Ctx.Id_self)
-                                   : false;
+
+    bool definedInCaptureList = false;
+    if (auto varDecl = dyn_cast_or_null<VarDecl>(DRE->getDecl())) {
+      definedInCaptureList = varDecl->isCaptureList();
+    }
+
+    if (requiresCaptureListRef && !definedInCaptureList) {
+      return false;
+    }
+
+    bool isVariableNamedSelf = false;
+    if (decl && decl->hasName()) {
+      isVariableNamedSelf = decl->getName().isSimpleName(Ctx.Id_self);
+    }
+
+    return isVariableNamedSelf;
   }
 
   return false;

lib/AST/UnqualifiedLookup.cpp

@@ -333,8 +333,7 @@ UnqualifiedLookupFactory::getBaseDeclForResult(const DeclContext *baseDC) const
 /// 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.
-static bool implicitSelfReferenceIsUnwrapped(const ValueDecl *selfDecl,
-                                             const AbstractClosureExpr *inClosure) {
+static bool implicitSelfReferenceIsUnwrapped(const ValueDecl *selfDecl) {
   ASTContext &Ctx = selfDecl->getASTContext();
 
   // Check if the implicit self decl refers to a var in a conditional stmt
@@ -352,71 +351,40 @@ static bool implicitSelfReferenceIsUnwrapped(const ValueDecl *selfDecl,
   return conditionalStmt->rebindsSelf(Ctx);
 }
 
-// Finds the nearest parent closure, which would define the
-// permitted usage of implicit self. In closures this is most
-// often just `dc` itself, but in functions defined in the
-// closure body this would be some parent context.
-static ClosureExpr *closestParentClosure(DeclContext *dc) {
-  if (!dc) {
-    return nullptr;
-  }
-
-  if (auto closure = dyn_cast<ClosureExpr>(dc)) {
-    return closure;
-  }
-
-  // Stop searching if we find a type decl, since types always
-  // redefine what 'self' means, even when nested inside a closure.
-  if (dc->getContextKind() == DeclContextKind::GenericTypeDecl) {
-    return nullptr;
-  }
-
-  return closestParentClosure(dc->getParent());
-}
-
 ValueDecl *UnqualifiedLookupFactory::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`.
-  auto closureExpr = closestParentClosure(DC);
+  // earlier in this closure or some outer closure.
+  auto closureExpr = DC->getInnermostClosureForSelfCapture();
   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(
       DC->getParentSourceFile(), DeclName(Ctx.Id_self), Loc);
 
   if (!selfDecl) {
     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
@@ -430,8 +398,29 @@ ValueDecl *UnqualifiedLookupFactory::lookupBaseDecl(const DeclContext *baseDC) c
   // 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 (!Ctx.LangOpts.isSwiftVersionAtLeast(6) &&
-      !implicitSelfReferenceIsUnwrapped(selfDecl, closureExpr)) {
+  if (capturesSelfWeakly && !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.
+  auto selfVD = dyn_cast<VarDecl>(selfDecl);
+  if (!selfVD) {
+    return nullptr;
+  }
+
+  if (!(selfVD->isCaptureList() || selfVD->getParentPatternStmt())) {
     return nullptr;
   }