[sema] Work around a double curry thunk actor isolation inference bug that is a knock on effect of ced96aa.

Specifically, there is currently a bug in TypeCheckConcurrency.cpp where we do
not visit autoclosures. This causes us to never set the autoclosure's
ActorIsolation field like all other closures. For a long time we were able to
get away with this just by relying on the isolation of the decl context of the
autoclosure... but with the introduction of nonisolated(nonsending), we found
cases where the generated single curry autoclosure would necessarily be
different than its decl context (e.x.: a synchronous outer curry thunk that is
nonisolated that returns an inner curry thunk that is
nonisolated(nonsending)). This problem caused us to hit asserts later in the
compiler since the inner closure was actually nonisolated(nonsending), but we
were thinking that it should have been concurrent.

To work around this problem, I changed the type checker in
ced96aa to explicitly set the isolation of
single curry thunk autoclosures when it generates them. The reason why we did
this is that it made it so that we did not have to have a potential large source
break in 6.2 by changing TypeCheckConcurrency.cpp to visit all autoclosures it
has not been visiting.

This caused a follow on issue where since we were now inferring the inner
autoclosure to have the correct isolation, in cases where we were creating a
double curry thunk for an access to a global actor isolated field of a
non-Sendable non-global actor isolated nominal type, we would have the outer
curry thunk have unspecified isolation instead of main actor isolation. An
example of this is the following:

```swift
class A {
  var block:  @mainactor () -> Void = {}
}

class B {
  let a = A()

  func d() {
    a.block = c // Error! Passing task isolated 'self' to @mainactor closure.
  }

  @mainactor
  func c() {}
}
```

This was unintentional. To work around this, this commit changes the type
checker to explicitly set the double curry thunk isolation to the correct value
when the type checker generates the double curry thunk in the same manner as it
does for single curry thunks and validates that if we do set the value to
something explicitly that it has the same value as the single curry thunk.

rdar://152522631

Author: gottesmm

lib/Sema/CSApply.cpp

@@ -1313,6 +1313,21 @@ namespace {
       return callExpr;
     }
 
+    std::optional<ActorIsolation>
+    getIsolationFromFunctionType(FunctionType *thunkTy) {
+      switch (thunkTy->getIsolation().getKind()) {
+      case FunctionTypeIsolation::Kind::NonIsolated:
+      case FunctionTypeIsolation::Kind::Parameter:
+      case FunctionTypeIsolation::Kind::Erased:
+        return {};
+      case FunctionTypeIsolation::Kind::GlobalActor:
+        return ActorIsolation::forGlobalActor(
+            thunkTy->getIsolation().getGlobalActorType());
+      case FunctionTypeIsolation::Kind::NonIsolatedCaller:
+        return ActorIsolation::forCallerIsolationInheriting();
+      }
+    }
+
     /// Build a "{ args in base.fn(args) }" single-expression curry thunk.
     ///
     /// \param baseExpr The base expression to be captured, if warranted.
