[region-isolation] Make sure that nonisolated(unsafe) works in all cases.

I made sure we match what we get without region isolation by turning off region
isolation in one of the test runs on the test for this.

There is one problem where for non-final classes with nonisolated(unsafe) var
fields, we currently do not properly squelch since I need to do more
infrastructure work. I am going to do that in the next commit.

rdar://128299305

Author: gottesmm

include/swift/SILOptimizer/Analysis/RegionAnalysis.h

@@ -163,6 +163,12 @@ class regionanalysisimpl::TrackableValueState {
     regionInfo = getIsolationRegionInfo().merge(newRegionInfo);
   }
 
+  void setDisconnectedNonisolatedUnsafe() {
+    auto oldRegionInfo = getIsolationRegionInfo();
+    assert(oldRegionInfo.isDisconnected());
+    regionInfo = oldRegionInfo.withUnsafeNonIsolated();
+  }
+
   Element getID() const { return Element(id); }
 
   void addFlag(TrackableValueFlag flag) { flagSet |= flag; }

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -616,6 +616,9 @@ SILValue createEmptyAndUndefValue(SILType ty, SILInstruction *insertionPoint,
 /// Check if a struct or its fields can have unreferenceable storage.
 bool findUnreferenceableStorage(StructDecl *decl, SILType structType,
                                 SILFunction *func);
+
+SILValue getInitOfTemporaryAllocStack(AllocStackInst *asi);
+
 } // end namespace swift
 
 #endif // SWIFT_SILOPTIMIZER_UTILS_INSTOPTUTILS_H

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -19,6 +19,7 @@
 #include "swift/Basic/LLVM.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILInstruction.h"
+#include "swift/SILOptimizer/Utils/InstOptUtils.h"
 #include "swift/SILOptimizer/Utils/SILIsolationInfo.h"
 
 #include "llvm/ADT/MapVector.h"
@@ -963,6 +964,14 @@ struct PartitionOpEvaluator {
     return Impl::getIsolationInfo(partitionOp);
   }
 
