[region-isolation] Assigns to var with structs/tuples with multiple fields should be merges for now.

The reason that this is being done is that since currently region based
isolation is not field sensitive, when we assign to the struct or tuple field of
the var, the new region relationship is set for the entire struct, not just a
specific field. This means that we effectively lose any region information from
the other fields. For example in the following at (1), given the current rules, we
lose that s.y was assigned to x:

```swift
struct S {
  var x: NonSendableKlass
  var y: NonSendableKlass
}

func foo() {
  var s = S()
  // Regions: [s]
  let x = NonSendableKlass()
  let y = NonSendableKlass()
  // Regions: [(s), (x), (y)]
  s.y = x
  // Regions: [(s, x), (y)]
  s.x = y                       (1)
  // Regions: [(x), (s, y)]
}
```

The best solution to this is to track such var like bindings in a field
sensitive manner where the regions of the aggregate are the union of the
individual fields. This would let us represent the regions of the above as
follows:

```swift
func foo() {
  var s = S()
  // Regions: [((s.x), (s.y))]
  let x = NonSendableKlass()
  let y = NonSendableKlass()
  // Regions: [((s.x), (s.y)), (x), (y)]
  s.y = x
  // Regions: [((s.x), (s.y, x)), (y)]
  s.x = y                       (1)
  // Regions: [((s.x, y), (s.y, x))]
}
```

We cannot do this today so to plug this hole, we instead treat these operations
as merges. This provides a conservative answer. Thus we would have:"

```swift
func foo() {
  var s = S()
  // Regions: [s]
  let x = NonSendableKlass()
  let y = NonSendableKlass()
  // Regions: [(s), (x), (y)]
  s.y = x
  // Regions: [(s, x), (y)]
  s.x = y                       (1)
  // Regions: [(s, x, y])
}
```

This is because we are treating the assignment into s.y and s.x as merges into
s, so we do not lose that y was assigned into s before we assigned y into
it. The unfortunate side-effect of this is that if a struct or tuple has
multiple fields, the merge makes it so that if we assign over the same field, we
do not lose the region of the old value:

```swift
func foo() {
  var s = S()
  // Regions: [s]
  let x = NonSendableKlass()
  let y = NonSendableKlass()
  // Regions: [(s), (x), (y)]
  s.y = x
  // Regions: [(s, x), (y)]
  s.y = y                    (1)
  // Regions: [(s, x, y])
}
```

If we were not to do this, then the s.y at (1) would result in s and x being in
different regions.

rdar://117273952

Author: gottesmm

lib/SILOptimizer/Mandatory/TransferNonSendable.cpp

