[region-isolation] Add support for async let.

Specifically:

1. If the value is transferred such that it becomes part of an actor region, the
value is permanently part of the actor region as one would normally have.

2. If the value is just used in an async let or is used by a nonisolated async
function within the async let then while the async let is alive it cannot be
used. But once the async let has been awaited upon, we allow for it to be used
again.

rdar://117506395

Author: gottesmm

include/swift/AST/DiagnosticsSIL.def

@@ -877,16 +877,19 @@ NOTE(sil_referencebinding_inout_binding_here, none,
 
 // Warnings arising from the flow-sensitive checking of Sendability of
 // non-Sendable values
-WARNING(arg_region_transferred, none,
+WARNING(regionbasedisolation_selforargtransferred, none,
         "call site passes `self` or a non-sendable argument of this function to another thread, potentially yielding a race with the caller", ())
-WARNING(transfer_yields_race, none,
-        "non-sendable value sent across isolation domains that could be concurrently accessed later in this function (%0 access site%select{|s}1 displayed%select{|, %3 more hidden}2)",
-        (unsigned, bool, bool, unsigned))
-WARNING(call_site_transfer_yields_race, none,
+WARNING(regionbasedisolation_transfer_yields_race_no_isolation, none,
+        "transferred value of non-Sendable type %0 that could race with later accesses",
+        (Type))
+WARNING(regionbasedisolation_transfer_yields_race_with_isolation, none,
         "passing argument of non-sendable type %0 from %1 context to %2 context at this call site could yield a race with accesses later in this function",
         (Type, ActorIsolation, ActorIsolation))
-NOTE(possible_racy_access_site, none,
+NOTE(regionbasedisolation_maybe_race, none,
      "access here could race", ())
+ERROR(regionbasedisolation_unknown_pattern, none,
+      "pattern that the region based isolation checker does not understand how to check. Please file a bug",
+      ())
 
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -105,6 +105,10 @@ enum class PartitionOpKind : uint8_t {
   /// Transfer the region of a value if not already transferred, takes one arg.
   Transfer,
 
+  /// Due to an async let or something like that a value that was transferred is
+  /// no longer transferred.
+  UndoTransfer,
+
   /// Require the region of a value to be non-transferred, takes one arg.
   Require,
 };
@@ -128,20 +132,26 @@ class PartitionOp {
   PartitionOp(PartitionOpKind opKind, Element arg1,
               SILInstruction *sourceInst = nullptr)
       : opKind(opKind), opArgs({arg1}), source(sourceInst) {
-    assert((opKind != PartitionOpKind::Transfer || sourceInst) &&
+    assert(((opKind != PartitionOpKind::Transfer &&
+             opKind != PartitionOpKind::UndoTransfer) ||
+            sourceInst) &&
            "Transfer needs a sourceInst");
   }
 
   PartitionOp(PartitionOpKind opKind, Element arg1, Operand *sourceOperand)
       : opKind(opKind), opArgs({arg1}), source(sourceOperand) {
-    assert((opKind != PartitionOpKind::Transfer || sourceOperand) &&
+    assert(((opKind != PartitionOpKind::Transfer &&
+             opKind != PartitionOpKind::UndoTransfer) ||
+            sourceOperand) &&
            "Transfer needs a sourceInst");
   }
 
   PartitionOp(PartitionOpKind opKind, Element arg1, Element arg2,
               SILInstruction *sourceInst = nullptr)
       : opKind(opKind), opArgs({arg1, arg2}), source(sourceInst) {
-    assert((opKind != PartitionOpKind::Transfer || sourceInst) &&
+    assert(((opKind != PartitionOpKind::Transfer &&
+             opKind != PartitionOpKind::UndoTransfer) ||
+            sourceInst) &&
            "Transfer needs a sourceInst");
   }
 
@@ -162,6 +172,11 @@ class PartitionOp {
     return PartitionOp(PartitionOpKind::Transfer, tgt, transferringOp);
   }
 
+  static PartitionOp UndoTransfer(Element tgt,
+                                  SILInstruction *untransferringInst) {
+    return PartitionOp(PartitionOpKind::UndoTransfer, tgt, untransferringInst);
+  }
+
   static PartitionOp Merge(Element tgt1, Element tgt2,
                            SILInstruction *sourceInst = nullptr) {
     return PartitionOp(PartitionOpKind::Merge, tgt1, tgt2, sourceInst);
@@ -222,6 +237,14 @@ class PartitionOp {
       os << "%%" << opArgs[0];
       break;
     }
+    case PartitionOpKind::UndoTransfer: {
+      constexpr static char extraSpaceLiteral[10] = "    ";
+      os << "undo_transfer ";
+      if (extraSpace)
+        os << extraSpaceLiteral;
+      os << "%%" << opArgs[0];
+      break;
+    }
     case PartitionOpKind::Merge: {
       constexpr static char extraSpaceLiteral[10] = "       ";
       os << "merge ";
@@ -354,6 +377,33 @@ class Partition {
     return iter2.second;
   }
 
+  /// If val was marked as transferred, unmark it as transfer. Returns true if
+  /// we found that \p val was transferred. We return false otherwise.
+  bool undoTransfer(Element val) {
+    // First see if our val is tracked. If it is not tracked, insert it.
+    if (!isTracked(val)) {
+      elementToRegionMap.insert_or_assign(val, fresh_label);
+      fresh_label = Region(fresh_label + 1);
+      canonical = false;
+      return true;
+    }
+
+    // Otherwise, we already have this value in the map. Remove it from the
+    // transferred map.
+    auto iter1 = elementToRegionMap.find(val);
+    assert(iter1 != elementToRegionMap.end());
+    return regionToTransferredOpMap.erase(iter1->second);
+  }
+
+  void addElement(Element newElt) {
+    // Map index newElt to a fresh label.
+    elementToRegionMap.insert_or_assign(newElt, fresh_label);
+
+    // Increment the fresh label so it remains fresh.
+    fresh_label = Region(fresh_label + 1);
+    canonical = false;
+  }
+
   /// Construct the partition corresponding to the union of the two passed
   /// partitions.
   ///
@@ -784,12 +834,7 @@ struct PartitionOpEvaluator {
       assert(op.getOpArgs().size() == 1 &&
              "AssignFresh PartitionOp should be passed 1 argument");
 
-      // map index op.getOpArgs()[0] to a fresh label
-      p.elementToRegionMap.insert_or_assign(op.getOpArgs()[0], p.fresh_label);
-
-      // increment the fresh label so it remains fresh
-      p.fresh_label = Region(p.fresh_label + 1);
-      p.canonical = false;
+      p.addElement(op.getOpArgs()[0]);
       break;
     case PartitionOpKind::Transfer: {
       assert(op.getOpArgs().size() == 1 &&
@@ -817,6 +862,7 @@ struct PartitionOpEvaluator {
       // actor derived, we need to treat as nontransferrable.
       if (isActorDerived(op.getOpArgs()[0]))
         return handleTransferNonTransferrable(op, op.getOpArgs()[0]);
+
       Region elementRegion = p.elementToRegionMap.at(op.getOpArgs()[0]);
       if (llvm::any_of(p.elementToRegionMap,
                        [&](const std::pair<Element, Region> &pair) -> bool {
@@ -830,6 +876,16 @@ struct PartitionOpEvaluator {
       p.markTransferred(op.getOpArgs()[0], op.getSourceOp());
       break;
     }
+    case PartitionOpKind::UndoTransfer: {
+      assert(op.getOpArgs().size() == 1 &&
+             "UndoTransfer PartitionOp should be passed 1 argument");
+      assert(p.elementToRegionMap.count(op.getOpArgs()[0]) &&
+             "UndoTransfer PartitionOp's argument should already be tracked");
+
+      // Mark op.getOpArgs()[0] as not transferred.
+      p.undoTransfer(op.getOpArgs()[0]);
+      break;
+    }
     case PartitionOpKind::Merge:
       assert(op.getOpArgs().size() == 2 &&
              "Merge PartitionOp should be passed 2 arguments");