Merge pull request #65355 from calda/cal--cp-SE-0356-fixes

xedin · web-flow · commit 3e97b1f7416e · 2023-05-16T07:09:23.000-07:00
[5.9] SE-0365 bug fixes
diff --git a/include/swift/AST/Stmt.h b/include/swift/AST/Stmt.h
@@ -556,11 +556,17 @@ class alignas(1 << PatternAlignInBits) StmtConditionElement {
     assert(getKind() == CK_PatternBinding && "Not a pattern binding condition");
     ThePattern = P;
   }
+  
+  /// Pattern Binding Accessors.
+  Expr *getInitializerOrNull() const {
+    return getKind() == CK_PatternBinding ? Condition.get<Expr *>() : nullptr;
+  }
 
   Expr *getInitializer() const {
     assert(getKind() == CK_PatternBinding && "Not a pattern binding condition");
     return Condition.get<Expr *>();
   }
+  
   void setInitializer(Expr *E) {
     assert(getKind() == CK_PatternBinding && "Not a pattern binding condition");
     Condition = E;
@@ -588,6 +594,11 @@ class alignas(1 << PatternAlignInBits) StmtConditionElement {
     Condition = Info;
   }
 
+  /// 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;
+
   SourceLoc getStartLoc() const;
   SourceLoc getEndLoc() const;
   SourceRange getSourceRange() const;
@@ -691,6 +702,11 @@ class LabeledConditionalStmt : public LabeledStmt {
   ///        stored in \c ASTNode.
   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;
+
   static bool classof(const Stmt *S) {
     return S->getKind() >= StmtKind::First_LabeledConditionalStmt &&
            S->getKind() <= StmtKind::Last_LabeledConditionalStmt;
diff --git a/lib/AST/Stmt.cpp b/lib/AST/Stmt.cpp
@@ -466,6 +466,66 @@ void LabeledConditionalStmt::setCond(StmtCondition e) {
   Cond = 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`.
+bool LabeledConditionalStmt::rebindsSelf(ASTContext &Ctx,
+                                         bool requireLoadExpr) const {
+  return llvm::any_of(getCond(), [&Ctx, requireLoadExpr](const auto &cond) {
+    return cond.rebindsSelf(Ctx, 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`.
+bool StmtConditionElement::rebindsSelf(ASTContext &Ctx,
+                                       bool requireLoadExpr) const {
+  auto pattern = getPatternOrNull();
+  if (!pattern) {
+    return false;
+  }
+
+  // Check whether or not this pattern defines a new `self` decl
+  bool isSelfRebinding = false;
+  if (pattern->getBoundName() == Ctx.Id_self) {
+    isSelfRebinding = true;
+  }
+
+  else if (auto OSP = dyn_cast<OptionalSomePattern>(pattern)) {
+    if (auto subPattern = OSP->getSubPattern()) {
+      isSelfRebinding = subPattern->getBoundName() == Ctx.Id_self;
+    }
+  }
+
+  if (!isSelfRebinding) {
+    return false;
+  }
+
+  // Check that the RHS expr is exactly `self` and not something else
+  Expr *exprToCheckForDRE = getInitializerOrNull();
+  if (!exprToCheckForDRE) {
+    return false;
+  }
+
+  if (requireLoadExpr && !isa<LoadExpr>(exprToCheckForDRE)) {
+    return false;
+  }
+
+  if (auto *load = dyn_cast<LoadExpr>(exprToCheckForDRE)) {
+    exprToCheckForDRE = load->getSubExpr();
+  }
+
+  if (auto *DRE = dyn_cast<DeclRefExpr>(
+          exprToCheckForDRE->getSemanticsProvidingExpr())) {
+    auto *decl = DRE->getDecl();
+    return decl && decl->hasName() ? decl->getName().isSimpleName(Ctx.Id_self)
+                                   : false;
+  }
+
+  return false;
+}
+
 PoundAvailableInfo *
 PoundAvailableInfo::create(ASTContext &ctx, SourceLoc PoundLoc,
                            SourceLoc LParenLoc,
diff --git a/lib/AST/UnqualifiedLookup.cpp b/lib/AST/UnqualifiedLookup.cpp
@@ -397,21 +397,29 @@ bool implicitSelfReferenceIsUnwrapped(const ValueDecl *selfDecl,
     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;
-      }
-    }
+  return conditionalStmt->rebindsSelf(Ctx);
+}
 
-    if (auto selfDRE = dyn_cast<DeclRefExpr>(cond.getInitializer())) {
-      return (selfDRE->getDecl()->getName().isSimpleName(Ctx.Id_self));
-    }
+// 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) {
+  if (!dc) {
+    return nullptr;
   }
 
-  return false;
+  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(
@@ -425,7 +433,7 @@ 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`.
-  auto closureExpr = dyn_cast<ClosureExpr>(factory->DC);
+  auto closureExpr = closestParentClosure(factory->DC);
   if (!closureExpr) {
     return nullptr;
   }
diff --git a/lib/Sema/MiscDiagnostics.cpp b/lib/Sema/MiscDiagnostics.cpp
@@ -1689,23 +1689,14 @@ static void diagnoseImplicitSelfUseInClosure(const Expr *E,
         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 OSP = dyn_cast<OptionalSomePattern>(cond.getPattern())) {
-          if (OSP->getSubPattern()->getBoundName() != Ctx.Id_self) {
-            continue;
-          }
-
-          if (auto LE = dyn_cast<LoadExpr>(cond.getInitializer())) {
-            if (auto selfDRE = dyn_cast<DeclRefExpr>(LE->getSubExpr())) {
-              return (selfDRE->getDecl()->getName().isSimpleName(Ctx.Id_self));
-            }
-          }
-        }
-      }
-
-      return false;
+      // Require `LoadExpr`s when validating the self binding.
+      // This lets us reject invalid examples like:
+      //
+      //   let `self` = self ?? .somethingElse
+      //   guard let self = self else { return }
+      //   method() // <- implicit self is not allowed
+      //
+      return conditionalStmt->rebindsSelf(Ctx, /*requireLoadExpr*/ true);
     }
 
     /// Return true if this is a closure expression that will require explicit
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
