[region-isolation] Ensure that actor methods without a result properly mark their operands as being actor derived.

I did this by abstracting the representative value of an equivalence class into
two cases: the normal case of actually having a value and a second case which is
used only to inject into a region an actor derived value.

rdar://119113563

Author: gottesmm

lib/SILOptimizer/Mandatory/TransferNonSendable.cpp

@@ -827,6 +827,87 @@ class TrackableValueState {
   SWIFT_DEBUG_DUMP { print(llvm::dbgs()); }
 };
 
+/// The representative value of the equivalence class that makes up a tracked
+/// value.
+///
+/// We use a wrapper struct here so that we can inject "fake" actor isolated
+/// values into the regions of values that become merged into an actor by
+/// calling a function without a non-sendable result.
+class RepresentativeValue {
+  friend llvm::DenseMapInfo<RepresentativeValue>;
+
+  using InnerType = PointerUnion<SILValue, SILInstruction *>;
+
+  /// If this is set to a SILValue then it is the actual represented value. If
+  /// it is set to a SILInstruction, then this is a "fake" representative value
+  /// used to inject actor isolatedness. The instruction stored is the
+  /// instruction that introduced the actor isolated-ness.
+  InnerType value;
+
+public:
+  RepresentativeValue() : value() {}
+  RepresentativeValue(SILValue value) : value(value) {}
+  RepresentativeValue(SILInstruction *actorRegionInst)
+      : value(actorRegionInst) {}
+
+  operator bool() const { return bool(value); }
+
+  void print(llvm::raw_ostream &os) const {
+    if (auto *inst = value.dyn_cast<SILInstruction *>()) {
+      os << "ActorRegionIntroducingInst: " << *inst;
+      return;
+    }
+
+    os << *value.get<SILValue>();
+  }
+
+  SILValue getValue() const { return value.get<SILValue>(); }
+  SILValue maybeGetValue() const { return value.dyn_cast<SILValue>(); }
+  SILInstruction *getActorRegionIntroducingInst() const {
+    return value.get<SILInstruction *>();
+  }
+
+  SWIFT_DEBUG_DUMP { print(llvm::dbgs()); }
+
+private:
+  RepresentativeValue(InnerType value) : value(value) {}
+};
+
+} // namespace
+
+namespace llvm {
+llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
+                              const RepresentativeValue &value) {
+  value.print(os);
+  return os;
+}
+
+template <>
+struct DenseMapInfo<::RepresentativeValue> {
+  using RepresentativeValue = ::RepresentativeValue;
+  using InnerType = RepresentativeValue::InnerType;
+  using InnerDenseMapInfo = DenseMapInfo<InnerType>;
+
+  static RepresentativeValue getEmptyKey() {
+    return RepresentativeValue(InnerDenseMapInfo::getEmptyKey());
+  }
+  static RepresentativeValue getTombstoneKey() {
+    return RepresentativeValue(InnerDenseMapInfo::getTombstoneKey());
+  }
+
+  static unsigned getHashValue(RepresentativeValue value) {
+    return InnerDenseMapInfo::getHashValue(value.value);
+  }
+
+  static bool isEqual(RepresentativeValue LHS, RepresentativeValue RHS) {
+    return InnerDenseMapInfo::isEqual(LHS.value, RHS.value);
+  }
+};
+
+} // namespace llvm
+
+namespace {
+
 /// A tuple consisting of a base value and its value state.
 ///
 /// DISCUSSION: We are computing regions among equivalence classes of values
@@ -843,11 +924,12 @@ class TrackableValueState {
 ///
 /// In the above example, %2 will be mapped to %0 by our value mapping.
 class TrackableValue {
-  SILValue representativeValue;
+  RepresentativeValue representativeValue;
   TrackableValueState valueState;
 
 public:
-  TrackableValue(SILValue representativeValue, TrackableValueState valueState)
+  TrackableValue(RepresentativeValue representativeValue,
+                 TrackableValueState valueState)
       : representativeValue(representativeValue), valueState(valueState) {}
 
   bool isMayAlias() const { return valueState.isMayAlias(); }
@@ -865,12 +947,12 @@ class TrackableValue {
   }
 
   /// Return the representative value of this equivalence class of values.
-  SILValue getRepresentative() const { return representativeValue; }
+  RepresentativeValue getRepresentative() const { return representativeValue; }
 
   void print(llvm::raw_ostream &os) const {
     os << "TrackableValue. State: ";
     valueState.print(os);
-    os << "\n    Rep Value: " << *getRepresentative();
+    os << "\n    Rep Value: " << getRepresentative();
   }
 
   SWIFT_DEBUG_DUMP { print(llvm::dbgs()); }
@@ -898,6 +980,8 @@ struct PartitionOpBuilder {
   TrackableValueID lookupValueID(SILValue value);
   bool valueHasID(SILValue value, bool dumpIfHasNoID = false);
 
+  TrackableValueID getActorIntroducingRepresentative();
+
   void addAssignFresh(SILValue value) {
     currentInstPartitionOps.emplace_back(
         PartitionOp::AssignFresh(lookupValueID(value), currentInst));
@@ -948,6 +1032,19 @@ struct PartitionOpBuilder {
         lookupValueID(fst), lookupValueID(snd), currentInst));
   }
 
+  /// Mark \p value artifically as being part of an actor isolated region by
+  /// introducing a new fake actor introducing representative and merging them.
+  void addActorIntroducingInst(SILValue value) {
+    assert(valueHasID(value, /*dumpIfHasNoID=*/true) &&
+           "merged values should already have been encountered");
+
+    auto elt = getActorIntroducingRepresentative();
+    currentInstPartitionOps.emplace_back(
+        PartitionOp::AssignFresh(elt, currentInst));
+    currentInstPartitionOps.emplace_back(
+        PartitionOp::Merge(lookupValueID(value), elt, currentInst));
+  }
+
   void addRequire(SILValue value) {
     assert(valueHasID(value, /*dumpIfHasNoID=*/true) &&
            "required value should already have been encountered");
@@ -990,8 +1087,9 @@ class PartitionOpTranslator {
   /// non-Sendable values. Implicit conversion from SILValue used pervasively.
   /// ensure getUnderlyingTrackedValue is called on SILValues before entering
   /// into this map
-  llvm::DenseMap<SILValue, TrackableValueState> equivalenceClassValuesToState;
-  llvm::DenseMap<unsigned, SILValue> stateIndexToEquivalenceClass;
+  llvm::DenseMap<RepresentativeValue, TrackableValueState>
+      equivalenceClassValuesToState;
+  llvm::DenseMap<unsigned, RepresentativeValue> stateIndexToEquivalenceClass;
 
   /// A list of values that can never be transferred.
   ///
@@ -1075,8 +1173,8 @@ class PartitionOpTranslator {
 
     // Check if our base is a ref_element_addr from an actor. In such a case,
     // mark this value as actor derived.
-    if (isa<LoadInst, LoadBorrowInst>(iter.first->first)) {
-      auto *svi = cast<SingleValueInstruction>(iter.first->first);
+    if (isa<LoadInst, LoadBorrowInst>(iter.first->first.getValue())) {
+      auto *svi = cast<SingleValueInstruction>(iter.first->first.getValue());
       auto storage = AccessStorageWithBase::compute(svi->getOperand(0));
       if (storage.storage && isa<RefElementAddrInst>(storage.base)) {
         if (storage.storage.getRoot()->getType().isActor()) {
@@ -1087,7 +1185,7 @@ class PartitionOpTranslator {
 
     // Check if we have an unsafeMutableAddressor from a global actor, mark the
     // returned value as being actor derived.
-    if (auto applySite = FullApplySite::isa(iter.first->first)) {
+    if (auto applySite = FullApplySite::isa(iter.first->first.getValue())) {
       if (auto *calleeFunction = applySite.getCalleeFunction()) {
         if (calleeFunction->isGlobalInit() &&
             isGlobalActorInit(calleeFunction)) {
@@ -1102,6 +1200,25 @@ class PartitionOpTranslator {
     return {iter.first->first, iter.first->second};
   }
 
+  TrackableValue
+  getActorIntroducingRepresentative(SILInstruction *introducingInst) const {
+    auto *self = const_cast<PartitionOpTranslator *>(this);
+    auto iter = self->equivalenceClassValuesToState.try_emplace(
+        introducingInst,
+        TrackableValueState(equivalenceClassValuesToState.size()));
+
+    // If we did not insert, just return the already stored value.
+    if (!iter.second) {
+      return {iter.first->first, iter.first->second};
+    }
+
+    // Otherwise, wire up the value.
+    self->stateIndexToEquivalenceClass[iter.first->second.getID()] =
+        introducingInst;
+    iter.first->getSecond().addFlag(TrackableValueFlag::isActorDerived);
+    return {iter.first->first, iter.first->second};
+  }
+
   bool markValueAsActorDerived(SILValue value) {
     value = getUnderlyingTrackedValue(value);
     auto iter = equivalenceClassValuesToState.find(value);
@@ -1300,7 +1417,7 @@ class PartitionOpTranslator {
   void dumpValues() const {
     // Since this is just used for debug output, be inefficient to make nicer
     // output.
-    std::vector<std::pair<unsigned, SILValue>> temp;
+    std::vector<std::pair<unsigned, RepresentativeValue>> temp;
     for (auto p : stateIndexToEquivalenceClass) {
       temp.emplace_back(p.first, p.second);
     }
@@ -1336,13 +1453,13 @@ class PartitionOpTranslator {
 
     for (SILValue src : sourceValues) {
       if (auto value = tryToTrackValue(src)) {
-        assignOperands.push_back(value->getRepresentative());
+        assignOperands.push_back(value->getRepresentative().getValue());
       }
     }
 
     for (SILValue result : resultValues) {
       if (auto value = tryToTrackValue(result)) {
-        assignResults.push_back(value->getRepresentative());
+        assignResults.push_back(value->getRepresentative().getValue());
         // TODO: Can we pass back a reference to value perhaps?
         if (options.contains(SILMultiAssignFlags::PropagatesActorSelf)) {
           markValueAsActorDerived(result);
@@ -1360,8 +1477,18 @@ class PartitionOpTranslator {
     }
 
     // If we do not have any non sendable results, return early.
-    if (assignResults.empty())
+    if (assignResults.empty()) {
+      // If we did not have any non-Sendable results and we did have
+      // non-Sendable operands and we are supposed to mark value as actor
+      // derived, introduce a fake element so we just propagate the actor
+      // region.
+      if (assignOperands.size() &&
+          options.contains(SILMultiAssignFlags::PropagatesActorSelf)) {
+        builder.addActorIntroducingInst(assignOperands.back());
+      }
+
       return;
+    }
 
     auto assignResultsRef = llvm::makeArrayRef(assignResults);
     SILValue front = assignResultsRef.front();
@@ -1426,7 +1553,8 @@ class PartitionOpTranslator {
                 }))
           continue;
 
-        builder.addUndoTransfer(trackedArgValue->getRepresentative(), ai);
+        builder.addUndoTransfer(trackedArgValue->getRepresentative().getValue(),
+                                ai);
       }
     }
   }
@@ -1444,8 +1572,9 @@ class PartitionOpTranslator {
     // useValue(x2)
     for (auto &op : ApplySite(pai).getArgumentOperands()) {
       if (auto trackedArgValue = tryToTrackValue(op.get())) {
-        builder.addRequire(trackedArgValue->getRepresentative());
-        builder.addTransfer(trackedArgValue->getRepresentative(), &op);
+        builder.addRequire(trackedArgValue->getRepresentative().getValue());
+        builder.addTransfer(trackedArgValue->getRepresentative().getValue(),
+                            &op);
       }
     }
   }
@@ -1538,19 +1667,19 @@ class PartitionOpTranslator {
     // require all operands
     for (auto op : applySite->getOperandValues())
       if (auto value = tryToTrackValue(op))
-        builder.addRequire(value->getRepresentative());
+        builder.addRequire(value->getRepresentative().getValue());
 
     auto handleSILOperands = [&](MutableArrayRef<Operand> operands) {
       for (auto &op : operands) {
         if (auto value = tryToTrackValue(op.get())) {
-          builder.addTransfer(value->getRepresentative(), &op);
+          builder.addTransfer(value->getRepresentative().getValue(), &op);
         }
       }
     };
 
     auto handleSILSelf = [&](Operand *self) {
       if (auto value = tryToTrackValue(self->get())) {
-        builder.addTransfer(value->getRepresentative(), self);
+        builder.addTransfer(value->getRepresentative().getValue(), self);
       }
     };
 
@@ -1568,7 +1697,7 @@ class PartitionOpTranslator {
     getApplyResults(*applySite, applyResults);
     for (auto result : applyResults)
       if (auto value = tryToTrackValue(result))
-        builder.addAssignFresh(value->getRepresentative());
+        builder.addAssignFresh(value->getRepresentative().getValue());
   }
 
   template <typename DestValues>
@@ -1636,8 +1765,8 @@ class PartitionOpTranslator {
       return;
     for (SILValue elt : collection) {
       if (auto trackableSrc = tryToTrackValue(elt)) {
-        builder.addMerge(trackableDest->getRepresentative(),
-                         trackableSrc->getRepresentative());
+        builder.addMerge(trackableDest->getRepresentative().getValue(),
+                         trackableSrc->getRepresentative().getValue());
       }
     }
   }
@@ -1701,7 +1830,7 @@ class PartitionOpTranslator {
 
   void translateSILRequire(SILValue val) {
     if (auto nonSendableVal = tryToTrackValue(val))
-      return builder.addRequire(nonSendableVal->getRepresentative());
+      return builder.addRequire(nonSendableVal->getRepresentative().getValue());
   }
 
   /// An enum select is just a multi assign.
@@ -2084,6 +2213,10 @@ TrackableValueID PartitionOpBuilder::lookupValueID(SILValue value) {
   return translator->lookupValueID(value);
 }
 
+TrackableValueID PartitionOpBuilder::getActorIntroducingRepresentative() {
+  return translator->getActorIntroducingRepresentative(currentInst).getID();
+}
+
 bool PartitionOpBuilder::valueHasID(SILValue value, bool dumpIfHasNoID) {
   return translator->valueHasID(value, dumpIfHasNoID);
 }
@@ -2129,7 +2262,7 @@ void PartitionOpBuilder::print(llvm::raw_ostream &os) const {
   sortUnique(opsToPrint);
   for (unsigned opArg : opsToPrint) {
     llvm::dbgs() << "          └╼ ";
-    SILValue value = translator->stateIndexToEquivalenceClass[opArg];
+    auto value = translator->stateIndexToEquivalenceClass[opArg];
     auto iter = translator->equivalenceClassValuesToState.find(value);
     assert(iter != translator->equivalenceClassValuesToState.end());
     llvm::dbgs() << "State: %%" << opArg << ". ";
@@ -2191,8 +2324,10 @@ class BlockPartitionState {
       auto iter = translator.getValueForId(element);
       if (!iter)
         return false;
-      return translator.isClosureCaptured(iter->getRepresentative(),
-                                          op->getUser());
+      auto value = iter->getRepresentative().maybeGetValue();
+      if (!value)
+        return false;
+      return translator.isClosureCaptured(value, op->getUser());
     };
     for (const auto &partitionOp : blockPartitionOps) {
       // By calling apply without providing a `handleFailure` closure, errors
@@ -2609,8 +2744,9 @@ class PartitionAnalysis {
               }
             }
 
-            auto rep =
-                translator.getValueForId(transferredVal)->getRepresentative();
+            auto rep = translator.getValueForId(transferredVal)
+                           ->getRepresentative()
+                           .getValue();
             LLVM_DEBUG(
                 llvm::dbgs()
                 << "    Emitting Use After Transfer Error!\n"
@@ -2630,7 +2766,8 @@ class PartitionAnalysis {
                        << "        ID:  %%" << transferredVal << "\n"
                        << "        Rep: "
                        << *translator.getValueForId(transferredVal)
-                               ->getRepresentative());
+                               ->getRepresentative()
+                               .getValue());
             diagnose(partitionOp,
                      diag::regionbasedisolation_selforargtransferred);
           };
@@ -2646,8 +2783,10 @@ class PartitionAnalysis {
         auto iter = translator.getValueForId(element);
         if (!iter)
           return false;
-        return translator.isClosureCaptured(iter->getRepresentative(),
-                                            op->getUser());
+        auto value = iter->getRepresentative().maybeGetValue();
+        if (!value)
+          return false;
+        return translator.isClosureCaptured(value, op->getUser());
       };
 
       // And then evaluate all of our partition ops on the entry partition.