Merge pull request #75590 from gottesmm/pr-66aa73e45aa508bc763af6ccc128f36e93a9e144

[region-isolation] Fix handling of coroutine apply results.

Author: gottesmm

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -493,6 +493,19 @@ class PartitionOp {
            "Transfer needs a sourceInst");
   }
 
+  template <typename T>
+  PartitionOp(PartitionOpKind opKind, T collectionOfIndices,
+              SILInstruction *sourceInst = nullptr)
+      : opKind(opKind), opArgs(), source(sourceInst) {
+    assert(((opKind != PartitionOpKind::Transfer &&
+             opKind != PartitionOpKind::UndoTransfer) ||
+            sourceInst) &&
+           "Transfer needs a sourceInst");
+    for (Element elt : collectionOfIndices) {
+      opArgs.push_back(elt);
+    }
+  }
+
   PartitionOp(PartitionOpKind opKind, Element arg1, Operand *sourceOperand)
       : opKind(opKind), opArgs({arg1}), source(sourceOperand) {
     assert(((opKind != PartitionOpKind::Transfer &&
@@ -529,9 +542,10 @@ class PartitionOp {
     return PartitionOp(PartitionOpKind::Assign, destElt, srcElt, srcOperand);
   }
 
-  static PartitionOp AssignFresh(Element tgt,
+  template <typename T>
+  static PartitionOp AssignFresh(T collection,
                                  SILInstruction *sourceInst = nullptr) {
-    return PartitionOp(PartitionOpKind::AssignFresh, tgt, sourceInst);
+    return PartitionOp(PartitionOpKind::AssignFresh, collection, sourceInst);
   }
 
   static PartitionOp Transfer(Element tgt, Operand *transferringOp) {
@@ -1162,12 +1176,18 @@ struct PartitionOpEvaluator {
       p.assignElement(op.getOpArgs()[0], op.getOpArgs()[1]);
       return;
     }
-    case PartitionOpKind::AssignFresh:
-      assert(op.getOpArgs().size() == 1 &&
-             "AssignFresh PartitionOp should be passed 1 argument");
-
-      p.trackNewElement(op.getOpArgs()[0]);
+    case PartitionOpKind::AssignFresh: {
+      auto arrayRef = op.getOpArgs();
+
+      Element front = arrayRef.front();
+      p.trackNewElement(front);
+      arrayRef = arrayRef.drop_front();
+      for (auto x : arrayRef) {
+        p.trackNewElement(x);
+        p.assignElement(x, front);
+      }
       return;
+    }
     case PartitionOpKind::Transfer: {
       // NOTE: We purposely do not check here if a transferred value is already
       // transferred. Callers are expected to put a require for that

include/swift/SILOptimizer/Utils/SILIsolationInfo.h

@@ -140,15 +140,26 @@ class ActorInstance {
 /// matching.
 class SILIsolationInfo {
 public:
-  /// The lattice is:
+  /// This forms a lattice of semantics. The lattice progresses from left ->
+  /// right below:
   ///
   /// Unknown -> Disconnected -> TransferringParameter -> Task -> Actor.
   ///
-  /// Unknown means no information. We error when merging on it.
   enum Kind : uint8_t {
+    /// Unknown means no information. We error when merging on it.
     Unknown,
+
+    /// An entity with disconnected isolation can be freely transferred into
+    /// another isolation domain. These are associated with "use after transfer"
+    /// diagnostics.
     Disconnected,
+
+    /// An entity that is in the same region as a task-isolated value. Cannot be
+    /// sent into another isolation domain.
     Task,
+
+    /// An entity that is in the same region as an actor-isolated value. Cannot
+    /// be sent into another isolation domain.
     Actor,
   };
 

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -1159,8 +1159,16 @@ struct PartitionOpBuilder {
   Element getActorIntroducingRepresentative(SILIsolationInfo actorIsolation);
 
   void addAssignFresh(SILValue value) {
+    std::array<Element, 1> values = {lookupValueID(value)};
     currentInstPartitionOps.emplace_back(
-        PartitionOp::AssignFresh(lookupValueID(value), currentInst));
+        PartitionOp::AssignFresh(values, currentInst));
+  }
+
+  void addAssignFresh(ArrayRef<SILValue> values) {
+    auto transformedCollection = makeTransformRange(
+        values, [&](SILValue value) { return lookupValueID(value); });
+    currentInstPartitionOps.emplace_back(
+        PartitionOp::AssignFresh(transformedCollection, currentInst));
   }
 
   void addAssign(SILValue destValue, Operand *srcOperand) {
@@ -1733,23 +1741,18 @@ class PartitionOpTranslator {
       return;
     }
 
-    auto assignResultsRef = llvm::ArrayRef(assignResults);
-    SILValue front = assignResultsRef.front();
-    assignResultsRef = assignResultsRef.drop_front();
-
+    // If we do not have any non-Sendable srcs, then all of our results get one
+    // large fresh region.
     if (assignOperands.empty()) {
-      // If no non-sendable srcs, non-sendable tgts get a fresh region.
-      builder.addAssignFresh(front);
-    } else {
-      builder.addAssign(front, assignOperands.front().first);
+      builder.addAssignFresh(assignResults);
+      return;
     }
 
-    // Assign all targets to the target region.
-    while (assignResultsRef.size()) {
-      SILValue next = assignResultsRef.front();
-      assignResultsRef = assignResultsRef.drop_front();
-
-      builder.addAssign(next, assignOperands.front().first);
+    // Otherwise, we need to assign all of the results to be in the same region
+    // as the operands. Without losing generality, we just use the first
+    // non-Sendable one.
+    for (auto result : assignResults) {
+      builder.addAssign(result, assignOperands.front().first);
     }
   }
 

test/Concurrency/transfernonsendable.sil

@@ -98,6 +98,8 @@ actor MyActor {
   var klass: NonSendableKlass { get set }
 }
 
+sil @beginApplyMultipleResultCallee : $@yield_once @convention(thin) () -> (@yields @guaranteed NonSendableKlass, @yields @guaranteed NonSendableKlass)
+
 /////////////////
 // MARK: Tests //
 /////////////////
@@ -395,3 +397,19 @@ bb0:
   %9999 = tuple ()
   return %9999 : $()
 }
+
+// Make sure that we do not crash on this.
+//
+// We used to crash on this since we would want to assign the region of an
+// operand to the results... but we do not have one and have multiple
+// results. This doesn't normally happen with most applies since applies do not
+// have multiple results, so in such a case, we would just assign fresh and not
+// try to do the assignment for the rest of the values.
+sil [ossa] @handleNoOperandToAssignToResults : $@convention(thin) () -> () {
+bb0:
+  %0 = function_ref @beginApplyMultipleResultCallee : $@yield_once @convention(thin) () -> (@yields @guaranteed NonSendableKlass, @yields @guaranteed NonSendableKlass)
+  (%1, %2, %3) = begin_apply %0() : $@yield_once @convention(thin) () -> (@yields @guaranteed NonSendableKlass, @yields @guaranteed NonSendableKlass)
+  end_apply %3 as $()
+  %9999 = tuple ()
+  return %9999 : $()
+}