[region-isolation] Use SILIsolationInfo::initializeTrackableValue instead of SILIsolationInfo::mergeIsolationRegionInfo to fix last issue

When merging SILIsolationInfo for regions, we want to drop
nonisolated(unsafe). This is important since nonisolated(unsafe) should only
apply to the specific "value" that it belongs to, not the entire region.

This creates a problem since in a few places in the code base we initialize a
value (producing a disconnected value) and then initialize it by merging in an
actor isolation. This no longer work since we will then always have
nonisolated(unsafe) stripped, so no values would ever be considered to be
nonisolated(unsafe). After analyzing the use case, I realized that these were
just initialization patterns and in this commit, I added a specific
initialization operation called SILIsolationInfo::initializeTrackableValue and
eliminated those calls to SILIsolationInfo::mergeIsolationRegionInfo.

Since SILIsolationInfo no longer has any merge operation on it, I then
eliminated that code in this commit. This completes the behavior split that I
put into the type system in the last commit. Specifically, I defined a
composition type called SILDynamicMergedIsolationInfo. It represents a
SILIsolationInfo that has been merged... that is why I called it the
DynamicMergedIsolationInfo. It could probably use a better name = (.

This fixes one of the last weird test case that I wrote where we were not letting through valid
nonisolated(unsafe) code.

At the same time, I discovered an additional issue (which can be seen in the
TODOs in this commit), where we are being too conservative around a non-Sendable
class var field. I am going to fix that in the next commit.

rdar://128299305

Author: gottesmm

include/swift/SILOptimizer/Analysis/RegionAnalysis.h

@@ -159,12 +159,6 @@ class regionanalysisimpl::TrackableValueState {
     return getIsolationRegionInfo().getActorIsolation();
   }
 
-  void mergeIsolationRegionInfo(SILIsolationInfo newRegionInfo) {
-    // TODO: Remove this.
-    regionInfo =
-        getIsolationRegionInfo().merge(newRegionInfo).getIsolationInfo();
-  }
-
   void setDisconnectedNonisolatedUnsafe() {
     auto oldRegionInfo = getIsolationRegionInfo();
     assert(oldRegionInfo.isDisconnected());
@@ -386,9 +380,17 @@ class RegionAnalysisValueMap {
   TrackableValue
   getActorIntroducingRepresentative(SILInstruction *introducingInst,
                                     SILIsolationInfo isolation) const;
-  bool mergeIsolationRegionInfo(SILValue value, SILIsolationInfo isolation);
   bool valueHasID(SILValue value, bool dumpIfHasNoID = false);
   Element lookupValueID(SILValue value);
+
+  /// Initialize a TrackableValue with a SILIsolationInfo that we already know
+  /// instead of inferring.
+  ///
+  /// If we successfully initialize \p value with \p info, returns
+  /// {TrackableValue(), true}. If we already had a TrackableValue, we return
+  /// {TrackableValue(), false}.
+  std::pair<TrackableValue, bool>
+  initializeTrackableValue(SILValue value, SILIsolationInfo info) const;
 };
 
 class RegionAnalysisFunctionInfo {

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -938,21 +938,24 @@ struct PartitionOpEvaluator {
   ///
   /// The bool result is if it is captured by a closure element. That only is
   /// computed if \p sourceOp is non-null.
-  std::pair<SILDynamicMergedIsolationInfo, bool>
+  std::optional<std::pair<SILDynamicMergedIsolationInfo, bool>>
   getIsolationRegionInfo(Region region, Operand *sourceOp) const {
     bool isClosureCapturedElt = false;
-    SILDynamicMergedIsolationInfo isolationRegionInfo;
+    std::optional<SILDynamicMergedIsolationInfo> isolationRegionInfo =
+        SILDynamicMergedIsolationInfo();
 
     for (const auto &pair : p.range()) {
       if (pair.second == region) {
         isolationRegionInfo =
-            isolationRegionInfo.merge(getIsolationRegionInfo(pair.first));
+            isolationRegionInfo->merge(getIsolationRegionInfo(pair.first));
+        if (!isolationRegionInfo)
+          return {};
         if (sourceOp)
           isClosureCapturedElt |= isClosureCaptured(pair.first, sourceOp);
       }
     }
 
-    return {isolationRegionInfo, isClosureCapturedElt};
+    return {{isolationRegionInfo.value(), isClosureCapturedElt}};
   }
 
   /// Overload of \p getIsolationRegionInfo without an Operand.
@@ -1062,8 +1065,13 @@ struct PartitionOpEvaluator {
       Region transferredRegion = p.getRegion(transferredElement);
       bool isClosureCapturedElt = false;
       SILDynamicMergedIsolationInfo transferredRegionIsolation;
-      std::tie(transferredRegionIsolation, isClosureCapturedElt) =
+      auto pairOpt =
           getIsolationRegionInfo(transferredRegion, op.getSourceOp());
+      if (!pairOpt) {
+        handleUnknownCodePattern(op);
+        return;
+      }
+      std::tie(transferredRegionIsolation, isClosureCapturedElt) = *pairOpt;
 
       // Before we do anything, see if our dynamic isolation kind is the same as
       // the isolation info for our partition op. If they match, this is not a
@@ -1089,8 +1097,12 @@ struct PartitionOpEvaluator {
       // Mark op.getOpArgs()[0] as transferred.
       TransferringOperandState &state = operandToStateMap.get(op.getSourceOp());
       state.isClosureCaptured |= isClosureCapturedElt;
-      state.isolationInfo =
-          state.isolationInfo.merge(transferredRegionIsolation);
+      if (auto newInfo =
+              state.isolationInfo.merge(transferredRegionIsolation)) {
+        state.isolationInfo = *newInfo;
+      } else {
+        handleUnknownCodePattern(op);
+      }
       assert(state.isolationInfo && "Cannot have unknown");
       state.isolationHistory.pushCFGHistoryJoin(p.getIsolationHistory());
       auto *ptrSet = ptrSetFactory.get(op.getSourceOp());

include/swift/SILOptimizer/Utils/SILIsolationInfo.h

@@ -102,8 +102,6 @@ class ActorInstance {
   }
 };
 
-class SILDynamicMergedIsolationInfo;
-
 class SILIsolationInfo {
 public:
   /// The lattice is:
@@ -232,13 +230,6 @@ class SILIsolationInfo {
     return (kind == Task || kind == Actor) && bool(isolatedValue);
   }
 
-  /// Merge this SILIsolationInfo with \p other and produce a
-  /// SILDynamicMergedIsolationInfo that represents the dynamic isolation of a
-  /// value found by merging the value's region's isolation info. This causes
-  /// nonisolated(unsafe) to be dropped.
-  [[nodiscard]] SILDynamicMergedIsolationInfo
-  merge(SILIsolationInfo other) const;
-
   static SILIsolationInfo getDisconnected(bool isUnsafeNonIsolated) {
     return {Kind::Disconnected, isUnsafeNonIsolated};
   }
@@ -395,10 +386,12 @@ class SILDynamicMergedIsolationInfo {
   SILDynamicMergedIsolationInfo(SILIsolationInfo innerInfo)
       : innerInfo(innerInfo) {}
 
-  [[nodiscard]] SILDynamicMergedIsolationInfo
+  /// Returns nullptr only if both this isolation info and \p other are actor
+  /// isolated to incompatible actors.
+  [[nodiscard]] std::optional<SILDynamicMergedIsolationInfo>
   merge(SILIsolationInfo other) const;
 
-  [[nodiscard]] SILDynamicMergedIsolationInfo
+  [[nodiscard]] std::optional<SILDynamicMergedIsolationInfo>
   merge(SILDynamicMergedIsolationInfo other) const {
     return merge(other.getIsolationInfo());
   }

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -1500,18 +1500,30 @@ class PartitionOpTranslator {
     return valueMap.tryToTrackValue(value);
   }
 
+  /// If \p value already has state associated with it, return that. Otherwise
+  /// begin tracking \p value and initialize its isolation info to be \p.
+  ///
+  /// This is in contrast to tryToTrackValue which infers isolation
+  /// information. This should only be used in exceptional cases when one 100%
+  /// knows what the isolation is already. E.x.: a partial_apply.
+  std::optional<std::pair<TrackableValue, bool>>
+  initializeTrackedValue(SILValue value, SILIsolationInfo info) const {
+    auto trackedValuePair = valueMap.initializeTrackableValue(value, info);
+    if (trackedValuePair.first.isNonSendable()) {
+      assert(trackedValuePair.second);
+      return trackedValuePair;
+    }
+    assert(trackedValuePair.second);
+    return {};
+  }
+
   TrackableValue
   getActorIntroducingRepresentative(SILInstruction *introducingInst,
                                     SILIsolationInfo actorIsolation) const {
     return valueMap.getActorIntroducingRepresentative(introducingInst,
                                                       actorIsolation);
   }
 
-  bool mergeIsolationRegionInfo(SILValue value,
-                                SILIsolationInfo isolationRegion) {
-    return valueMap.mergeIsolationRegionInfo(value, isolationRegion);
-  }
-
   bool valueHasID(SILValue value, bool dumpIfHasNoID = false) {
     return valueMap.valueHasID(value, dumpIfHasNoID);
   }
@@ -1542,9 +1554,6 @@ class PartitionOpTranslator {
       if (tryApplyInst->getErrorBB()->getNumArguments() > 0) {
         foundResults.emplace_back(tryApplyInst->getErrorBB()->getArgument(0));
       }
-      for (auto indirectResults : tryApplyInst->getIndirectSILResults()) {
-        foundResults.emplace_back(indirectResults);
-      }
       return;
     }
 
@@ -1568,11 +1577,20 @@ class PartitionOpTranslator {
     }
 
     for (SILValue result : resultValues) {
-      if (auto value = tryToTrackValue(result)) {
-        assignResults.push_back(value->getRepresentative().getValue());
-        // TODO: Can we pass back a reference to value perhaps?
-        if (isolationInfo) {
-          mergeIsolationRegionInfo(result, isolationInfo);
+      if (isolationInfo) {
+        // We only get back result if it is non-Sendable.
+        if (auto nonSendableValue =
+                initializeTrackedValue(result, isolationInfo)) {
+          if (!nonSendableValue->second) {
+            builder.addUnknownPatternError(result);
+            continue;
+          }
+          assignResults.push_back(
+              nonSendableValue->first.getRepresentative().getValue());
+        }
+      } else {
+        if (auto value = tryToTrackValue(result)) {
+          assignResults.push_back(value->getRepresentative().getValue());
         }
       }
     }
@@ -1741,10 +1759,12 @@ class PartitionOpTranslator {
     if (pai->getFunctionType()->isSendable())
       return;
 
-    auto paiValue = tryToTrackValue(pai).value();
-    SILValue rep = paiValue.getRepresentative().getValue();
-    mergeIsolationRegionInfo(rep, actorIsolation);
-    translateSILAssignFresh(rep);
+    auto trackedValue = initializeTrackedValue(pai, actorIsolation);
+    assert(trackedValue && "Should not have happened already");
+    assert(
+        trackedValue->second &&
+        "Insert should have happened since this could not happen before this");
+    translateSILAssignFresh(trackedValue->first.getRepresentative().getValue());
   }
 
   void translateSILPartialApply(PartialApplyInst *pai) {
@@ -3258,6 +3278,27 @@ RegionAnalysisValueMap::getIsolationRegion(SILValue value) const {
   return iter->getSecond().getIsolationRegionInfo();
 }
 
+std::pair<TrackableValue, bool>
+RegionAnalysisValueMap::initializeTrackableValue(
+    SILValue value, SILIsolationInfo newInfo) const {
+  auto info = getUnderlyingTrackedValue(value);
+  value = info.value;
+
+  auto *self = const_cast<RegionAnalysisValueMap *>(this);
+  auto iter = self->equivalenceClassValuesToState.try_emplace(
+      value, TrackableValueState(equivalenceClassValuesToState.size()));
+
+  // If we did not insert, just return the already stored value.
+  if (!iter.second) {
+    return {{iter.first->first, iter.first->second}, false};
+  }
+
+  self->stateIndexToEquivalenceClass[iter.first->second.getID()] = value;
+  iter.first->getSecond().setIsolationRegionInfo(newInfo);
+
+  return {{iter.first->first, iter.first->second}, true};
+}
+
 /// If \p isAddressCapturedByPartialApply is set to true, then this value is
 /// an address that is captured by a partial_apply and we want to treat it as
 /// may alias.
@@ -3275,6 +3316,7 @@ TrackableValue RegionAnalysisValueMap::getTrackableValue(
     return {iter.first->first, iter.first->second};
   }
 
+  // If we did not insert, just return the already stored value.
   self->stateIndexToEquivalenceClass[iter.first->second.getID()] = value;
 
   // Otherwise, we need to compute our flags.
@@ -3409,20 +3451,10 @@ TrackableValue RegionAnalysisValueMap::getActorIntroducingRepresentative(
   // Otherwise, wire up the value.
   self->stateIndexToEquivalenceClass[iter.first->second.getID()] =
       introducingInst;
-  iter.first->getSecond().mergeIsolationRegionInfo(actorIsolation);
+  iter.first->getSecond().setIsolationRegionInfo(actorIsolation);
   return {iter.first->first, iter.first->second};
 }
 
-bool RegionAnalysisValueMap::mergeIsolationRegionInfo(
-    SILValue value, SILIsolationInfo actorIsolation) {
-  value = getUnderlyingTrackedValue(value).value;
-  auto iter = equivalenceClassValuesToState.find(value);
-  if (iter == equivalenceClassValuesToState.end())
-    return false;
-  iter->getSecond().mergeIsolationRegionInfo(actorIsolation);
-  return true;
-}
-
 bool RegionAnalysisValueMap::valueHasID(SILValue value, bool dumpIfHasNoID) {
   assert(getTrackableValue(value).isNonSendable() &&
          "Can only accept non-Sendable values");

lib/SILOptimizer/Utils/SILIsolationInfo.cpp

@@ -640,36 +640,6 @@ void SILIsolationInfo::print(llvm::raw_ostream &os) const {
   }
 }
 
-SILDynamicMergedIsolationInfo
-SILIsolationInfo::merge(SILIsolationInfo other) const {
-  // If we are greater than the other kind, then we are further along the
-  // lattice. We ignore the change.
-  if (unsigned(other.kind) < unsigned(kind))
-    return {*this};
-
-  // TODO: Make this failing mean that we emit an unknown SIL error instead of
-  // asserting.
-  assert((!other.isActorIsolated() || !isActorIsolated() ||
-          hasSameIsolation(other)) &&
-         "Actor can only be merged with the same actor");
-
-  // If we are both disconnected and other has the unsafeNonIsolated bit set,
-  // drop that bit and return that.
-  //
-  // DISCUSSION: We do not want to preserve the unsafe non isolated bit after
-  // merging. These bits should not propagate through merging and should instead
-  // always be associated with non-merged infos.
-  //
-  // TODO: We should really bake the above into the type system by having merged
-  // and non-merged SILIsolationInfo.
-  if (other.isDisconnected() && other.isUnsafeNonIsolated()) {
-    return other.withUnsafeNonIsolated(false);
-  }
-
-  // Otherwise, just return other.
-  return other;
-}
-
 bool SILIsolationInfo::hasSameIsolation(ActorIsolation actorIsolation) const {
   if (getKind() != Kind::Actor)
     return false;
@@ -846,28 +816,25 @@ bool SILIsolationInfo::isNonSendableType(SILType type, SILFunction *fn) {
 //                    MARK: SILDynamicMergedIsolationInfo
 //===----------------------------------------------------------------------===//
 
-SILDynamicMergedIsolationInfo
+std::optional<SILDynamicMergedIsolationInfo>
 SILDynamicMergedIsolationInfo::merge(SILIsolationInfo other) const {
   // If we are greater than the other kind, then we are further along the
   // lattice. We ignore the change.
   if (unsigned(other.getKind()) < unsigned(innerInfo.getKind()))
     return {*this};
 
-  // TODO: Make this failing mean that we emit an unknown SIL error instead of
-  // asserting.
-  assert((!other.isActorIsolated() || !innerInfo.isActorIsolated() ||
-          innerInfo.hasSameIsolation(other)) &&
-         "Actor can only be merged with the same actor");
+  // If we are both actor isolated and our isolations are not
+  // compatible... return None.
+  if (other.isActorIsolated() && innerInfo.isActorIsolated() &&
+      !innerInfo.hasSameIsolation(other))
+    return {};
 
   // If we are both disconnected and other has the unsafeNonIsolated bit set,
   // drop that bit and return that.
   //
   // DISCUSSION: We do not want to preserve the unsafe non isolated bit after
   // merging. These bits should not propagate through merging and should instead
   // always be associated with non-merged infos.
-  //
-  // TODO: We should really bake the above into the type system by having merged
-  // and non-merged SILIsolationInfo.
   if (other.isDisconnected() && other.isUnsafeNonIsolated()) {
     return other.withUnsafeNonIsolated(false);
   }