Merge pull request #84206 from rjmccall/capture-defer-isolation

Make sure that we capture the isolation of a defer function that implicitly uses the isolation

Author: rjmccall

lib/Sema/TypeCheckCaptures.cpp

@@ -66,6 +66,7 @@ class FindCapturedVars : public ASTWalker {
   DeclContext *CurDC;
   bool NoEscape, ObjC;
   bool HasGenericParamCaptures;
+  bool HasUsesOfCurrentIsolation = false;
 
 public:
   FindCapturedVars(SourceLoc CaptureLoc,
@@ -91,6 +92,10 @@ class FindCapturedVars : public ASTWalker {
                        CapturedTypes);
   }
 
+  bool hasUsesOfCurrentIsolation() const {
+    return HasUsesOfCurrentIsolation;
+  }
+
   bool hasGenericParamCaptures() const {
     return HasGenericParamCaptures;
   }
@@ -693,6 +698,31 @@ class FindCapturedVars : public ASTWalker {
       checkType(typeValue->getParamType(), E->getLoc());
     }
 
+    // Record that we saw an #isolation expression that hasn't been filled in.
+    if (auto currentIsolation = dyn_cast<CurrentContextIsolationExpr>(E)) {
+      if (!currentIsolation->getActor())
+        HasUsesOfCurrentIsolation = true;
+    }
+
+    // Record that we saw an apply of a function with caller isolation.
+    if (auto apply = dyn_cast<ApplyExpr>(E)) {
+      if (auto type = apply->getFn()->getType()) {
+        if (auto fnType = type->getAs<AnyFunctionType>();
+            fnType && fnType->getIsolation().isNonIsolatedCaller()) {
+          HasUsesOfCurrentIsolation = true;
+        }
+      }
+    }
+
+    // Look into caller-side default arguments.
+    if (auto defArg = dyn_cast<DefaultArgumentExpr>(E)) {
+      if (defArg->isCallerSide()) {
+        if (auto callerSideExpr = defArg->getCallerSideDefaultExpr()) {
+          callerSideExpr->walk(*this);
+        }
+      }
+    }
+
     return Action::Continue(E);
   }
 
