Look through defer bodies in actor isolation checking

CHANGELOG.md

@@ -19,6 +19,31 @@ _**Note:** This is in reverse chronological order, so newer entries are added to
 Swift 5.6
 ---------
 
+* Actor isolation checking now understands that `defer` bodies share the isolation of their enclosing function.
+
+  ```swift
+  // Works on global actors
+  @MainActor
+  func runAnimation(controller: MyViewController) async {
+    controller.hasActiveAnimation = true
+    defer { controller.hasActiveAnimation = false }
+
+    // do the animation here...
+  }
+
+  // Works on actor instances
+  actor OperationCounter {
+    var activeOperationCount = 0
+
+    func operate() async {
+      activeOperationCount += 1
+      defer { activeOperationCount -= 1 }
+
+      // do work here...
+    }
+  }
+  ```
+
 * [SE-0335][]:
 
   Swift now allows existential types to be explicitly written with the `any`

lib/SILGen/SILGenProlog.cpp

@@ -552,11 +552,23 @@ void SILGenFunction::emitProlog(CaptureInfo captureInfo,
     return false;
   };
 
-  // Initialize ExpectedExecutor if the function is an async
-  // function or closure.
+  // Initialize ExpectedExecutor if:
+  // - this function is async or
+  // - this function is sync and isolated to an actor, and we want to
+  //   dynamically check that we're on the right executor.
+  //
+  // Actor destructors are isolated in the sense that we now have a
+  // unique reference to the actor, but we probably aren't running on
+  // the actor's executor, so we cannot safely do this check.
+  //
+  // 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.
   bool wantDataRaceChecks = getOptions().EnableActorDataRaceChecks &&
       !F.isAsync() &&
-      !isInActorDestructor(FunctionDC);
+      !isInActorDestructor(FunctionDC) &&
+      !F.isDefer();
 
   // Local function to load the expected executor from a local actor
   auto loadExpectedExecutorForLocalVar = [&](VarDecl *var) {

lib/Sema/TypeCheckConcurrency.cpp

@@ -1733,6 +1733,12 @@ namespace {
           }
         }
 
