Tweak, improve, and debug the PartitionAnalysis engine until a fairly comprehensive suite of simple tests passes (region_based_sendability.swift)

Author: jturcotti

include/swift/AST/DiagnosticsSIL.def

@@ -856,11 +856,12 @@ NOTE(sil_referencebinding_inout_binding_here, none,
       "inout binding here",
       ())
 
-// SendNonSendable checking
-
-WARNING(send_non_sendable, none,
-        "Non-Sendable value %0 consumed, then used at this site; could be race",
-        (unsigned))
+// 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", ())
+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", ())
 
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -30,7 +30,7 @@ class PartitionOp {
   PartitionOpKind OpKind;
   llvm::SmallVector<unsigned, 2> OpArgs;
 
-  // record the SILInstruction that this PartitionOp was generated from, if
+  // Record the SILInstruction that this PartitionOp was generated from, if
   // generated during compilation from a SILBasicBlock
   SILInstruction *sourceInst;
 
@@ -75,7 +75,7 @@ class PartitionOp {
     return sourceInst;
   }
 
-  void dump() const {
+  void dump() const LLVM_ATTRIBUTE_USED {
     switch (OpKind) {
     case PartitionOpKind::Assign:
       llvm::dbgs() << "assign %" << OpArgs[0] << " = %" << OpArgs[1] << "\n";
@@ -110,6 +110,7 @@ static void horizontalUpdate(std::map<unsigned, signed> &map, unsigned key,
   }
 
   signed oldVal = map[key];
+  if (val == oldVal) return;
 
   for (auto [otherKey, otherVal] : map)
     if (otherVal == oldVal)
@@ -118,20 +119,49 @@ static void horizontalUpdate(std::map<unsigned, signed> &map, unsigned key,
 
 class Partition {
 private:
-  // label each index with a non-negative (unsigned) label if it is associated
+  // Label each index with a non-negative (unsigned) label if it is associated
   // with a valid region, and with -1 if it is associated with a consumed region
-  // in-order traversal relied upon
+  // in-order traversal relied upon.
   std::map<unsigned, signed> labels;
 
-  // track a label that is guaranteed to be fresh
+  // Track a label that is guaranteed to be strictly larger than all in use,
+  // and therefore safe for use as a fresh label.
   unsigned fresh_label = 0;
 
-  // in a canonical partition, all regions are labelled with the smallest index
+  // In a canonical partition, all regions are labelled with the smallest index
   // of any member. Certain operations like join and equals rely on canonicality
   // so when it's invalidated this boolean tracks that, and it must be
-  // reestablished by a call to canonicalize()
+  // reestablished by a call to canonicalize().
   bool canonical;
 
+  // Used only in assertions, check that Partitions promised to be canonical
+  // are actually canonical
+  bool is_canonical_correct() {
+    if (!canonical) return true; // vacuously correct
+
+    auto fail = [&](unsigned i, int type) {
+      llvm::dbgs() << "FAIL(i=" << i << "; type=" << type << "): ";
+      dump();
+      return false;
+    };
+
+    for (auto &[i, label] : labels) {
+      // correctness vacuous at consumed indices
+      if (label < 0) continue;
+
+      // this label should not exceed fresh_label
+      if (label >= fresh_label) return fail(i, 0);
+
+      // the label of a region should be at most as large as each index in it
+      if (i < label) return fail(i, 1);
+
+      // each region label should refer to an index in that region
+      if (labels[label] != label) return fail(i, 2);
+    }
+
+    return true;
+  }
+
   // linear time - For each region label that occurs, find the first index
   // at which it occurs and relabel all instances of it to that index.
   // This excludes the -1 label for consumed regions.
@@ -163,11 +193,32 @@ class Partition {
       // set fresh_label
       fresh_label = i + 1;
     }
+
+    assert(is_canonical_correct());
+  }
+
+  // linear time - merge the regions of two indices, maintaining canonicality
+  void merge(unsigned fst, unsigned snd) {
+    assert(labels.count(fst) && labels.count(snd));
+    if (labels[fst] == labels[snd])
+      return;
+
+    // maintain canonicality by renaming the greater-numbered region
+    if (labels[fst] < labels[snd])
+      horizontalUpdate(labels, snd, labels[fst]);
+    else
+      horizontalUpdate(labels, fst, labels[snd]);
+
+    assert(is_canonical_correct());
   }
 
 public:
   Partition() : labels({}), canonical(true) {}
 
+  // 1-arg constructor used when canonicality will be immediately invalidated,
+  // so set to false to begin with
+  Partition(bool canonical) : labels({}), canonical(canonical) {}
+
   static Partition singleRegion(std::vector<unsigned> indices) {
     Partition p;
     if (!indices.empty()) {
@@ -177,17 +228,9 @@ class Partition {
         p.labels[index] = min_index;
       }
     }
-    return p;
-  }
 
-  void dump() const {
-    llvm::dbgs() << "Partition";
-    if (canonical)
-      llvm::dbgs() << "(canonical)";
-    llvm::dbgs() << "{";
-    for (const auto &[i, label] : labels)
-      llvm::dbgs() << "[" << i << ": " << label << "] ";
-    llvm::dbgs() << "}\n";
+    assert(p.is_canonical_correct());
+    return p;
   }
 
   // linear time - Test two partititons for equality by first putting them
@@ -204,54 +247,52 @@ class Partition {
   // and two indices are in the same region of the join iff they are in the same
   // region in either operand.
   static Partition join(Partition &fst, Partition &snd) {
-    fst.canonicalize();
-    snd.canonicalize();
-
-    std::map<unsigned, signed> relabel_fst;
-    std::map<unsigned, signed> relabel_snd;
-    auto lookup_fst = [&](unsigned i) {
-      // signed to unsigned conversion... ?
-      return relabel_fst.count(fst.labels[i]) ? relabel_fst[fst.labels[i]]
-                                              : fst.labels[i];
-    };
-
-    auto lookup_snd = [&](unsigned i) {
-      // signed to unsigned conversion... safe?
-      return relabel_snd.count(snd.labels[i]) ? relabel_snd[snd.labels[i]]
-                                              : snd.labels[i];
-    };
-
-    for (const auto &[i, _] : fst.labels) {
-      // only consider indices present in both fst and snd
-      if (!snd.labels.count(i))
-        continue;
-
-      signed label_joined = std::min(lookup_fst(i), lookup_snd(i));
-
-      horizontalUpdate(relabel_fst, fst.labels[i], label_joined);
-      horizontalUpdate(relabel_snd, snd.labels[i], label_joined);
+    //ensure copies are made
+    Partition fst_reduced = false;
+    Partition snd_reduced = false;
+
+    // make canonical copies of fst and snd, reduced to their intersected domain
+    for (auto [i, _] : fst.labels)
+      if (snd.labels.count(i)) {
+        fst_reduced.labels[i] = fst.labels[i];
+        snd_reduced.labels[i] = snd.labels[i];
+      }
+    fst_reduced.canonicalize();
+    snd_reduced.canonicalize();
+
+    // merging each index in fst with its label in snd ensures that all pairs
+    // of indices that are in the same region in snd are also in the same region
+    // in fst - the desired property
+    for (const auto [i, snd_label] : snd_reduced.labels) {
+      if (snd_label < 0)
+        // if snd says that the region has been consumed, mark it consumed in fst
+        horizontalUpdate(fst_reduced.labels, i, -1);
+      else
+        fst_reduced.merge(i, snd_label);
     }
 
-    Partition joined;
-    joined.canonical = true;
-    for (const auto &[i, _] : fst.labels) {
-      if (!snd.labels.count(i))
-        continue;
-      signed label_i = lookup_fst(i);
-      joined.labels[i] = label_i;
-      joined.fresh_label = std::max(joined.fresh_label, (unsigned) label_i + 1);
-    }
+    assert(fst_reduced.is_canonical_correct());
 
-    return joined;
+    // fst_reduced is now the join
+    return fst_reduced;
   }
 
   // Apply the passed PartitionOp to this partition, performing its action.
   // A `handleFailure` closure can optionally be passed in that will be called
   // if a consumed region is required. The closure is given the PartitionOp that
   // failed, and the index of the SIL value that was required but consumed.
+  // Additionally, a list of "nonconsumable" indices can be passed in along with
+  // a handleConsumeNonConsumable closure. In the event that a region containing
+  // one of the nonconsumable indices is consumed, the closure will be called
+  // with the offending Consume.
   void apply(
-      PartitionOp op, llvm::function_ref<void(const PartitionOp&, unsigned)> handleFailure =
-                          [](const PartitionOp&, unsigned) {}) {
+      PartitionOp op,
+      llvm::function_ref<void(const PartitionOp&, unsigned)>
+      handleFailure = [](const PartitionOp&, unsigned) {},
+      std::vector<unsigned> nonconsumables = {},
+      llvm::function_ref<void(const PartitionOp&, unsigned)>
+      handleConsumeNonConsumable = [](const PartitionOp&, unsigned) {}
+  ) {
     switch (op.OpKind) {
     case PartitionOpKind::Assign:
       assert(op.OpArgs.size() == 2 &&
@@ -290,6 +331,19 @@ class Partition {
 
       // mark region as consumed
       horizontalUpdate(labels, op.OpArgs[0], -1);
+
+      // check if any nonconsumables were consumed, and handle the failure if so
+      for (unsigned nonconsumable : nonconsumables) {
+        assert(labels.count(nonconsumable) &&
+               "nonconsumables should be function args and self, and therefore"
+               "always present in the label map because of initialization at "
+               "entry");
+        if (labels[nonconsumable] < 0) {
+          handleConsumeNonConsumable(op, nonconsumable);
+          break;
+        }
+      }
+
       break;
     case PartitionOpKind::Merge:
       assert(op.OpArgs.size() == 2 &&
@@ -302,14 +356,7 @@ class Partition {
       if (labels[op.OpArgs[1]] < 0)
           handleFailure(op, op.OpArgs[1]);
 
-      if (labels[op.OpArgs[0]] == labels[op.OpArgs[1]])
-        break;
-
-      // maintain canonicality by renaming the greater-numbered region
-      if (labels[op.OpArgs[0]] < labels[op.OpArgs[1]])
-        horizontalUpdate(labels, op.OpArgs[1], labels[op.OpArgs[0]]);
-      else
-        horizontalUpdate(labels, op.OpArgs[0], labels[op.OpArgs[1]]);
+      merge(op.OpArgs[0], op.OpArgs[1]);
       break;
     case PartitionOpKind::Require:
       assert(op.OpArgs.size() == 1 &&
@@ -319,9 +366,21 @@ class Partition {
       if (labels[op.OpArgs[0]] < 0)
         handleFailure(op, op.OpArgs[0]);
     }
+
+    assert(is_canonical_correct());
+  }
+
+  void dump_labels() const LLVM_ATTRIBUTE_USED {
+    llvm::dbgs() << "Partition";
+    if (canonical)
+      llvm::dbgs() << "(canonical)";
+    llvm::dbgs() << "(fresh=" << fresh_label << "){";
+    for (const auto &[i, label] : labels)
+      llvm::dbgs() << "[" << i << ": " << label << "] ";
+    llvm::dbgs() << "}\n";
   }
 
-  void dump() {
+  void dump() LLVM_ATTRIBUTE_USED {
     std::map<signed, std::vector<unsigned>> buckets;
 
     for (auto [i, label] : labels) {
@@ -331,12 +390,15 @@ class Partition {
     llvm::dbgs() << "[";
     for (auto [label, indices] : buckets) {
       llvm::dbgs() << (label < 0 ? "{" : "(");
+      int j = 0;
       for (unsigned i : indices) {
-        llvm::dbgs() << i << " ";
+        llvm::dbgs() << (j++? " " : "") << i;
       }
       llvm::dbgs() << (label < 0 ? "}" : ")");
     }
-    llvm::dbgs() << "]\n";
+    llvm::dbgs() << "] | ";
+
+    dump_labels();
   }
 };
 }