Begin to fill in the SendNonSendable SIL pass using a PartitionAnalysis engine, works for some simple examples but needs to address address values in SIL.

Author: jturcotti

include/swift/AST/DiagnosticsSIL.def

@@ -856,5 +856,11 @@ 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))
+
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -1,52 +1,98 @@
 #ifndef SWIFT_PARTITIONUTILS_H
 #define SWIFT_PARTITIONUTILS_H
 
-#include "llvm/Support/Debug.h"
+#include "swift/Basic/LLVM.h"
+#include "swift/SIL/SILInstruction.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Debug.h"
+#include <algorithm>
 
 namespace swift {
 
 enum class PartitionOpKind : uint8_t {
+  // Assign one value to the region of another, takes two args, second arg
+  // must already be tracked with a non-consumed region
   Assign,
+  // Assign one value to a fresh region, takes one arg.
   AssignFresh,
+  // Consume the region of a value, takes one arg
   Consume,
   Merge,
   Require
 };
 
+// PartitionOp represents a primitive operation that can be performed on
+// Partitions. This is part of the SendNonSendable SIL pass workflow:
+// first SILBasicBlocks are compiled to vectors of PartitionOps, then a fixed
+// point partition is found over the CFG.
 class PartitionOp {
 private:
   PartitionOpKind OpKind;
   llvm::SmallVector<unsigned, 2> OpArgs;
 
-  // TODO: can the following two declarations be merged?
-  PartitionOp(PartitionOpKind OpKind, unsigned arg1)
-      : OpKind(OpKind), OpArgs({arg1}) {}
+  // record the SILInstruction that this PartitionOp was generated from, if
+  // generated during compilation from a SILBasicBlock
+  SILInstruction *sourceInst;
+
+  // TODO: can the following declarations be merged?
+  PartitionOp(PartitionOpKind OpKind, unsigned arg1,
+              SILInstruction *sourceInst = nullptr)
+      : OpKind(OpKind), OpArgs({arg1}), sourceInst(sourceInst) {}
 
-  PartitionOp(PartitionOpKind OpKind, unsigned arg1, unsigned arg2)
-      : OpKind(OpKind), OpArgs({arg1, arg2}) {}
+  PartitionOp(PartitionOpKind OpKind, unsigned arg1, unsigned arg2,
+              SILInstruction *sourceInst = nullptr)
+      : OpKind(OpKind), OpArgs({arg1, arg2}), sourceInst(sourceInst) {}
 
   friend class Partition;
 
 public:
-  static PartitionOp Assign(unsigned tgt, unsigned src) {
-    return PartitionOp(PartitionOpKind::Assign, tgt, src);
+  static PartitionOp Assign(unsigned tgt, unsigned src,
+                            SILInstruction *sourceInst = nullptr) {
+    return PartitionOp(PartitionOpKind::Assign, tgt, src, sourceInst);
+  }
+
+  static PartitionOp AssignFresh(unsigned tgt,
+                                 SILInstruction *sourceInst = nullptr) {
+    return PartitionOp(PartitionOpKind::AssignFresh, tgt, sourceInst);
   }
 
-  static PartitionOp AssignFresh(unsigned tgt) {
-    return PartitionOp(PartitionOpKind::AssignFresh, tgt);
+  static PartitionOp Consume(unsigned tgt,
+                             SILInstruction *sourceInst = nullptr) {
+    return PartitionOp(PartitionOpKind::Consume, tgt, sourceInst);
   }
 
-  static PartitionOp Consume(unsigned tgt) {
-    return PartitionOp(PartitionOpKind::Consume, tgt);
+  static PartitionOp Merge(unsigned tgt1, unsigned tgt2,
+                           SILInstruction *sourceInst = nullptr) {
+    return PartitionOp(PartitionOpKind::Merge, tgt1, tgt2, sourceInst);
   }
 
-  static PartitionOp Merge(unsigned tgt1, unsigned tgt2) {
-    return PartitionOp(PartitionOpKind::Merge, tgt1, tgt2);
+  static PartitionOp Require(unsigned tgt,
+                             SILInstruction *sourceInst = nullptr) {
+    return PartitionOp(PartitionOpKind::Require, tgt, sourceInst);
   }
 
-  static PartitionOp Require(unsigned tgt) {
-    return PartitionOp(PartitionOpKind::Require, tgt);
+  SILInstruction *getSourceInst() const {
+    return sourceInst;
+  }
+
+  void dump() const {
+    switch (OpKind) {
+    case PartitionOpKind::Assign:
+      llvm::dbgs() << "assign %" << OpArgs[0] << " = %" << OpArgs[1] << "\n";
+      break;
+    case PartitionOpKind::AssignFresh:
+      llvm::dbgs() << "assign_fresh %" << OpArgs[0] << "\n";
+      break;
+    case PartitionOpKind::Consume:
+      llvm::dbgs() << "consume %" << OpArgs[0] << "\n";
+      break;
+    case PartitionOpKind::Merge:
+      llvm::dbgs() << "merge %" << OpArgs[0] << " with %" << OpArgs[1] << "\n";
+      break;
+    case PartitionOpKind::Require:
+      llvm::dbgs() << "require %" << OpArgs[0] << "\n";
+      break;
+    }
   }
 };
 
@@ -72,22 +118,33 @@ static void horizontalUpdate(std::map<unsigned, signed> &map, unsigned key,
 
 class Partition {
 private:
-  std::map<unsigned, signed> labels = {};
+  // 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
+  std::map<unsigned, signed> labels;
+
+  // track a label that is guaranteed to be fresh
+  unsigned fresh_label = 0;
 
-  bool canonical = true;
+  // 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()
+  bool canonical;
 
   // 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 missing region.
+  // This excludes the -1 label for consumed regions.
   void canonicalize() {
     if (canonical)
       return;
     canonical = true;
 
     std::map<signed, unsigned> relabel;
 
+    // relies on in-order traversal of labels
     for (auto &[i, label] : labels) {
-      // leave -1 (missing region) as is
+      // leave -1 (consumed region) as is
       if (label < 0)
         continue;
 
@@ -98,11 +155,31 @@ class Partition {
         relabel[label] = i;
       }
 
+      // update this label with either its own index, or a prior index that
+      // shared a region with it
       label = relabel[label];
+
+      // the maximum index iterated over will be used here to appropriately
+      // set fresh_label
+      fresh_label = i + 1;
     }
   }
 
 public:
+  Partition() : labels({}), canonical(true) {}
+
+  static Partition singleRegion(std::vector<unsigned> indices) {
+    Partition p;
+    if (!indices.empty()) {
+      unsigned min_index = *std::min_element(indices.begin(), indices.end());
+      p.fresh_label = min_index + 1;
+      for (unsigned index : indices) {
+        p.labels[index] = min_index;
+      }
+    }
+    return p;
+  }
+
   void dump() const {
     llvm::dbgs() << "Partition";
     if (canonical)
@@ -160,42 +237,35 @@ class Partition {
     for (const auto &[i, _] : fst.labels) {
       if (!snd.labels.count(i))
         continue;
-      joined.labels[i] = lookup_fst(i);
+      signed label_i = lookup_fst(i);
+      joined.labels[i] = label_i;
+      joined.fresh_label = std::max(joined.fresh_label, (unsigned) label_i + 1);
     }
 
     return joined;
   }
 
-  // It's possible for all PartitionOps' to maintain canonicality,
-  // but it comes at the cost of making Assign operations worst-case
-  // linear time instead of constant. This is likely not worth it,
-  // so it's disabled by default, but leaving this flag here in case
-  // it becomes useful as a performance optimization.
-  static const bool ALWAYS_CANONICAL = false;
-
+  // 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.
   void apply(
-      PartitionOp op, std::function<void(const PartitionOp &)> handleFailure =
-                          [](const PartitionOp &_) {}) {
+      PartitionOp op, llvm::function_ref<void(const PartitionOp&, unsigned)> handleFailure =
+                          [](const PartitionOp&, unsigned) {}) {
     switch (op.OpKind) {
     case PartitionOpKind::Assign:
       assert(op.OpArgs.size() == 2 &&
              "Assign PartitionOp should be passed 2 arguments");
-      assert(labels.count(op.OpArgs[0]) && labels.count(op.OpArgs[1]) &&
-             "Assign PartitionOp's arguments should be already tracked");
+      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)
-        handleFailure(op);
+        handleFailure(op, op.OpArgs[1]);
 
       labels[op.OpArgs[0]] = labels[op.OpArgs[1]];
 
-      if (ALWAYS_CANONICAL) {
-        // if seeking to maintain canonicality, then do so
-        if (op.OpArgs[0] < labels[op.OpArgs[0]])
-          horizontalUpdate(labels, op.OpArgs[0], op.OpArgs[0]);
-        break;
-      }
-
-      // assignment could have invalidated canonicality
+      // assignment could have invalidated canonicality of either the old region
+      // of op.OpArgs[0] or the region of op.OpArgs[1], or both
       canonical = false;
       break;
     case PartitionOpKind::AssignFresh:
@@ -204,30 +274,33 @@ class Partition {
       assert(!labels.count(op.OpArgs[0]) &&
              "AssignFresh PartitionOp's argument should NOT already be tracked");
 
-      // fresh region generated by mapping the passed index to itself
-      labels[op.OpArgs[0]] = op.OpArgs[0];
+      // map index op.OpArgs[0] to a fresh label
+      labels[op.OpArgs[0]] = fresh_label++;
+      canonical = false;
       break;
     case PartitionOpKind::Consume:
       assert(op.OpArgs.size() == 1 &&
              "Consume PartitionOp should be passed 1 argument");
       assert(labels.count(op.OpArgs[0]) &&
              "Consume PartitionOp's argument should already be tracked");
 
-      // if attempting to consume a missing region, handle the failure
+      // if attempting to consume a consumed region, handle the failure
       if (labels[op.OpArgs[0]] < 0)
-        handleFailure(op);
+        handleFailure(op, op.OpArgs[0]);
 
-      // mark region as missing
+      // mark region as consumed
       horizontalUpdate(labels, op.OpArgs[0], -1);
       break;
     case PartitionOpKind::Merge:
       assert(op.OpArgs.size() == 2 &&
              "Merge PartitionOp should be passed 2 arguments");
       assert(labels.count(op.OpArgs[0]) && labels.count(op.OpArgs[1]) &&
              "Merge PartitionOp's arguments should already be tracked");
-      // if attempting to merge a missing region, handle the failure
-      if (labels[op.OpArgs[0]] < 0 || labels[op.OpArgs[1]] < 0)
-        handleFailure(op);
+      // if attempting to merge a consumed region, handle the failure
+      if (labels[op.OpArgs[0]] < 0)
+          handleFailure(op, op.OpArgs[0]);
+      if (labels[op.OpArgs[1]] < 0)
+          handleFailure(op, op.OpArgs[1]);
 
       if (labels[op.OpArgs[0]] == labels[op.OpArgs[1]])
         break;
@@ -244,8 +317,26 @@ class Partition {
       assert(labels.count(op.OpArgs[0]) &&
              "Require PartitionOp's argument should already be tracked");
       if (labels[op.OpArgs[0]] < 0)
-        handleFailure(op);
+        handleFailure(op, op.OpArgs[0]);
+    }
+  }
+
+  void dump() {
+    std::map<signed, std::vector<unsigned>> buckets;
+
+    for (auto [i, label] : labels) {
+      buckets[label].push_back(i);
+    }
+
+    llvm::dbgs() << "[";
+    for (auto [label, indices] : buckets) {
+      llvm::dbgs() << (label < 0 ? "{" : "(");
+      for (unsigned i : indices) {
+        llvm::dbgs() << i << " ";
+      }
+      llvm::dbgs() << (label < 0 ? "}" : ")");
     }
+    llvm::dbgs() << "]\n";
   }
 };
 }