+        // Look through defers.
+        // FIXME: should this be covered automatically by the logic below?
+        if (auto func = dyn_cast<FuncDecl>(dc))
+          if (func->isDeferBody())
+            continue;
+
         if (auto func = dyn_cast<AbstractFunctionDecl>(dc)) {
           // @Sendable functions are nonisolated.
           if (func->isSendable())
@@ -1784,12 +1790,26 @@ namespace {
       return nullptr;
     }
 
+    static FuncDecl *findAnnotatableFunction(DeclContext *dc) {
+      auto fn = dyn_cast<FuncDecl>(dc);
+      if (!fn) return nullptr;
+      if (fn->isDeferBody())
+        return findAnnotatableFunction(fn->getDeclContext());
+      return fn;
+    }
+
     /// Note when the enclosing context could be put on a global actor.
     void noteGlobalActorOnContext(DeclContext *dc, Type globalActor) {
       // If we are in a synchronous function on the global actor,
       // suggest annotating with the global actor itself.
-      if (auto fn = dyn_cast<FuncDecl>(dc)) {
-        if (!isa<AccessorDecl>(fn) && !fn->isAsyncContext()) {
+      if (auto fn = findAnnotatableFunction(dc)) {
+        // Suppress this for accesssors because you can't change the
+        // actor isolation of an individual accessor.  Arguably we could
+        // add this to the entire storage declaration, though.
+        // Suppress this for async functions out of caution; but don't
+        // suppress it if we looked through a defer.
+        if (!isa<AccessorDecl>(fn) &&
+            (!fn->isAsyncContext() || fn != dc)) {
           switch (getActorIsolation(fn)) {
           case ActorIsolation::ActorInstance:
           case ActorIsolation::DistributedActorInstance:
@@ -3468,6 +3488,14 @@ ActorIsolation ActorIsolationRequest::evaluate(
   // If this is a local function, inherit the actor isolation from its
   // context if it global or was captured.
   if (auto func = dyn_cast<FuncDecl>(value)) {
+    // If this is a defer body, inherit unconditionally; we don't
+    // care if the enclosing function captures the isolated parameter.
+    if (func->isDeferBody()) {
+      auto enclosingIsolation =
+                        getActorIsolationOfContext(func->getDeclContext());
+      return inferredIsolation(enclosingIsolation);
+    }
+
     if (func->isLocalCapture() && !func->isSendable()) {
       switch (auto enclosingIsolation =
                   getActorIsolationOfContext(func->getDeclContext())) {

test/Concurrency/actor_defer.swift

@@ -0,0 +1,132 @@
+// RUN: %target-swift-frontend -parse-as-library -emit-sil -DNEGATIVES -verify %s
+// RUN: %target-swift-frontend -parse-as-library -emit-sil -enable-actor-data-race-checks -o - %s | %FileCheck %s
+
+// REQUIRES: concurrency
+
+func doSomething() {}
+
+// expected-note @+1 4 {{calls to global function 'requiresMainActor()' from outside of its actor context are implicitly asynchronous}}
+@MainActor func requiresMainActor() {}
+
+@MainActor func testNonDefer_positive() {
+  requiresMainActor()
+}
+
+#if NEGATIVES
+// expected-note @+1 {{add '@MainActor' to make global function 'testNonDefer_negative()' part of global actor 'MainActor'}}
+func testNonDefer_negative() {
+  // expected-error @+1 {{call to main actor-isolated global function 'requiresMainActor()' in a synchronous nonisolated context}}
+  requiresMainActor()
+}
+#endif
+
+@MainActor func testGlobalActor_positive() {
+  defer {
+    requiresMainActor()
+  }
+  doSomething()
+}
+// Don't include a data race check at the start of the defer
+// CHECK-LABEL: sil private @$s11actor_defer24testGlobalActor_positiveyyF6$deferL_yyF
+// CHECK-NEXT:  bb0:
+// CHECK-NEXT:    // function_ref
+// CHECK-NEXT:    function_ref
+// CHECK-NEXT:    apply
+
+#if NEGATIVES
+// expected-note @+1 {{add '@MainActor' to make global function 'testGlobalActor_negative()' part of global actor 'MainActor'}}
+func testGlobalActor_negative() {
+  defer {
+    // expected-error @+1 {{call to main actor-isolated global function 'requiresMainActor()' in a synchronous nonisolated context}}
+    requiresMainActor()
+  }
+  doSomething()
+}
+#endif
+
+@available(SwiftStdlib 5.1, *)
+@MainActor func testGlobalActorAsync_positive() async {
+  defer {
+    requiresMainActor()
+  }
+  doSomething()
+}
+
+#if NEGATIVES
+// expected-note @+2 {{add '@MainActor' to make global function 'testGlobalActorAsync_negative()' part of global actor 'MainActor'}}
+@available(SwiftStdlib 5.1, *)
+func testGlobalActorAsync_negative() async {
+  defer {
+    // expected-error @+1 {{call to main actor-isolated global function 'requiresMainActor()' in a synchronous nonisolated context}}
+    requiresMainActor()
+  }
+  doSomething()
+}
+#endif
+
+@available(SwiftStdlib 5.1, *)
+actor Actor {
+  // expected-note @+1 3 {{mutation of this property is only permitted within the actor}}
+  var actorProperty = 0
+
+  func testActor_positive() {
+    defer {
+      actorProperty += 1
+    }
+    doSomething()
+  }
+
+#if NEGATIVES
+  nonisolated func testActor_negative() {
+    defer {
+      // expected-error @+1 {{actor-isolated property 'actorProperty' can not be mutated from a non-isolated context}}
+      actorProperty += 1
+    }
+    doSomething()
+  }
+  @MainActor func testActor_negative_globalActor() {
+    defer {
+      // expected-error @+1 {{actor-isolated property 'actorProperty' can not be mutated from the main actor}}
+      actorProperty += 1
+    }
+    doSomething()
+  }
+#endif
+
+  @MainActor func testGlobalActor_positive() {
+    defer {
+      requiresMainActor()
+    }
+    doSomething()
+  }
+
+#if NEGATIVES
+  func testGlobalActor_negative() {
+    defer {
+      // expected-error @+1 {{call to main actor-isolated global function 'requiresMainActor()' in a synchronous actor-isolated context}}
+      requiresMainActor()
+    }
+    doSomething()
+  }
+#endif
+}
+
+@available(SwiftStdlib 5.1, *)
+func testIsolatedActor_positive(actor: isolated Actor) {
+  actor.actorProperty += 1
+  defer {
+    actor.actorProperty += 1
+  }
+  doSomething()
+}
+
+#if NEGATIVES
+@available(SwiftStdlib 5.1, *)
+func testIsolatedActor_negative(actor: Actor) {
+  defer {
+    // expected-error @+1 {{actor-isolated property 'actorProperty' can not be mutated from a non-isolated context}}
+    actor.actorProperty += 1
+  }
+  doSomething()
+}
+#endif