[send-non-sendable] Ensure that we properly warn if a field of a final actor is transferred.

Previously, we were not recognizing that a ref_element_addr from an actor object
is equivalent to the actor object and we shouldn't allow for it to be consumed.

rdar://115132118

Author: gottesmm

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -394,7 +394,7 @@ class Partition {
       llvm::function_ref<void(const PartitionOp &, Element)> handleFailure =
           [](const PartitionOp &, Element) {},
 
-      std::vector<Element> nonconsumables = {},
+      ArrayRef<Element> nonconsumables = {},
 
       llvm::function_ref<void(const PartitionOp &, Element)>
           handleConsumeNonConsumable = [](const PartitionOp &, Element) {}) {

lib/SILOptimizer/Mandatory/SendNonSendable.cpp

@@ -30,6 +30,8 @@
 
 using namespace swift;
 
+#pragma clang optimize off
+
 //===----------------------------------------------------------------------===//
 //                              MARK: Utilities
 //===----------------------------------------------------------------------===//
@@ -240,6 +242,12 @@ class PartitionOpTranslator {
   //       utilities would make it easier than handrolling
   llvm::DenseSet<SILValue> capturedUIValues;
 
+  /// A list of values that can never be consumed.
+  ///
+  /// This includes all function arguments as well as ref_element_addr from
+  /// actors.
+  std::vector<TrackableValueID> neverConsumedValueIDs;
+
   TrackableValue getTrackableValue(SILValue value) const {
     value = getUnderlyingTrackedValue(value);
 
@@ -280,6 +288,21 @@ class PartitionOpTranslator {
     if (!isNonSendableType(value->getType()))
       iter.first->getSecond().addFlag(TrackableValueFlag::isSendable);
 
+    // Check if our base is a ref_element_addr from an actor value. In such a
+    // case, add this id to the neverConsumedValueIDs.
+    if (isa<LoadInst, LoadBorrowInst>(iter.first->first)) {
+      auto *svi = cast<SingleValueInstruction>(iter.first->first);
+      auto storage = AccessStorageWithBase::compute(svi->getOperand(0));
+      if (storage.storage && isa<RefElementAddrInst>(storage.base)) {
+        if (storage.storage.getRoot()->getType().getASTType()->isActorType()) {
+          self->neverConsumedValueIDs.push_back(iter.first->second.getID());
+        }
+      }
+    }
+
+    // If our access storage is from a class, then see if we have an actor. In
+    // such a case, we need to add this id to the neverConsumed set.
+
     return {iter.first->first, iter.first->second};
   }
 
@@ -417,12 +440,9 @@ class PartitionOpTranslator {
   std::vector<TrackableValueID> getArgIDs() {
     std::vector<TrackableValueID> argIDs;
 
-    for (SILArgument *arg : function->getArguments())
-      if (auto state = trackIfNonSendable(arg))
-        argIDs.push_back(state->getID());
-
-    if (auto *selfArg = function->maybeGetSelfArgument()) {
-      if (auto state = trackIfNonSendable(selfArg)) {
+    for (SILArgument *arg : function->getArguments()) {
+      if (auto state = trackIfNonSendable(arg)) {
+        neverConsumedValueIDs.push_back(state->getID());
         argIDs.push_back(state->getID());
       }
     }
@@ -442,11 +462,8 @@ class PartitionOpTranslator {
   // Get the vector of IDs that cannot be legally consumed at any point in
   // this function. Since we place all args and self in a single region right
   // now, it is only necessary to choose a single representative of the set.
-  std::vector<TrackableValueID> getNonConsumables() {
-    if (const auto &argIDs = getArgIDs(); !argIDs.empty()) {
-      return {argIDs.front()};
-    }
-    return {};
+  ArrayRef<TrackableValueID> getNeverConsumedValues() {
+    return llvm::makeArrayRef(neverConsumedValueIDs);
   }
 
   // get the results of an apply instruction. This is the single result value
@@ -1049,7 +1066,7 @@ class BlockPartitionState {
     Partition workingPartition = entryPartition;
     for (auto &partitionOp : blockPartitionOps) {
       workingPartition.apply(partitionOp, handleFailure,
-                             translator.getNonConsumables(),
+                             translator.getNeverConsumedValues(),
                              handleConsumeNonConsumable);
     }
   }

test/Concurrency/sendnonsendable_basic.swift

@@ -15,10 +15,48 @@ class NonSendableKlass {
 actor Actor {
   var klass = NonSendableKlass()
   final var finalKlass = NonSendableKlass()
+
+  func useKlass(_ x: NonSendableKlass) {}
 }
 
 final actor FinalActor {
   var klass = NonSendableKlass()
+  func useKlass(_ x: NonSendableKlass) {}
+}
+
+func useInOut<T>(_ x: inout T) {}
+func useValue<T>(_ x: T) {}
+
+////////////////////////////
+// MARK: Actor Self Tests //
+////////////////////////////
+
+extension Actor {
+  func noWarningIfCallingGetter() async {
+    // We do not emit a warning in this case since we are calling transferring
+    // the result of a getter of an actor. The result of this getter could not
+    // have been returned from the actor without us emitting a warning as shown
+    // by the test klassGetter below.
+    await self.klass.asyncCall()
+  }
+
+  func warningIfCallingAsyncOnFinalField() async {
+    // Since we are calling finalKlass directly, we emit a warning here.
+    await self.finalKlass.asyncCall() // expected-warning {{}}
+  }
+
+  // We warn on this since we are escaping an actor field out of its isolation
+  // domain.
+  var klassGetter: NonSendableKlass {
+    self.finalKlass
+  }
+}
+
+extension FinalActor {
+  func warningIfCallingAsyncOnFinalField() async {
+    // Since our whole class is final, we emit the error directly here.
+    await self.klass.asyncCall() // expected-warning {{}}
+  }
 }
 
 /////////////////////////////
@@ -30,3 +68,61 @@ func formClosureWithoutCrashing() {
   var c = NonSendableKlass() // expected-warning {{variable 'c' was never mutated; consider changing to 'let' constant}}
   let _ = { print(c) }
 }
+
+// In this test, closure is not promoted into its box form. As a result, we
+// treat assignments into contents as a merge operation rather than an assign.
+func closureInOut(_ a: Actor) async {
+  var contents = NonSendableKlass()
+  let ns0 = NonSendableKlass()
+  let ns1 = NonSendableKlass()
+
+  contents = ns0
+  contents = ns1
+
+  var closure = {}
+  closure = { useInOut(&contents) }
+
+  await a.useKlass(ns0) // expected-warning {{passing argument of non-sendable type 'NonSendableKlass' from nonisolated context to actor-isolated context at this call site could yield a race with accesses later in this function}}
+  await a.useKlass(ns1) // expected-note {{access here could race}}
+
+  closure() // expected-note {{access here could race}}
+}
+
+func closureInOut2(_ a: Actor) async {
+  var contents = NonSendableKlass()
+  let ns0 = NonSendableKlass()
+  let ns1 = NonSendableKlass()
+
+  contents = ns0
+  contents = ns1
+
+  var closure = {}
+
+  await a.useKlass(ns0) // expected-warning {{passing argument of non-sendable type 'NonSendableKlass' from nonisolated context to actor-isolated context at this call site could yield a race with accesses later in this function}}
+
+  closure = { useInOut(&contents) } // expected-note {{access here could race}}
+
+  await a.useKlass(ns1) // expected-note {{access here could race}}
+
+  closure()
+}
+
+func closureNonInOut(_ a: Actor) async {
+  var contents = NonSendableKlass()
+  let ns0 = NonSendableKlass()
+  let ns1 = NonSendableKlass()
+
+  contents = ns0
+  contents = ns1
+
+  var closure = {}
+
+  await a.useKlass(ns0)
+
+  closure = { useValue(contents) }
+
+  await a.useKlass(ns1) // expected-warning {{passing argument of non-sendable type 'NonSendableKlass' from nonisolated context to actor-isolated context at this call site could yield a race with accesses later in this function}}
+
+  closure() // expected-note {{access here could race}}
+}
+