Fix edge cases related to nested autoclosures, invalid weak self unwrapping

calda · calda · commit 46731e08386a · 2024-03-10T08:56:18.000-07:00
diff --git a/include/swift/AST/DeclContext.h b/include/swift/AST/DeclContext.h
@@ -504,6 +504,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
+  DeclContext *getInnermostClosureForSelfCapture();
+  const DeclContext *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 {
diff --git a/include/swift/AST/Expr.h b/include/swift/AST/Expr.h
@@ -4356,7 +4356,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
diff --git a/include/swift/AST/Stmt.h b/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 &&
diff --git a/lib/AST/DeclContext.cpp b/lib/AST/DeclContext.cpp
@@ -263,6 +263,25 @@ DeclContext *DeclContext::getInnermostSkippedFunctionContext() {
   return nullptr;
 }
 
+DeclContext *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
diff --git a/lib/AST/Expr.cpp b/lib/AST/Expr.cpp
@@ -1350,22 +1350,31 @@ 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) {
+      if (excludeWeakCaptures) {
+        return false;
+      }
+
+      // If not excluding weak captures, look through any InjectIntoOptionalExpr
+      if (auto injectExpr = dyn_cast_or_null<InjectIntoOptionalExpr>(expr)) {
+        expr = injectExpr->getSubExpr();
+      }
+    }
+  }
+
   if (auto *DRE = dyn_cast<DeclRefExpr>(expr)) {
     if (auto *VD = dyn_cast<VarDecl>(DRE->getDecl())) {
       return VD->getName() == ctx.Id_self;
diff --git a/lib/AST/Stmt.cpp b/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;
diff --git a/lib/AST/UnqualifiedLookup.cpp b/lib/AST/UnqualifiedLookup.cpp
@@ -411,34 +411,18 @@ 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.
-ClosureExpr *closestParentClosure(DeclContext *dc) {
+ValueDecl *UnqualifiedLookupFactory::ResultFinderForTypeContext::lookupBaseDecl(
+    const DeclContext *baseDC) const {
+  auto dc = factory->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::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 this closure or some outer closure.
-  auto closureExpr = closestParentClosure(factory->DC);
+  auto closureExpr =
+      dyn_cast_or_null<ClosureExpr>(dc->getInnermostClosureForSelfCapture());
   if (!closureExpr) {
     return nullptr;
   }
diff --git a/lib/Sema/MiscDiagnostics.cpp b/lib/Sema/MiscDiagnostics.cpp
diff --git a/test/expr/closure/closures.swift b/test/expr/closure/closures.swift
diff --git a/test/expr/closure/closures_swift6.swift b/test/expr/closure/closures_swift6.swift
