Prevent nested closure in 'extension Optional' from unexpectedly being allowed

Author: calda

lib/Sema/MiscDiagnostics.cpp

@@ -1681,12 +1681,21 @@ static void diagRecursivePropertyAccess(const Expr *E, const DeclContext *DC) {
 }
 
 /// The `weak self` capture of this closure if present
-static VarDecl *weakSelfCapture(const AbstractClosureExpr *ACE) {
+static VarDecl *selfCapture(const AbstractClosureExpr *ACE) {
   if (auto closureExpr = dyn_cast<ClosureExpr>(ACE)) {
     if (auto selfDecl = closureExpr->getCapturedSelfDecl()) {
-      if (selfDecl->getInterfaceType()->is<WeakStorageType>()) {
-        return selfDecl;
-      }
+      return selfDecl;
+    }
+  }
+
+  return nullptr;
+}
+
+/// The `weak self` capture of this closure if present
+static VarDecl *weakSelfCapture(const AbstractClosureExpr *ACE) {
+  if (auto selfDecl = selfCapture(ACE)) {
+    if (selfDecl->getInterfaceType()->is<WeakStorageType>()) {
+      return selfDecl;
     }
   }
 
@@ -2034,6 +2043,25 @@ class ImplicitSelfUsageChecker : public BaseDiagnosticWalker {
         return nullptr;
       }
 
+      // If this closure has a non-weak capture, the parent closure
+      // has a weak capture, the implicit self decl refers directly
+      // to this closure's self capture, and that capture directly
+      // refers to the parent closure's weak self capture, then
+      // this is invalid.
+      //  - Normally this is rejected by the type checker because
+      //    self is optional here, but that doesn't happen when calling
+      //    a method defined on `Optional<Self>`.
+      if (!closureHasWeakSelfCapture(inClosure)) {
+        if (auto closureSelfCapture = selfCapture(inClosure)) {
+          if (auto parentWeakSelfCapture = weakSelfCapture(outerClosure)) {
+            if (selfDecl == closureSelfCapture &&
+                parentWeakSelfCapture == outerSelfDecl) {
+              return outerClosure;
+            }
+          }
+        }
+      }
+
       // Check if this closure contains the self decl.
       //  - If the self decl is defined in the closure's body, its
       //    decl context will be the closure itself.

test/expr/closure/closures_swift6.swift

@@ -959,14 +959,62 @@ extension TestExtensionOnOptionalSelf? {
 
     _ = { [weak self] in
       _ = { [self] in
-        foo()
+        foo() // expected-error {{call to method 'foo' in closure requires explicit use of 'self' to make capture semantics explicit}}
         self.foo()
         self?.bar()
       }
     }
   }
 }
 
+// non-optional self in this extension, but on a type with members defined on optional self
+extension TestExtensionOnOptionalSelf {
+  func foo() {
+    _ = { [weak self] in
+      foo() // expected-error {{implicit use of 'self' in closure; use 'self.' to make capture semantics explicit}}
+      self.foo()
+      self?.bar()
+    }
+
+    _ = { // expected-note {{capture 'self' explicitly to enable implicit 'self' in this closure}}
+      foo() // expected-error {{call to method 'foo' in closure requires explicit use of 'self' to make capture semantics explicit}} expected-note {{reference 'self.' explicitly}}
+      self.foo()
+    }
+
+    _ = { [weak self] in
+      _ = {
+        foo()
+        self.foo()
+      }
+    }
+
+    _ = { [weak self] in
+      _ = { [self] in
+        foo() // expected-error {{call to method 'foo' in closure requires explicit use of 'self' to make capture semantics explicit}}
+        self.foo()
+      }
+    }
+
+    _ = { [weak self] in
+      _ = { [self] in
+        _ = { [self] in
+          foo() // expected-error {{call to method 'foo' in closure requires explicit use of 'self' to make capture semantics explicit}}
+          self.foo()
+        }
+      }
+    }
+
+    _ = { [weak self] in
+      doVoidStuffNonEscaping {
+        _ = { [self] in
+          foo() // expected-error {{call to method 'foo' in closure requires explicit use of 'self' to make capture semantics explicit}}
+          self.foo()
+        }
+      }
+    }
+  }
+}
+
 actor TestActor {
     func setUp() {
         doVoidStuff { [weak self] in