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

In this part of the code, we are attempting to merge all of the operands into
the same region and then assigning all non-Sendable results of the function to
that same region. The problem that was occuring here was a thinko due to the
control flow of the code here not separating nicely the case of whether or not
we had operands or not. Previously this did not matter, since we just used the
first result in such a case... but since we changed to assign to the first
operand element in some cases, it matters now. To fix this, I split the confused
logic into two different easy to follow control paths... one if we have operands
and one where we do not have an operand. In the case where we have a first
operand, we merge our elements into its region. If we do not have any operands,
then we just perform one large region assign fresh.

This was not exposed by code that used non-coroutines since in SIL only
coroutines today have multiple results.

rdar://132767643
(cherry picked from commit 541863d)

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

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -1158,8 +1158,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) {
@@ -1732,23 +1740,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

@@ -99,6 +99,8 @@ actor MyActor {
   var klass: NonSendableKlass { get set }
 }
 
+sil @beginApplyMultipleResultCallee : $@yield_once @convention(thin) () -> (@yields @guaranteed NonSendableKlass, @yields @guaranteed NonSendableKlass)
+
 /////////////////
 // MARK: Tests //
 /////////////////
@@ -396,3 +398,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 : $()
+}