@@ -747,13 +777,18 @@ class FindCapturedVars : public ASTWalker {
 /// 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) {
+                                              VarDecl *var,
+                                              bool hasUsesOfCurrentIsolation) {
   assert(isa<ParamDecl>(var));
 
   // Don't try to capture an isolated parameter of the function itself.
   if (var->getDeclContext() == AFD)
     return false;
 
+  // Force capture if we have uses of the isolation in the function body.
+  if (hasUsesOfCurrentIsolation)
+    return true;
+
   // 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
@@ -792,7 +827,8 @@ CaptureInfo CaptureInfoRequest::evaluate(Evaluator &evaluator,
     auto actorIsolation = getActorIsolation(AFD);
     if (actorIsolation.getKind() == ActorIsolation::ActorInstance) {
       if (auto *var = actorIsolation.getActorInstance()) {
-        if (shouldCaptureIsolationInLocalFunc(AFD, var))
+        if (shouldCaptureIsolationInLocalFunc(AFD, var,
+                                              finder.hasUsesOfCurrentIsolation()))
           finder.addCapture(CapturedValue(var, 0, AFD->getLoc()));
       }
     }

test/Concurrency/isolated_nonsending_isolation_macro_sil.swift renamed to test/Concurrency/isolation_macro_sil.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -parse-as-library -emit-sil %s | %FileCheck %s
+// RUN: %target-swift-frontend -parse-as-library -emit-sil %s -module-name test | %FileCheck %s
 
 // REQUIRES: concurrency
 
@@ -9,16 +9,18 @@ nonisolated(nonsending) func nonisolatedNonsending() async {
 
 func take(iso: (any Actor)?) {}
 
+func takeDefaulted(iso: isolated (any Actor)? = #isolation) {}
+
 // CHECK-LABEL: // nonisolatedNonsending()
 // CHECK-NEXT: // Isolation: caller_isolation_inheriting
-// CHECK-NEXT: sil hidden @$s39isolated_nonsending_isolation_macro_sil21nonisolatedNonsendingyyYaF : $@convention(thin) @async (@sil_isolated @sil_implicit_leading_param @guaranteed Optional<any Actor>) -> () {
+// CHECK-NEXT: sil hidden @$s4test21nonisolatedNonsendingyyYaF : $@convention(thin) @async (@sil_isolated @sil_implicit_leading_param @guaranteed Optional<any Actor>) -> () {
 // CHECK:      bb0([[ACTOR:%.*]] : $Optional<any Actor>):
 // CHECK:   retain_value [[ACTOR]]
 // CHECK:   debug_value [[ACTOR]], let, name "iso"
-// CHECK:   [[FUNC:%.*]] = function_ref @$s39isolated_nonsending_isolation_macro_sil4take3isoyScA_pSg_tF : $@convention(thin) (@guaranteed Optional<any Actor>) -> ()
+// CHECK:   [[FUNC:%.*]] = function_ref @$s4test4take3isoyScA_pSg_tF : $@convention(thin) (@guaranteed Optional<any Actor>) -> ()
 // CHECK:   apply [[FUNC]]([[ACTOR]]) : $@convention(thin) (@guaranteed Optional<any Actor>) -> ()
 // CHECK:   release_value [[ACTOR]]
-// CHECK: } // end sil function '$s39isolated_nonsending_isolation_macro_sil21nonisolatedNonsendingyyYaF'
+// CHECK: } // end sil function '$s4test21nonisolatedNonsendingyyYaF'
 
 //   Check that we emit #isolation correctly in closures.
 // CHECK-LABEL: // closure #1 in containsClosure()
@@ -33,6 +35,49 @@ func containsClosure() {
   }
 }
 
+//   Check that we capture variables as necessary to emit #isolation
+//   correctly in defer bodies.
+func deferWithIsolatedParam(_ iso: isolated (any Actor)) {
+  defer {
+    take(iso: #isolation)
+  }
+  do {}
+}
+// CHECK-LABEL: sil hidden @$s4test22deferWithIsolatedParamyyScA_pYiF :
+// CHECK:       bb0(%0 : $any Actor)
+// CHECK:         [[DEFER:%.*]] = function_ref @$s4test22deferWithIsolatedParamyyScA_pYiF6$deferL_yyF :
+// CHECK-NEXT:    apply [[DEFER]](%0)
+
+// CHECK-LABEL: sil private @$s4test22deferWithIsolatedParamyyScA_pYiF6$deferL_yyF :
+// CHECK:       bb0(%0 : @closureCapture $any Actor):
+// CHECK:         [[T0:%.*]] = enum $Optional<any Actor>, #Optional.some!enumelt, %0
+// CHECK:         [[FN:%.*]] = function_ref @$s4test4take3isoyScA_pSg_tF :
+// CHECK-NEXT:    apply [[FN]]([[T0]])
+
+//   Check that that happens even with uses in caller-side default
+//   arguments, which capture analysis was not previously walking into.
+func deferWithIsolatedParam_defaultedUse(_ iso: isolated (any Actor)) {
+  defer {
+    takeDefaulted()
+  }
+  do {}
+}
+
+// CHECK-LABEL: sil hidden @$s4test35deferWithIsolatedParam_defaultedUseyyScA_pYiF :
+// CHECK:       bb0(%0 : $any Actor):
+// CHECK:         [[DEFER:%.*]] = function_ref @$s4test35deferWithIsolatedParam_defaultedUseyyScA_pYiF6$deferL_yyF :
+// CHECK-NEXT:    apply [[DEFER]](%0)
+
+// CHECK-LABEL: sil private @$s4test35deferWithIsolatedParam_defaultedUseyyScA_pYiF6$deferL_yyF :
+// CHECK:       bb0(%0 : @closureCapture $any Actor):
+// CHECK:         [[T0:%.*]] = enum $Optional<any Actor>, #Optional.some!enumelt, %0
+// CHECK:         [[FN:%.*]] = function_ref @$s4test13takeDefaulted3isoyScA_pSgYi_tF :
+// CHECK-NEXT:    apply [[FN]]([[T0]])
+
+// TODO: we can't currently call nonisolated(nonsending) functions in
+// defer bodies because they have to be async, but that should be
+// tested here as well.
+
 //   Check that we emit #isolation correctly in defer bodies.
 nonisolated(nonsending)
 func hasDefer() async {
@@ -41,7 +86,7 @@ func hasDefer() async {
   }
   do {}
 }
-// CHECK-LABEL: sil hidden @$s39isolated_nonsending_isolation_macro_sil8hasDeferyyYaF :
+// CHECK-LABEL: sil hidden @$s4test8hasDeferyyYaF :
 // CHECK:       bb0(%0 : $Optional<any Actor>):
 // CHECK:         // function_ref $defer
 // CHECK-NEXT:    [[DEFER:%.*]] = function_ref
@@ -66,7 +111,7 @@ func hasNestedDefer() async {
   do {}
 }
 
-// CHECK-LABEL: sil hidden @$s39isolated_nonsending_isolation_macro_sil14hasNestedDeferyyYaF :
+// CHECK-LABEL: sil hidden @$s4test14hasNestedDeferyyYaF :
 // CHECK:       bb0(%0 : $Optional<any Actor>):
 // CHECK:         // function_ref $defer #1 () in hasNestedDefer()
 // CHECK-NEXT:    [[DEFER:%.*]] = function_ref