[Concurrency] NonisolatedNonsendingByDefault: Extend nonisolated(nonsending) to withoutActuallyEscaping

`withoutActuallyEscaping` is type-checked in a special way which
means that we need to explicitly inject `nonisolated(nonsending)`
isolation when forming a reference to this builtin.

Author: xedin

lib/Sema/TypeCheckConcurrency.cpp

@@ -3527,6 +3527,28 @@ namespace {
       if (auto isolationExpr = dyn_cast<CurrentContextIsolationExpr>(expr))
         recordCurrentContextIsolation(isolationExpr);
 
+      // `withoutActuallyEscaping` parameter types are set to be
+      // `nonisolated(nonsending)` when the `NonisolatedNonsendingByDefault`
+      // feature is enabled, which means that we need to make the argument
+      // as `nonisolated(nonsending)` if it's a closure. This cannot be done
+      // sooner because we need to make sure that closure is definitely
+      // nonisolated and due to how AST is structured we cannot do this in
+      // `determineClosureIsolation`.
+      if (ctx.LangOpts.hasFeature(Feature::NonisolatedNonsendingByDefault)) {
+        if (auto *MTEE = dyn_cast<MakeTemporarilyEscapableExpr>(expr)) {
+          if (auto *call = dyn_cast<CallExpr>(MTEE->getSubExpr())) {
+            if (auto *closure = dyn_cast<ClosureExpr>(call->getFn())) {
+              if (auto closureTy = closure->getType()->getAs<FunctionType>()) {
+                if (closureTy->isAsync() &&
+                    closure->getActorIsolation().isNonisolated())
+                  closure->setActorIsolation(
+                      ActorIsolation::forCallerIsolationInheriting());
+              }
+            }
+          }
+        }
+      }
+
       return Action::Continue(expr);
     }
 

lib/Sema/TypeOfReference.cpp

@@ -2427,9 +2427,17 @@ static DeclReferenceType getTypeOfReferenceWithSpecialTypeCheckingSemantics(
         CS.getConstraintLocator(locator, ConstraintLocator::ThrownErrorType),
         0, preparedOverload);
     FunctionType::Param arg(escapeClosure);
+
+    auto bodyParamIsolation = FunctionTypeIsolation::forNonIsolated();
+    if (CS.getASTContext().LangOpts.hasFeature(
+            Feature::NonisolatedNonsendingByDefault)) {
+      bodyParamIsolation = FunctionTypeIsolation::forNonIsolatedCaller();
+    }
+
     auto bodyClosure = FunctionType::get(arg, result,
                                          FunctionType::ExtInfoBuilder()
                                              .withNoEscape(true)
+                                             .withIsolation(bodyParamIsolation)
                                              .withAsync(true)
                                              .withThrows(true, thrownError)
                                              .build());
@@ -2438,9 +2446,16 @@ static DeclReferenceType getTypeOfReferenceWithSpecialTypeCheckingSemantics(
       FunctionType::Param(bodyClosure, CS.getASTContext().getIdentifier("do")),
     };
 
+    auto withoutEscapingIsolation = FunctionTypeIsolation::forNonIsolated();
+    if (CS.getASTContext().LangOpts.hasFeature(
+            Feature::NonisolatedNonsendingByDefault)) {
+      withoutEscapingIsolation = FunctionTypeIsolation::forNonIsolatedCaller();
+    }
+
     auto refType = FunctionType::get(args, result,
                                      FunctionType::ExtInfoBuilder()
                                          .withNoEscape(false)
+                                         .withIsolation(withoutEscapingIsolation)
                                          .withAsync(true)
                                          .withThrows(true, thrownError)
                                          .build());

test/Concurrency/attr_execution/attr_execution.swift

@@ -80,3 +80,30 @@ func testClosure() {
   takesClosure {
   }
 }
+
+func testWithoutActuallyEscaping(_ f: () async -> ()) async {
+  // CHECK-LABEL: // closure #1 in testWithoutActuallyEscaping(_:)
+  // CHECK-NEXT: // Isolation: caller_isolation_inheriting
+  await withoutActuallyEscaping(f) {
+    await $0()
+  }
+
+  // CHECK-LABEL: // closure #2 in testWithoutActuallyEscaping(_:)
+  // CHECK-NEXT: // Isolation: global_actor. type: MainActor
+  await withoutActuallyEscaping(f) { @MainActor in
+    await $0()
+  }
+
+  actor Test {
+    // CHECK-LABEL: // closure #1 in testActorIsolatedCapture() in Test #1 in testWithoutActuallyEscaping(_:)
+    // CHECK-NEXT: // Isolation: actor_instance. name: 'self'
+    func testActorIsolatedCapture() async {
+      await withoutActuallyEscaping(compute) {
+        _ = self
+        await $0()
+      }
+    }
+
+    func compute() async {}
+  }
+}