[region-isolation] Disconnected regions in global actor isolated function are not accessible again after call to nonisolated async function

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -1121,6 +1121,13 @@ struct PartitionOpEvaluator {
   /// if they need to.
   static SILLocation getLoc(Operand *op) { return Impl::getLoc(op); }
 
+  /// Some evaluators pass in mock operands that one cannot call getUser()
+  /// upon. So to allow for this, provide a routine that our impl can override
+  /// if they need to.
+  static SILIsolationInfo getIsolationInfo(const PartitionOp &partitionOp) {
+    return Impl::getIsolationInfo(partitionOp);
+  }
+
   /// Apply \p op to the partition op.
   void apply(const PartitionOp &op) const {
     if (shouldEmitVerboseLogging()) {
@@ -1192,13 +1199,9 @@ struct PartitionOpEvaluator {
       //
       // DISCUSSION: We couldn't not emit this earlier since we needed the
       // dynamic isolation info of our value.
-      if (transferredRegionIsolation.isActorIsolated()) {
-        if (auto calleeIsolationInfo =
-                SILIsolationInfo::get(op.getSourceInst())) {
-          if (transferredRegionIsolation.hasSameIsolation(
-                  calleeIsolationInfo)) {
-            return;
-          }
+      if (auto calleeIsolationInfo = getIsolationInfo(op)) {
+        if (transferredRegionIsolation.hasSameIsolation(calleeIsolationInfo)) {
+          return;
         }
       }
 
@@ -1291,7 +1294,7 @@ struct PartitionOpEvaluator {
   void handleLocalUseAfterTransferHelper(const PartitionOp &op, Element elt,
                                          Operand *transferringOp) const {
     if (shouldTryToSquelchErrors()) {
-      if (auto isolationInfo = SILIsolationInfo::get(op.getSourceInst())) {
+      if (auto isolationInfo = getIsolationInfo(op)) {
         if (isolationInfo.isActorIsolated() &&
             isolationInfo.hasSameIsolation(
                 SILIsolationInfo::get(transferringOp->getUser())))
@@ -1389,6 +1392,9 @@ struct PartitionOpEvaluatorBaseImpl : PartitionOpEvaluator<Subclass> {
 
   static SILLocation getLoc(SILInstruction *inst) { return inst->getLoc(); }
   static SILLocation getLoc(Operand *op) { return op->getUser()->getLoc(); }
+  static SILIsolationInfo getIsolationInfo(const PartitionOp &partitionOp) {
+    return SILIsolationInfo::get(partitionOp.getSourceInst());
+  }
 };
 
 /// A subclass of PartitionOpEvaluatorBaseImpl that doesn't have any special

lib/SILOptimizer/Utils/PartitionUtils.cpp

@@ -12,6 +12,7 @@
 
 #include "swift/SILOptimizer/Utils/PartitionUtils.h"
 
+#include "swift/AST/ASTWalker.h"
 #include "swift/AST/Expr.h"
 #include "swift/SIL/ApplySite.h"
 #include "swift/SIL/InstructionUtils.h"
@@ -70,6 +71,37 @@ getGlobalActorInitIsolation(SILFunction *fn) {
   return getActorIsolation(globalDecl);
 }
 
+static DeclRefExpr *getDeclRefExprFromExpr(Expr *expr) {
+  struct LocalWalker final : ASTWalker {
+    DeclRefExpr *result = nullptr;
+
+    PreWalkResult<Expr *> walkToExprPre(Expr *expr) override {
+      assert(!result && "Shouldn't have a result yet");
+
+      if (auto *dre = dyn_cast<DeclRefExpr>(expr)) {
+        result = dre;
+        return Action::Stop();
+      }
+
+      if (isa<CoerceExpr, MemberRefExpr, ImplicitConversionExpr, IdentityExpr>(
+              expr))
+        return Action::Continue(expr);
+
+      return Action::Stop();
+    }
+  };
+
+  LocalWalker walker;
+
+  if (auto *ae = dyn_cast<AssignExpr>(expr)) {
+    ae->getSrc()->walk(walker);
+  } else {
+    expr->walk(walker);
+  }
+
+  return walker.result;
+}
+
 SILIsolationInfo SILIsolationInfo::get(SILInstruction *inst) {
   if (auto fas = FullApplySite::isa(inst)) {
     if (auto crossing = fas.getIsolationCrossing()) {
@@ -176,15 +208,30 @@ SILIsolationInfo SILIsolationInfo::get(SILInstruction *inst) {
     // actor isolated method. Use the AST to compute the actor isolation and
     // check if we are self. If we are not self, we want this to be
     // disconnected.
-    if (auto *declRefExpr = cmi->getLoc().getAsASTNode<DeclRefExpr>()) {
-      if (auto isolation = swift::getActorIsolation(declRefExpr->getDecl())) {
-        if (isolation.isActorIsolated() &&
-            (isolation.getKind() != ActorIsolation::ActorInstance ||
-             isolation.getActorInstanceParameter() == 0)) {
-          auto actor = cmi->getOperand()->getType().isAnyActor()
-                           ? cmi->getOperand()
-                           : SILValue();
-          return SILIsolationInfo::getActorIsolated(cmi, actor, isolation);
+    if (auto *expr = cmi->getLoc().getAsASTNode<Expr>()) {
+      if (auto *dre = getDeclRefExprFromExpr(expr)) {
+        if (auto isolation = swift::getActorIsolation(dre->getDecl())) {
+          if (isolation.isActorIsolated() &&
+              (isolation.getKind() != ActorIsolation::ActorInstance ||
+               isolation.getActorInstanceParameter() == 0)) {
+            auto actor = cmi->getOperand()->getType().isAnyActor()
+                             ? cmi->getOperand()
+                             : SILValue();
+            return SILIsolationInfo::getActorIsolated(cmi, actor, isolation);
+          }
+        }
+
+        if (auto type = dre->getType()->getNominalOrBoundGenericNominal()) {
+          if (auto isolation = swift::getActorIsolation(type)) {
+            if (isolation.isActorIsolated() &&
+                (isolation.getKind() != ActorIsolation::ActorInstance ||
+                 isolation.getActorInstanceParameter() == 0)) {
+              auto actor = cmi->getOperand()->getType().isAnyActor()
+                               ? cmi->getOperand()
+                               : SILValue();
+              return SILIsolationInfo::getActorIsolated(cmi, actor, isolation);
+            }
+          }
         }
       }
     }
@@ -266,9 +313,15 @@ SILIsolationInfo SILIsolationInfo::get(SILInstruction *inst) {
   // of the actor.
   if (ApplyExpr *apply = inst->getLoc().getAsASTNode<ApplyExpr>()) {
     if (auto crossing = apply->getIsolationCrossing()) {
-      if (crossing->getCalleeIsolation().isActorIsolated())
+      auto calleeIsolation = crossing->getCalleeIsolation();
+      if (calleeIsolation.isActorIsolated()) {
         return SILIsolationInfo::getActorIsolated(
             SILValue(), SILValue(), crossing->getCalleeIsolation());
+      }
+
+      if (calleeIsolation.isNonisolated()) {
+        return SILIsolationInfo::getDisconnected();
+      }
     }
   }
 

test/Concurrency/experimental_feature_strictconcurrency.swift

@@ -28,12 +28,7 @@ struct Test2: TestProtocol { // expected-warning{{conformance of 'C2' to 'Sendab
 @MainActor
 func iterate(stream: AsyncStream<Int>) async {
   nonisolated(unsafe) var it = stream.makeAsyncIterator()
-  // FIXME: Region isolation should consider a value from a 'nonisolated(unsafe)'
-  // declaration to be in a disconnected region
 
-  // expected-region-isolation-warning @+3 {{transferring 'it' may cause a data race}}
-  // expected-region-isolation-note @+2 {{transferring disconnected 'it' to nonisolated callee could cause races in between callee nonisolated and local main actor-isolated uses}}
-  // expected-region-isolation-note @+1 {{use here could race}}
   while let element = await it.next() {
     print(element)
   }

test/Concurrency/transfernonsendable.swift

@@ -1552,9 +1552,8 @@ func functionArgumentIntoClosure(_ x: @escaping () -> ()) async {
   var c = NonSendableKlass()
   for _ in 0..<1024 {
     await useValueAsync(c) // expected-tns-warning {{transferring 'c' may cause a data race}}
-    // expected-tns-note @-1 {{transferring disconnected 'c' to nonisolated callee could cause races in between callee nonisolated and local main actor-isolated uses}}
-    // expected-tns-note @-2 {{use here could race}}
-    // expected-complete-warning @-3 {{passing argument of non-sendable type 'NonSendableKlass' outside of main actor-isolated context may introduce data races}}
+    // expected-tns-note @-1 {{transferring main actor-isolated 'c' to nonisolated callee could cause races between nonisolated and main actor-isolated uses}}
+    // expected-complete-warning @-2 {{passing argument of non-sendable type 'NonSendableKlass' outside of main actor-isolated context may introduce data races}}
     c = a.klass
   }
 }
@@ -1563,11 +1562,9 @@ func functionArgumentIntoClosure(_ x: @escaping () -> ()) async {
   let a = MainActorIsolatedKlass()
   var c = NonSendableKlass()
   for _ in 0..<1024 {
-    await useValueAsync(c) // expected-tns-warning 2{{transferring 'c' may cause a data race}}
-    // expected-tns-note @-1 {{transferring disconnected 'c' to nonisolated callee could cause races in between callee nonisolated and local main actor-isolated uses}}
-    // expected-tns-note @-2 {{use here could race}}
-    // expected-tns-note @-3 {{transferring main actor-isolated 'c' to nonisolated callee could cause races between nonisolated and main actor-isolated uses}}
-    // expected-complete-warning @-4 {{passing argument of non-sendable type 'NonSendableKlass' outside of main actor-isolated context may introduce data races}}
+    await useValueAsync(c) // expected-tns-warning {{transferring 'c' may cause a data race}}
+    // expected-tns-note @-1 {{transferring main actor-isolated 'c' to nonisolated callee could cause races between nonisolated and main actor-isolated uses}}
+    // expected-complete-warning @-2 {{passing argument of non-sendable type 'NonSendableKlass' outside of main actor-isolated context may introduce data races}}
     c = a.klassLet
   }
 }

test/Concurrency/transfernonsendable_global_actor.swift

@@ -97,11 +97,12 @@ private class NonSendableLinkedListNode<T> { // expected-complete-note 3{{}}
 @CustomActor func useCustomActor5() async {
   let x = NonSendableLinkedListNode<Int>()
 
-  await transferToNonIsolated(x) // expected-tns-warning {{transferring 'x' may cause a data race}}
-  // expected-tns-note @-1 {{transferring disconnected 'x' to nonisolated callee could cause races in between callee nonisolated and local global actor 'CustomActor'-isolated uses}}
-  // expected-complete-warning @-2 {{passing argument of non-sendable type 'NonSendableLinkedListNode<Int>' outside of global actor 'CustomActor'-isolated context may introduce data races}}
+  // This is ok since the nonisolated function cannot transfer x, so once we
+  // return x will be isolated again.
+  await transferToNonIsolated(x)
+  // expected-complete-warning @-1 {{passing argument of non-sendable type 'NonSendableLinkedListNode<Int>' outside of global actor 'CustomActor'-isolated context may introduce data races}}
 
-  useValue(x) // expected-tns-note {{use here could race}}
+  useValue(x)
 }
 
 private struct StructContainingValue { // expected-complete-note 2{{}}
@@ -113,11 +114,12 @@ private struct StructContainingValue { // expected-complete-note 2{{}}
   var x = StructContainingValue()
   x = StructContainingValue()
 
-  await transferToNonIsolated(x) // expected-tns-warning {{transferring 'x' may cause a data race}}
-  // expected-tns-note @-1 {{transferring disconnected 'x' to nonisolated callee could cause races in between callee nonisolated and local global actor 'CustomActor'-isolated uses}}
-  // expected-complete-warning @-2 {{passing argument of non-sendable type 'StructContainingValue' outside of global actor 'CustomActor'-isolated context may introduce data races}}
+  // This is ok since the nonisolated function cannot transfer x meaning after
+  // we return we know that x will be disconnected upon return as well.
+  await transferToNonIsolated(x)
+  // expected-complete-warning @-1 {{passing argument of non-sendable type 'StructContainingValue' outside of global actor 'CustomActor'-isolated context may introduce data races}}
 
-  useValue(x) // expected-tns-note {{use here could race}}
+  useValue(x)
 }
 
 @CustomActor func useCustomActor7() async {
@@ -135,12 +137,12 @@ private struct StructContainingValue { // expected-complete-note 2{{}}
   var x = (NonSendableLinkedList<Int>(), NonSendableLinkedList<Int>())
   x = (NonSendableLinkedList<Int>(), NonSendableLinkedList<Int>())
 
-  await transferToNonIsolated(x) // expected-tns-warning {{transferring 'x' may cause a data race}}
-  // expected-tns-note @-1 {{transferring disconnected 'x' to nonisolated callee could cause races in between callee nonisolated and local global actor 'CustomActor'-isolated uses}}
+  // This is safe since the nonisolated function cannot transfer x further.
+  await transferToNonIsolated(x)
+  // expected-complete-warning @-1 {{passing argument of non-sendable type '(NonSendableLinkedList<Int>, NonSendableLinkedList<Int>)' outside of global actor 'CustomActor'-isolated context may introduce data races}}
   // expected-complete-warning @-2 {{passing argument of non-sendable type '(NonSendableLinkedList<Int>, NonSendableLinkedList<Int>)' outside of global actor 'CustomActor'-isolated context may introduce data races}}
-  // expected-complete-warning @-3 {{passing argument of non-sendable type '(NonSendableLinkedList<Int>, NonSendableLinkedList<Int>)' outside of global actor 'CustomActor'-isolated context may introduce data races}}
 
-  useValue(x) // expected-tns-note {{use here could race}}
+  useValue(x)
 }
 
 @CustomActor func useCustomActor9() async {

test/Concurrency/transfernonsendable_nonisolated.swift

@@ -117,3 +117,16 @@ func callMyFinalActor_NonIsolatedSync_SyncUseAfter(_ a: MyFinalActor) async {
   a.syncNonisolated(x)
   useValueSync(x)
 }
+
+func callIsolatedFunction() async {
+  let x = NonSendable()
+  await useValueAsync(x)
+  useValueSync(x)
+}
+
+@MainActor
+func callMainActorIsolatedFunction() async {
+  let x = NonSendable()
+  await useValueAsync(x)
+  useValueSync(x)
+}

unittests/SILOptimizer/PartitionUtilsTest.cpp

@@ -57,6 +57,10 @@ struct MockedPartitionOpEvaluator final
   }
 
   static SILLocation getLoc(Operand *op) { return SILLocation::invalid(); }
+
+  static SILIsolationInfo getIsolationInfo(const PartitionOp &partitionOp) {
+    return {};
+  }
 };
 
 } // namespace
@@ -95,6 +99,10 @@ struct MockedPartitionOpEvaluatorWithFailureCallback final
   }
 
   static SILLocation getLoc(Operand *op) { return SILLocation::invalid(); }
+
+  static SILIsolationInfo getIsolationInfo(const PartitionOp &partitionOp) {
+    return {};
+  }
 };
 
 } // namespace