[6.2] Don't force a capture of an isolated parameter in defer bodies

lib/SILGen/SILGenConcurrency.cpp

@@ -136,7 +136,8 @@ void SILGenFunction::emitExpectedExecutorProlog() {
   // Defer bodies are always called synchronously within their enclosing
   // function, so the check is unnecessary; in addition, we cannot
   // necessarily perform the check because the defer may not have
-  // captured the isolated parameter of the enclosing function.
+  // captured the isolated parameter of the enclosing function, and
+  // forcing a capture would cause DI problems in actor initializers.
   bool wantDataRaceChecks = [&] {
     if (F.isAsync() || F.isDefer())
       return false;

lib/Sema/TypeCheckCaptures.cpp

@@ -745,6 +745,31 @@ class FindCapturedVars : public ASTWalker {
 
 } // end anonymous namespace
 
+/// Given that a local function is isolated to the given var, should we
+/// force a capture of the var?
+static bool shouldCaptureIsolationInLocalFunc(AbstractFunctionDecl *AFD,
+                                              VarDecl *var) {
+  assert(isa<ParamDecl>(var));
+
+  // Don't try to capture an isolated parameter of the function itself.
+  if (var->getDeclContext() == AFD)
+    return false;
+
+  // We only *need* to force a capture of the isolation in an async function
+  // (in which case it's needed for executor switching) or if we're in the
+  // mode that forces an executor check in all synchronous functions. But
+  // it's a simpler rule if we just do it unconditionally.
+
+  // However, don't do it for the implicit functions that represent defer
+  // bodies, where it is both unnecessary and likely to lead to bad diagnostics.
+  // We already suppress the executor check in defer bodies.
+  if (auto FD = dyn_cast<FuncDecl>(AFD))
+    if (FD->isDeferBody())
+      return false;
+
+  return true;
+}
+
 CaptureInfo CaptureInfoRequest::evaluate(Evaluator &evaluator,
                                          AbstractFunctionDecl *AFD) const {
   auto type = AFD->getInterfaceType();
@@ -768,10 +793,7 @@ CaptureInfo CaptureInfoRequest::evaluate(Evaluator &evaluator,
     auto actorIsolation = getActorIsolation(AFD);
     if (actorIsolation.getKind() == ActorIsolation::ActorInstance) {
       if (auto *var = actorIsolation.getActorInstance()) {
-        assert(isa<ParamDecl>(var));
-        // Don't capture anything if the isolation parameter is a parameter
-        // of the local function.
-        if (var->getDeclContext() != AFD)
+        if (shouldCaptureIsolationInLocalFunc(AFD, var))
           finder.addCapture(CapturedValue(var, 0, AFD->getLoc()));
       }
     }

test/SILGen/local_function_isolation.swift

@@ -71,7 +71,7 @@ actor GenericActor<K> {
 // Make sure defer doesn't capture anything.
 actor DeferInsideInitActor {
   init(foo: ()) async throws {
-    // CHECK-LABEL: sil private [ossa] @$s24local_function_isolation20DeferInsideInitActorC3fooACyt_tYaKcfc6$deferL_yyF : $@convention(thin) (@sil_isolated @guaranteed DeferInsideInitActor) -> () {
+    // CHECK-LABEL: sil private [ossa] @$s24local_function_isolation20DeferInsideInitActorC3fooACyt_tYaKcfc6$deferL_yyF : $@convention(thin) () -> () {
     defer {}
     self.init()
   }

test/SILOptimizer/definite_init_actor.swift

@@ -10,6 +10,7 @@
 
 func neverReturn() -> Never { fatalError("quit!") }
 func arbitraryAsync() async {}
+func makeIntOrThrow() throws -> Int { return 0 }
 
 actor BoringActor {
 
@@ -402,3 +403,40 @@ actor Ahmad {
   // CHECK: } // end sil function '$s4test5AhmadCACyYacfc'
   nonisolated init() async {}
 }
+
+// This should not complain about needing self in the defer prior to it being
+// fully initialized.
+actor Customer {
+  var x: Int
+  var y: Int
+
+  // CHECK-LABEL: sil hidden @$s4test8CustomerCACyYaKcfc :
+  init() async throws {
+    // CHECK:      [[GENERIC:%[0-9]+]] = enum $Optional<Builtin.Executor>, #Optional.none!enumelt
+    // CHECK-NEXT: hop_to_executor [[GENERIC]]
+    // CHECK:      [[SELF:%.*]] = end_init_let_ref %0 : $Customer
+
+    defer { print("I have a complaint") }
+
+    // CHECK:       try_apply {{.*}}, error [[FAIL1:bb[0-9]+]]
+    self.x = try makeIntOrThrow()
+
+    // CHECK:       try_apply {{.*}}, error [[FAIL2:bb[0-9]+]]
+    // CHECK:       hop_to_executor [[SELF]] : $Customer
+    self.y = try makeIntOrThrow()
+
+    // CHECK:      [[DEFER:%.*]] = function_ref @$s4test8CustomerCACyYaKcfc6$deferL_yyF :
+    // CHECK-NEXT: apply [[DEFER]]()
+    // CHECK-NEXT: return [[SELF]] : $Customer
+
+    // CHECK:    [[FAIL1]]({{%.*}} : $any Error):
+    // CHECK-NEXT: // function_ref
+    // CHECK-NEXT: [[DEFER:%.*]] = function_ref @$s4test8CustomerCACyYaKcfc6$deferL_yyF :
+    // CHECK-NEXT: apply [[DEFER]]()
+
+    // CHECK:    [[FAIL2]]({{%.*}} : $any Error):
+    // CHECK-NEXT: // function_ref
+    // CHECK-NEXT: [[DEFER:%.*]] = function_ref @$s4test8CustomerCACyYaKcfc6$deferL_yyF :
+    // CHECK-NEXT: apply [[DEFER]]()
+  }
+}