@@ -55,6 +55,7 @@ namespace {
 
 struct UseDefChainVisitor
     : public AccessUseDefChainVisitor<UseDefChainVisitor, SILValue> {
+  bool isMerge = false;
 
   SILValue visitAll(SILValue sourceAddr) {
     SILValue result = visit(sourceAddr);
@@ -91,6 +92,8 @@ struct UseDefChainVisitor
 
   SILValue visitStorageCast(SingleValueInstruction *, Operand *sourceAddr,
                             AccessStorageCast castType) {
+    // If we do not have an identity cast, mark this as a merge.
+    isMerge |= castType != AccessStorageCast::Identity;
     return sourceAddr->get();
   }
 
@@ -108,13 +111,22 @@ struct UseDefChainVisitor
       case ProjectionKind::Class:
         llvm_unreachable("Shouldn't see this here");
       case ProjectionKind::Index:
+        // Index is always a merge.
+        isMerge = true;
         break;
       case ProjectionKind::Enum:
+        // Enum is never a merge since it always has a single field.
         break;
       case ProjectionKind::Tuple: {
+        // These are merges if we have multiple fields.
+        auto *tti = cast<TupleElementAddrInst>(inst);
+        isMerge |= tti->getOperand()->getType().getNumTupleElements() > 1;
         break;
       }
       case ProjectionKind::Struct:
+        // These are merges if we have multiple fields.
+        auto *sea = cast<StructElementAddrInst>(inst);
+        isMerge |= sea->getOperand()->getType().getNumNominalFields() > 1;
         break;
       }
     }
@@ -163,6 +175,13 @@ static Expr *getExprForPartitionOp(const PartitionOp &op) {
   return expr;
 }
 
+static bool isProjectedFromAggregate(SILValue value) {
+  assert(value->getType().isAddress());
+  UseDefChainVisitor visitor;
+  visitor.visitAll(value);
+  return visitor.isMerge;
+}
+
 //===----------------------------------------------------------------------===//
 //                           MARK: Main Computation
 //===----------------------------------------------------------------------===//
@@ -873,8 +892,18 @@ class PartitionOpTranslator {
   /// merges, to ensure any aliases of tgt are also updated.
   void translateSILStore(SILValue dest, SILValue src) {
     if (auto nonSendableTgt = tryToTrackValue(dest)) {
-      // Stores to unaliased storage can be treated as assignments, not merges.
-      if (nonSendableTgt.value().isNoAlias())
+      // In the following situations, we can perform an assign:
+      //
+      // 1. A store to unaliased storage.
+      // 2. A store that is to an entire value.
+      //
+      // DISCUSSION: If we have case 2, we need to merge the regions since we
+      // are not overwriting the entire region of the value. This does mean that
+      // we artificially include the previous region that was stored
+      // specifically in this projection... but that is better than
+      // miscompiling. For memory like this, we probably need to track it on a
+      // per field basis to allow for us to assign.
+      if (nonSendableTgt.value().isNoAlias() && !isProjectedFromAggregate(dest))
         return translateSILAssign(dest, src);
 
       // Stores to possibly aliased storage must be treated as merges.
@@ -982,8 +1011,13 @@ class PartitionOpTranslator {
     // NOTE: translateSILLookThrough asserts that this property is true.
     case SILInstructionKind::BeginAccessInst:
     case SILInstructionKind::BeginBorrowInst:
+    case SILInstructionKind::BeginDeallocRefInst:
+    case SILInstructionKind::RefToBridgeObjectInst:
+    case SILInstructionKind::BridgeObjectToRefInst:
     case SILInstructionKind::CopyValueInst:
+    case SILInstructionKind::EndCOWMutationInst:
     case SILInstructionKind::ProjectBoxInst:
+    case SILInstructionKind::EndInitLetRefInst:
     case SILInstructionKind::InitEnumDataAddrInst:
     case SILInstructionKind::OpenExistentialAddrInst:
     case SILInstructionKind::StructElementAddrInst:
@@ -993,6 +1027,11 @@ class PartitionOpTranslator {
     case SILInstructionKind::UpcastInst:
       return translateSILLookThrough(inst->getResult(0), inst->getOperand(0));
 
+    case SILInstructionKind::UnconditionalCheckedCastInst:
+      if (SILDynamicCastInst(inst).isRCIdentityPreserving())
+        return translateSILLookThrough(inst->getResult(0), inst->getOperand(0));
+      return translateSILAssign(inst);
+
     // Just make the result part of the operand's region without requiring.
     //
     // This is appropriate for things like object casts and object
@@ -1011,7 +1050,6 @@ class PartitionOpTranslator {
     case SILInstructionKind::OpenExistentialRefInst:
     case SILInstructionKind::PointerToAddressInst:
     case SILInstructionKind::ProjectBlockStorageInst:
-    case SILInstructionKind::RefElementAddrInst:
     case SILInstructionKind::RefToUnmanagedInst:
     case SILInstructionKind::StructExtractInst:
     case SILInstructionKind::TailAddrInst:
@@ -1021,6 +1059,14 @@ class PartitionOpTranslator {
     case SILInstructionKind::UncheckedEnumDataInst:
     case SILInstructionKind::UncheckedOwnershipConversionInst:
     case SILInstructionKind::UnmanagedToRefInst:
+
+    // RefElementAddrInst is not considered to be a lookThrough since we want to
+    // consider the address projected from the class to be a separate value that
+    // is in the same region as the parent operand. The reason that we want to
+    // do this is to ensure that if we assign into the ref_element_addr memory,
+    // we do not consider writes into the struct that contains the
+    // ref_element_addr to be merged into.
+    case SILInstructionKind::RefElementAddrInst:
       return translateSILAssign(inst);
 
     /// Enum inst is handled specially since if it does not have an argument,

test/Concurrency/sendnonsendable_basic.swift

@@ -44,6 +44,15 @@ func useValue<T>(_ x: T) {}
 
 var booleanFlag: Bool { false }
 
+struct SingleFieldKlassBox {
+  var k = NonSendableKlass()
+}
+
+struct TwoFieldKlassBox {
+  var k1 = NonSendableKlass()
+  var k2 = NonSendableKlass()
+}
+
 ////////////////////////////
 // MARK: Actor Self Tests //
 ////////////////////////////
@@ -187,10 +196,9 @@ extension Actor {
     let closure: () -> () = {
       print(self.klass)
     }
-    // NOTE: We do not error on this today since we assign into 1 and that makes
-    // x assign fresh. It will be fixed in a forthcoming commit.
+
     let x = (closure, 1)
-    await transferToMain(x)
+    await transferToMain(x) // expected-tns-warning {{call site passes `self` or a non-sendable argument of this function to another thread, potentially yielding a race with the caller}}
     // expected-complete-warning @-1 {{passing argument of non-sendable type '(() -> (), Int)' into main actor-isolated context may introduce data races}}
     // expected-complete-note @-2 {{a function type must be marked '@Sendable' to conform to 'Sendable'}}
   }
@@ -491,3 +499,107 @@ func testConversionsAndSendable(a: Actor, f: @Sendable () -> Void, f2: () -> Voi
   // expected-complete-warning @-1 {{passing argument of non-sendable type '() -> Void' into actor-isolated context may introduce data races}}
   // expected-complete-note @-2 {{a function type must be marked '@Sendable' to conform to 'Sendable'}}
 }
+
+///////////////////////////////////////////////
+// Multiple Field Var Assignment Merge Tests //
+///////////////////////////////////////////////
+
+func singleFieldVarMergeTest() async {
+  var box = SingleFieldKlassBox()
+  box = SingleFieldKlassBox()
+
+  // This transfers the entire region.
+  await transferToMain(box.k)
+
+  // But since box has only a single element, this doesn't race.
+  box.k = NonSendableKlass()
+  useValue(box.k)
+  useValue(box)
+
+  // We transfer the box back to main.
+  await transferToMain(box)
+
+  // And reassign over the entire box, so again we can use it again.
+  box = SingleFieldKlassBox()
+
+  useValue(box)
+  useValue(box.k)
+
+  await transferToMain(box) // expected-tns-warning {{passing argument of non-sendable type 'SingleFieldKlassBox' from nonisolated context to main actor-isolated context at this call site could yield a race with accesses later in this function}}
+
+  // But if we use box.k here, we emit an error since we didn't reinitialize at
+  // all.
+  useValue(box.k) // expected-note {{access here could race}}
+}
+
+func multipleFieldVarMergeTest1() async {
+  var box = TwoFieldKlassBox()
+  box = TwoFieldKlassBox()
+
+  // This transfers the entire region.
+  await transferToMain(box.k1) // expected-tns-warning {{passing argument of non-sendable type 'NonSendableKlass' from nonisolated context to main actor-isolated context at this call site could yield a race with accesses later in this function}}
+
+  // So even if we reassign over k1, since we did a merge, this should error.
+  box.k1 = NonSendableKlass() // expected-note {{access here could race}}
+  useValue(box) // expected-note {{access here could race}}
+}
+
+func multipleFieldVarMergeTest2() async {
+  var box = TwoFieldKlassBox()
+  box = TwoFieldKlassBox()
+
+  // This transfers the entire region.
+  await transferToMain(box.k1) 
+
+  // But if we assign over box completely, we can use it again.
+  box = TwoFieldKlassBox()
+
+  useValue(box.k1)
+  useValue(box.k2)
+  useValue(box)
+
+  await transferToMain(box.k2)
+
+  // But if we assign over box completely, we can use it again.
+  box = TwoFieldKlassBox()
+
+  useValue(box.k1)
+  useValue(box.k2)
+  useValue(box)
+}
+
+func multipleFieldTupleMergeTest1() async {
+  var box = (NonSendableKlass(), NonSendableKlass())
+  box = (NonSendableKlass(), NonSendableKlass())
+
+  // This transfers the entire region.
+  await transferToMain(box.0) // expected-tns-warning {{passing argument of non-sendable type 'NonSendableKlass' from nonisolated context to main actor-isolated context at this call site could yield a race with accesses later in this function}}
+
+  // So even if we reassign over k1, since we did a merge, this should error.
+  box.0 = NonSendableKlass() // expected-note {{access here could race}}
+  useValue(box) // expected-note {{access here could race}}
+}
+
+func multipleFieldTupleMergeTest2() async {
+  var box = (NonSendableKlass(), NonSendableKlass())
+
+  // This transfers the entire region.
+  await transferToMain(box.0)
+
+  let box2 = (NonSendableKlass(), NonSendableKlass())
+  // But if we assign over box completely, we can use it again.
+  box = box2
+
+  useValue(box.0)
+  useValue(box.1)
+  useValue(box)
+
+  await transferToMain(box.1)
+
+  // But if we assign over box completely, we can use it again.
+  box = (NonSendableKlass(), NonSendableKlass())
+
+  useValue(box.0)
+  useValue(box.1)
+  useValue(box)
+}