implement searching for a ConsumeReason instead of just diagnosing accesses to consumed values

Author: jturcotti

include/swift/AST/DiagnosticsSIL.def

@@ -859,9 +859,13 @@ NOTE(sil_referencebinding_inout_binding_here, none,
 // Warnings arising from the flow-sensitive checking of Sendability of
 // non-Sendable values
 WARNING(consumed_value_used, none,
-        "Non-Sendable value consumed, then used at this site; could yield race with another thread", ())
+        "non-Sendable value consumed, then used at this site; could yield race with another thread", ())
 WARNING(arg_region_consumed, none,
-        "This application could pass `self` or a Non-Sendable argument of this function to another thread, potentially yielding a race with the caller", ())
+        "this application could pass `self` or a Non-Sendable argument of this function to another thread, potentially yielding a race with the caller", ())
+WARNING(consumption_yields_race, none,
+        "non-Sendable value sent across isolation domains here, but could be accessed later in this function (%0 access site%select{|s}1 displayed%select{|, %3 more hidden}2)", (unsigned, bool, bool, unsigned))
+NOTE(possible_racy_access_site, none,
+    "access here could race with non-Sendable value send above", ())
 
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -81,7 +81,22 @@ class PartitionOp {
     return PartitionOp(PartitionOpKind::Require, tgt, sourceInst);
   }
 
-  SILInstruction *getSourceInst() const {
+  bool operator==(const PartitionOp& other) const = default;
+  // implemented for insertion into std::map
+  bool operator<(const PartitionOp& other) const {
+    if (OpKind != other.OpKind)
+      return OpKind < other.OpKind;
+    return sourceInst < other.sourceInst;
+  }
+
+  PartitionOpKind getKind() const {
+    return OpKind;
+  }
+
+  SILInstruction *getSourceInst(bool assertNonNull = false) const {
+    assert(!assertNonNull || sourceInst
+                             && "PartitionOps should be assigned SILInstruction"
+                                    " sources when used for the core analysis");
     return sourceInst;
   }
 
@@ -196,7 +211,7 @@ class Partition {
       return;
     canonical = true;
 
-    std::map<signed, unsigned> relabel;
+    std::map<unsigned, unsigned> relabel;
 
     // relies on in-order traversal of labels
     for (auto &[i, label] : labels) {
@@ -268,6 +283,14 @@ class Partition {
     return fst.labels == snd.labels;
   }
 
+  bool isTracked(unsigned val) const {
+    return labels.count(val);
+  }
+
+  bool isConsumed(unsigned val) const {
+    return isTracked(val) && labels.at(val) < 0;
+  }
+
   // quadratic time - Construct the partition corresponding to the join of the
   // two passed partitions; the join labels each index labelled by both operands
   // and two indices are in the same region of the join iff they are in the same
@@ -320,25 +343,17 @@ class Partition {
           nonconsumables = {},
 
       llvm::function_ref<void(const PartitionOp&, unsigned)>
-      handleConsumeNonConsumable = [](const PartitionOp&, unsigned) {},
-
-      bool reviveAfterFailure = false
+      handleConsumeNonConsumable = [](const PartitionOp&, unsigned) {}
   ) {
-    auto handleFailureAndRevive =
-        [&](const PartitionOp& partitionOp, unsigned consumedVal) {
-          if (reviveAfterFailure)
-            horizontalUpdate(labels, consumedVal, fresh_label++);
-          handleFailure(partitionOp, consumedVal);
-        };
     switch (op.OpKind) {
     case PartitionOpKind::Assign:
       assert(op.OpArgs.size() == 2 &&
              "Assign PartitionOp should be passed 2 arguments");
       assert(labels.count(op.OpArgs[1]) &&
              "Assign PartitionOp's source argument should be already tracked");
       // if assigning to a missing region, handle the failure
-      if (labels[op.OpArgs[1]] < 0)
-        handleFailureAndRevive(op, op.OpArgs[1]);
+      if (isConsumed(op.OpArgs[1]))
+        handleFailure(op, op.OpArgs[1]);
 
       labels[op.OpArgs[0]] = labels[op.OpArgs[1]];
 
@@ -361,8 +376,8 @@ class Partition {
              "Consume PartitionOp's argument should already be tracked");
 
       // if attempting to consume a consumed region, handle the failure
-      if (labels[op.OpArgs[0]] < 0)
-        handleFailureAndRevive(op, op.OpArgs[0]);
+      if (isConsumed(op.OpArgs[0]))
+        handleFailure(op, op.OpArgs[0]);
 
       // mark region as consumed
       horizontalUpdate(labels, op.OpArgs[0], -1);
@@ -387,10 +402,10 @@ class Partition {
              "Merge PartitionOp's arguments should already be tracked");
 
       // if attempting to merge a consumed region, handle the failure
-      if (labels[op.OpArgs[0]] < 0)
-        handleFailureAndRevive(op, op.OpArgs[0]);
-      if (labels[op.OpArgs[1]] < 0)
-        handleFailureAndRevive(op, op.OpArgs[1]);
+      if (isConsumed(op.OpArgs[0]))
+        handleFailure(op, op.OpArgs[0]);
+      if (isConsumed(op.OpArgs[1]))
+        handleFailure(op, op.OpArgs[1]);
 
       merge(op.OpArgs[0], op.OpArgs[1]);
       break;
@@ -399,13 +414,40 @@ class Partition {
              "Require PartitionOp should be passed 1 argument");
       assert(labels.count(op.OpArgs[0]) &&
              "Require PartitionOp's argument should already be tracked");
-      if (labels[op.OpArgs[0]] < 0)
-        handleFailureAndRevive(op, op.OpArgs[0]);
+      if (isConsumed(op.OpArgs[0]))
+        handleFailure(op, op.OpArgs[0]);
     }
 
     assert(is_canonical_correct());
   }
 
+  // return a vector of the consumed values in this partition
+  std::vector<unsigned> getConsumedVals() const {
+    // for effeciency, this could return an iterator not a vector
+    std::vector<unsigned> consumedVals;
+    for (auto [i, _] : labels)
+      if (isConsumed(i))
+        consumedVals.push_back(i);
+    return consumedVals;
+  }
+
+  // return a vector of the non-consumed regions in this partition, each
+  // represented as a vector of values
+  std::vector<std::vector<unsigned>> getNonConsumedRegions() const {
+    // for effeciency, this could return an iterator not a vector
+    std::map<signed, std::vector<unsigned>> buckets;
+
+    for (auto [i, label] : labels)
+      buckets[label].push_back(i);
+
+    std::vector<std::vector<unsigned>> doubleVec;
+
+    for (auto [_, bucket] : buckets)
+      doubleVec.push_back(bucket);
+
+    return doubleVec;
+  }
+
   void dump_labels() const LLVM_ATTRIBUTE_USED {
     llvm::dbgs() << "Partition";
     if (canonical)
@@ -419,9 +461,8 @@ class Partition {
   void dump() LLVM_ATTRIBUTE_USED {
     std::map<signed, std::vector<unsigned>> buckets;
 
-    for (auto [i, label] : labels) {
+    for (auto [i, label] : labels)
       buckets[label].push_back(i);
-    }
 
     llvm::dbgs() << "[";
     for (auto [label, indices] : buckets) {
@@ -432,9 +473,9 @@ class Partition {
       }
       llvm::dbgs() << (label < 0 ? "}" : ")");
     }
-    llvm::dbgs() << "]";
+    llvm::dbgs() << "]\n";
   }
 };
 }
 
-#endif
+#endif // SWIFT_PARTITIONUTILS_H