@@ -1355,19 +1370,8 @@ namespace {
       // we do not visit the inner autoclosure in the ActorIsolation checker
       // meaning that we do not properly call setActorIsolation on the
       // AbstractClosureExpr that we produce here.
-      switch (thunkTy->getIsolation().getKind()) {
-      case FunctionTypeIsolation::Kind::NonIsolated:
-      case FunctionTypeIsolation::Kind::Parameter:
-      case FunctionTypeIsolation::Kind::Erased:
-        break;
-      case FunctionTypeIsolation::Kind::GlobalActor:
-        thunk->setActorIsolation(ActorIsolation::forGlobalActor(
-            thunkTy->getIsolation().getGlobalActorType()));
-        break;
-      case FunctionTypeIsolation::Kind::NonIsolatedCaller:
-        thunk->setActorIsolation(
-            ActorIsolation::forCallerIsolationInheriting());
-        break;
+      if (auto isolation = getIsolationFromFunctionType(thunkTy)) {
+        thunk->setActorIsolation(*isolation);
       }
 
       cs.cacheType(thunk);
@@ -1536,6 +1540,31 @@ namespace {
         cs.cacheType(selfOpenedRef);
       }
 
+      auto outerActorIsolation = [&]() -> std::optional<ActorIsolation> {
+        auto resultType = outerThunkTy->getResult();
+        // Look through all optionals.
+        while (auto opt = resultType->getOptionalObjectType())
+          resultType = opt;
+        auto f =
+            getIsolationFromFunctionType(resultType->castTo<FunctionType>());
+        if (!f)
+          return {};
+
+        // If we have a non-async function and our "inferred" isolation is
+        // caller isolation inheriting, then we do not infer isolation and just
+        // use the default. The reason why we are doing this is:
+        //
+        // 1. nonisolated for synchronous functions is equivalent to
+        // nonisolated(nonsending).
+        //
+        // 2. There is a strong invariant in the compiler today that caller
+        // isolation inheriting is always async. By using nonisolated here, we
+        // avoid breaking that invariant.
+        if (!outerThunkTy->isAsync() && f->isCallerIsolationInheriting())
+          return {};
+        return f;
+      }();
+
       Expr *outerThunkBody = nullptr;
 
       // For an @objc optional member or a member found via dynamic lookup,
@@ -1566,6 +1595,11 @@ namespace {
         auto *innerThunk = buildSingleCurryThunk(
             selfOpenedRef, memberRef, cast<DeclContext>(member),
             outerThunkTy->getResult()->castTo<FunctionType>(), memberLocator);
+        assert((!outerActorIsolation ||
+                innerThunk->getActorIsolation().getKind() ==
+                    outerActorIsolation->getKind()) &&
+               "If we have an isolation for our double curry thunk it should "
+               "match our single curry thunk");
 
         // Rewrite the body to close the existential if warranted.
         if (hasOpenedExistential) {
@@ -1587,6 +1621,11 @@ namespace {
       outerThunk->setThunkKind(AutoClosureExpr::Kind::DoubleCurryThunk);
       outerThunk->setParameterList(
           ParameterList::create(ctx, SourceLoc(), selfParamDecl, SourceLoc()));
+
+      if (outerActorIsolation) {
+        outerThunk->setActorIsolation(*outerActorIsolation);
+      }
+
       cs.cacheType(outerThunk);
 
       return outerThunk;

lib/Sema/TypeCheckConcurrency.cpp

@@ -3219,6 +3219,65 @@ namespace {
       return LazyInitializerWalking::InAccessor;
     }
 
+    /// This function is a stripped down version of checkApply that only is
+    /// applied to curry thunks generated by the type checker that explicitly
+    /// have isolation put upon them by the typechecker to work around a bug in
+    /// 6.2. We do not perform any sort of actual inference... we only use it to
+    /// mark the apply as being isolation crossing if we have an autoclosure
+    /// with mismatching isolation.
+    ///
+    /// We take advantage that we only can have two types of isolation on such
+    /// an autoclosure, global actor isolation and nonisolated(nonsending).
+    ///
+    /// For more information, see the comment in buildSingleCurryThunk.
+    void perform62AutoclosureCurryThunkChecking(ApplyExpr *apply,
+                                                AutoClosureExpr *fn) {
+      // The isolation of the context that we are in.
+      std::optional<ActorIsolation> contextIsolation;
+      auto getContextIsolation = [&]() -> ActorIsolation {
+        if (contextIsolation)
+          return *contextIsolation;
+
+        auto declContext = const_cast<DeclContext *>(getDeclContext());
+        contextIsolation =
+            getInnermostIsolatedContext(declContext, getClosureActorIsolation);
+        return *contextIsolation;
+      };
+
+      std::optional<ActorIsolation> unsatisfiedIsolation;
+
+      // NOTE: Normally autoclosures did not have ActorIsolation set on it since
+      // we do not visit the function of the partial apply due to a bug. The
+      // only reason why it is set is b/c we are explicitly setting this in the
+      // type checker when we generate the single and double curry thunks.
+      auto fnTypeIsolation = fn->getActorIsolation();
+      if (fnTypeIsolation.isGlobalActor()) {
+        Type globalActor = fnTypeIsolation.getGlobalActor();
+        if (!(getContextIsolation().isGlobalActor() &&
+              getContextIsolation().getGlobalActor()->isEqual(globalActor)))
+          unsatisfiedIsolation = ActorIsolation::forGlobalActor(globalActor);
+      }
+
+      // If there was no unsatisfied actor isolation, we're done.
+      if (!unsatisfiedIsolation)
+        return;
+
+      // Record whether the callee isolation or the context isolation
+      // is preconcurrency, which is used later to downgrade errors to
+      // warnings in minimal checking.
+      auto calleeDecl = apply->getCalledValue(/*skipFunctionConversions=*/true);
+      bool preconcurrency =
+          getContextIsolation().preconcurrency() ||
+          (calleeDecl && getActorIsolation(calleeDecl).preconcurrency());
+      unsatisfiedIsolation =
+          unsatisfiedIsolation->withPreconcurrency(preconcurrency);
+
+      // At this point, we know a jump is made to the callee that yields
+      // an isolation requirement unsatisfied by the calling context, so
+      // set the unsatisfiedIsolationJump fields of the ApplyExpr appropriately
+      apply->setIsolationCrossing(getContextIsolation(), *unsatisfiedIsolation);
+    }
+
     PreWalkResult<Pattern *> walkToPatternPre(Pattern *pattern) override {
       // Walking into patterns leads to nothing good because then we
       // end up visiting the AccessorDecls of a top-level
@@ -3336,6 +3395,19 @@ namespace {
 
             partialApply->base->walk(*this);
 
+            // See if we have an autoclosure as our function. If so, check if we
+            // have a difference in isolation. If so, make this apply an
+            // isolation crossing apply.
+            //
+            // NOTE: This is just a work around for 6.2 to make checking of
+            // double curry thunks work correctly in the face of us not
+            // performing full type checking of autoclosures that are functions
+            // of the apply. We are doing this to make sure that we do not
+            // increase the surface area too much.
+            if (auto *fn = dyn_cast<AutoClosureExpr>(apply->getFn())) {
+              perform62AutoclosureCurryThunkChecking(apply, fn);
+            }
+
             return Action::SkipNode(expr);
           }
         }

test/Concurrency/double_curry_thunk.swift

@@ -0,0 +1,21 @@
+// RUN: %target-swift-frontend -typecheck -strict-concurrency=complete -target %target-swift-5.1-abi-triple %s
+
+// REQUIRES: concurrency
+
+// We used to crash on this when processing double curry thunks. Make sure that
+// we do not do crash in the future.
+extension AsyncStream {
+  @Sendable func myCancel() {
+  }
+  func myNext2(_ continuation: UnsafeContinuation<Element?, Never>) {
+  }
+  func myNext() async -> Element? {
+    await withTaskCancellationHandler {
+      unsafe await withUnsafeContinuation {
+        unsafe myNext2($0)
+      }
+    } onCancel: { [myCancel] in
+      myCancel()
+    }
+  }
+}

test/Concurrency/transfernonsendable.swift

@@ -2084,11 +2084,14 @@ func inferLocationOfCapturedTaskIsolatedSelfCorrectly() {
 
     func d() {
       a.block = c // expected-warning {{converting non-Sendable function value to '@MainActor @Sendable () -> Void' may introduce data races}}
-      // expected-tns-warning @-1 {{sending 'self' risks causing data races}}
-      // expected-tns-note @-2 {{task-isolated 'self' is captured by a main actor-isolated closure. main actor-isolated uses in closure may race against later nonisolated uses}}
+      // expected-tns-warning @-1 {{non-Sendable '@MainActor () -> ()'-typed result can not be returned from main actor-isolated function to nonisolated context}}
+      // expected-tns-note @-2 {{a function type must be marked '@Sendable' to conform to 'Sendable'}}
+      // expected-tns-warning @-3 {{sending 'self' risks causing data races}}
+      // expected-tns-note @-4 {{task-isolated 'self' is captured by a main actor-isolated closure. main actor-isolated uses in closure may race against later nonisolated uses}}
     }
 
     @MainActor
     func c() {}
   }
 }
+

test/Concurrency/transfernonsendable_global_actor.swift

@@ -338,3 +338,22 @@ func localCaptureDataRace5() {
 
   x = 2 // expected-tns-note {{access can happen concurrently}}
 }
+
+func inferLocationOfCapturedActorIsolatedSelfCorrectly() {
+  class A {
+    var block:  @MainActor () -> Void = {}
+  }
+  @CustomActor
+  class B {
+    let a = A()
+
+    func d() {
+      a.block = c // expected-warning {{converting non-Sendable function value to '@MainActor @Sendable () -> Void' may introduce data races}}
+      // expected-warning @-1 {{non-Sendable '@MainActor () -> ()'-typed result can not be returned from main actor-isolated function to global actor 'CustomActor'-isolated context}}
+      // expected-note @-2 {{a function type must be marked '@Sendable' to conform to 'Sendable'}}
+    }
+
+    @MainActor
+    func c() {}
+  }
+}