+  std::optional<Element> getElement(SILValue value) const {
+    return asImpl().getElement(value);
+  }
+
+  SILValue getRepresentative(SILValue value) const {
+    return asImpl().getRepresentative(value);
+  }
+
   /// Apply \p op to the partition op.
   void apply(const PartitionOp &op) const {
     if (shouldEmitVerboseLogging()) {
@@ -1018,10 +1027,16 @@ struct PartitionOpEvaluator {
       assert(p.isTrackingElement(op.getOpArgs()[0]) &&
              "Transfer PartitionOp's argument should already be tracked");
 
+      // Before we do any further work, see if we have a nonisolated(unsafe)
+      // element. In such a case, this is also not a real transfer point.
+      Element transferredElement = op.getOpArgs()[0];
+      if (getIsolationRegionInfo(transferredElement).isUnsafeNonIsolated()) {
+        return;
+      }
+
       // Otherwise, we need to merge our isolation region info with the
       // isolation region info of everything else in our region. This is the
       // dynamic isolation region info found by the dataflow.
-      Element transferredElement = op.getOpArgs()[0];
       Region transferredRegion = p.getRegion(transferredElement);
       bool isClosureCapturedElt = false;
       SILIsolationInfo transferredRegionIsolation;
@@ -1045,8 +1060,8 @@ struct PartitionOpEvaluator {
       // region.
       if (bool(transferredRegionIsolation) &&
           !transferredRegionIsolation.isDisconnected()) {
-        return handleTransferNonTransferrable(op, op.getOpArgs()[0],
-                                              transferredRegionIsolation);
+        return handleTransferNonTransferrableHelper(op, op.getOpArgs()[0],
+                                                    transferredRegionIsolation);
       }
 
       // Mark op.getOpArgs()[0] as transferred.
@@ -1136,6 +1151,24 @@ struct PartitionOpEvaluator {
           return;
       }
 
+      // If we have a temporary that is initialized with an unsafe nonisolated
+      // value... squelch the error like if we were that value.
+      //
+      // TODO: This goes away with opaque values.
+      if (SILValue equivalenceClassRep =
+              getRepresentative(transferringOp->get())) {
+        if (auto *asi = dyn_cast<AllocStackInst>(equivalenceClassRep)) {
+          if (SILValue value = getInitOfTemporaryAllocStack(asi)) {
+            if (auto elt = getElement(value)) {
+              SILIsolationInfo eltIsolationInfo = getIsolationRegionInfo(*elt);
+              if (eltIsolationInfo.isUnsafeNonIsolated()) {
+                return;
+              }
+            }
+          }
+        }
+      }
+
       // If our instruction does not have any isolation info associated with it,
       // it must be nonisolated. See if our function has a matching isolation to
       // our transferring operand. If so, we can squelch this.
@@ -1151,6 +1184,35 @@ struct PartitionOpEvaluator {
     // Ok, we actually need to emit a call to the callback.
     return handleLocalUseAfterTransfer(op, elt, transferringOp);
   }
+
+  // Private helper that squelches the error if our transfer instruction and our
+  // use have the same isolation.
+  void handleTransferNonTransferrableHelper(
+      const PartitionOp &op, Element elt,
+      SILIsolationInfo dynamicMergedIsolationInfo) const {
+    if (shouldTryToSquelchErrors()) {
+      // If we have a temporary that is initialized with an unsafe nonisolated
+      // value... squelch the error like if we were that value.
+      //
+      // TODO: This goes away with opaque values.
+      if (SILValue equivalenceClassRep =
+              getRepresentative(op.getSourceOp()->get())) {
+        if (auto *asi = dyn_cast<AllocStackInst>(equivalenceClassRep)) {
+          if (SILValue value = getInitOfTemporaryAllocStack(asi)) {
+            if (auto elt = getElement(value)) {
+              SILIsolationInfo eltIsolationInfo = getIsolationRegionInfo(*elt);
+              if (eltIsolationInfo.isUnsafeNonIsolated()) {
+                return;
+              }
+            }
+          }
+        }
+      }
+    }
+
+    // Ok, we actually need to emit a call to the callback.
+    return handleTransferNonTransferrable(op, elt, dynamicMergedIsolationInfo);
+  }
 };
 
 /// A base implementation that can be used to default initialize CRTP
@@ -1213,6 +1275,15 @@ struct PartitionOpEvaluatorBaseImpl : PartitionOpEvaluator<Subclass> {
     return SILIsolationInfo();
   }
 
+  /// If we are able to, return the element associated with \p value. If
+  /// unsupported, returns none.
+  std::optional<Element> getElement(SILValue value) const { return {}; }
+
+  /// If supported, returns the representative in \p value's equivalence
+  /// class. Returns an empty SILValue if this is unsupported or if it does not
+  /// have one.
+  SILValue getRepresentative(SILValue value) const { return SILValue(); }
+
   /// Check if the representative value of \p elt is closure captured at \p
   /// op.
   ///

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -35,6 +35,7 @@
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/Test.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
+#include "swift/SILOptimizer/Utils/InstOptUtils.h"
 #include "swift/SILOptimizer/Utils/PartitionUtils.h"
 #include "swift/SILOptimizer/Utils/VariableNameUtils.h"
 #include "llvm/ADT/DenseMap.h"
@@ -384,6 +385,7 @@ static UnderlyingTrackedValueInfo getUnderlyingTrackedValue(SILValue value) {
   assert(base);
   if (base->getType().isObject())
     return {getUnderlyingTrackedValue(base).value, visitor.actorIsolation};
+
   return {base, visitor.actorIsolation};
 }
 
@@ -2981,6 +2983,17 @@ bool BlockPartitionState::recomputeExitFromEntry(
       return translator.getValueMap().getIsolationRegion(elt);
     }
 
+    std::optional<Element> getElement(SILValue value) const {
+      return translator.getValueMap().getTrackableValue(value).getID();
+    }
+
+    SILValue getRepresentative(SILValue value) const {
+      return translator.getValueMap()
+          .getTrackableValue(value)
+          .getRepresentative()
+          .maybeGetValue();
+    }
+
     bool isClosureCaptured(Element elt, Operand *op) const {
       auto iter = translator.getValueForId(elt);
       if (!iter)

lib/SILOptimizer/Mandatory/TransferNonSendable.cpp

@@ -1616,6 +1616,17 @@ struct DiagnosticEvaluator final
     return info->getValueMap().getIsolationRegion(element);
   }
 
+  std::optional<Element> getElement(SILValue value) const {
+    return info->getValueMap().getTrackableValue(value).getID();
+  }
+
+  SILValue getRepresentative(SILValue value) const {
+    return info->getValueMap()
+        .getTrackableValue(value)
+        .getRepresentative()
+        .maybeGetValue();
+  }
+
   bool isClosureCaptured(Element element, Operand *op) const {
     auto value = info->getValueMap().maybeGetRepresentative(element);
     if (!value)