Consolidate LabeledConditionalStmt::rebindsSelf logic into shared helper

Author: calda

include/swift/AST/Stmt.h

@@ -583,6 +583,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;
@@ -686,6 +691,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;

lib/AST/Stmt.cpp

@@ -440,6 +440,71 @@ 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;
+  }
+
+  // The RHS could be a `LoadExpr` or `DeclRefExpr`
+  Expr *loadExprSubExpr = nullptr;
+  if (auto LE = dyn_cast<LoadExpr>(exprToCheckForDRE)) {
+    loadExprSubExpr = LE->getSubExpr();
+  }
+
+  if (requireLoadExpr) {
+    exprToCheckForDRE = loadExprSubExpr;
+  } else {
+    if (loadExprSubExpr) {
+      exprToCheckForDRE = loadExprSubExpr;
+    }
+    exprToCheckForDRE = exprToCheckForDRE->getSemanticsProvidingExpr();
+  }
+
+  DeclRefExpr *condDRE = dyn_cast_or_null<DeclRefExpr>(exprToCheckForDRE);
+  if (!condDRE || !condDRE->getDecl()->hasName()) {
+    return false;
+  }
+
+  return condDRE->getDecl()->getName().isSimpleName(Ctx.Id_self);
+}
+
 PoundAvailableInfo *
 PoundAvailableInfo::create(ASTContext &ctx, SourceLoc PoundLoc,
                            SourceLoc LParenLoc,

lib/AST/UnqualifiedLookup.cpp

@@ -397,25 +397,7 @@ 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 (cond.getKind() != StmtConditionElement::CK_PatternBinding) {
-      continue;
-    }
-    
-    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;
+  return conditionalStmt->rebindsSelf(Ctx);
 }
 
 ValueDecl *UnqualifiedLookupFactory::ResultFinderForTypeContext::lookupBaseDecl(

lib/Sema/MiscDiagnostics.cpp

@@ -1700,27 +1700,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 (cond.getKind() != StmtConditionElement::CK_PatternBinding) {
-          continue;
-        }
-        
-        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