[region-isolation] Implement the dataflow correctly.

This involved fixing a few different bugs.

1. We were just performing dataflow by setting that only the initial block needs
to be updated. This means that if there isn’t anything in the initial dataflow
block, we won’t visit any successor blocks. Instead the correct thing to do here
is to visit all blocks in the initial round.

2. I also needed to fix a separate issue where we were updating our union-find
data structure incorrectly as found by an assert on transfernonsendable.swift
that was triggered once I fixed 1. Put simply, we needed to set a new max label
+ 1 when our new max element is less than or equal to the old max label + 1…
before we just did less than so if we had a new max element that is the same as
our fresh label, we wouldn’t increment the fresh label.

rdar://119584497

Author: gottesmm

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -577,7 +577,9 @@ class Partition {
             fstReduced.elementToRegionMap.find(Element(sndRegionNumber));
         if (iter != fstReduced.elementToRegionMap.end()) {
           fstReduced.elementToRegionMap.insert({sndEltNumber, iter->second});
-          if (fstReduced.fresh_label < Region(sndEltNumber))
+          // We want fresh_label to always be one element larger than our
+          // maximum element.
+          if (fstReduced.fresh_label <= Region(sndEltNumber))
             fstReduced.fresh_label = Region(sndEltNumber + 1);
           continue;
         }
@@ -597,15 +599,10 @@ class Partition {
           fstIter.first->getSecond() =
               fstIter.first->second->merge(sndIter->second);
       }
-      if (fstReduced.fresh_label < sndRegionNumber)
+      if (fstReduced.fresh_label <= sndRegionNumber)
         fstReduced.fresh_label = Region(sndEltNumber + 1);
     }
 
-    LLVM_DEBUG(llvm::dbgs() << "JOIN PEFORMED: \nFST: ";
-               fst.print(llvm::dbgs()); llvm::dbgs() << "SND: ";
-               snd.print(llvm::dbgs()); llvm::dbgs() << "RESULT: ";
-               fstReduced.print(llvm::dbgs()););
-
     assert(fstReduced.is_canonical_correct());
 
     // fst_reduced is now the join

lib/SILOptimizer/Mandatory/TransferNonSendable.cpp

@@ -2639,9 +2639,6 @@ class BlockPartitionState {
   /// Set if this block in the next iteration needs to be visited.
   bool needsUpdate = false;
 
-  /// Set if we have ever visited this block at all.
-  bool reached = false;
-
   /// The partition of elements into regions at the top of the block.
   Partition entryPartition;
 
@@ -2698,8 +2695,12 @@ class BlockPartitionState {
       // will be suppressed
       eval.apply(partitionOp);
     }
+    LLVM_DEBUG(llvm::dbgs() << "    Working Partition: ";
+               workingPartition.print(llvm::dbgs()));
     bool exitUpdated = !Partition::equals(exitPartition, workingPartition);
     exitPartition = workingPartition;
+    LLVM_DEBUG(llvm::dbgs() << "    Exit Partition: ";
+               exitPartition.print(llvm::dbgs()));
     return exitUpdated;
   }
 
@@ -2713,8 +2714,8 @@ class BlockPartitionState {
   SWIFT_DEBUG_DUMP { print(llvm::dbgs()); }
 
   void print(llvm::raw_ostream &os) const {
-    os << SEP_STR << "BlockPartitionState[reached=" << reached
-       << ", needsUpdate=" << needsUpdate << "]\nid: ";
+    os << SEP_STR << "BlockPartitionState[needsUpdate=" << needsUpdate
+       << "]\nid: ";
 #ifndef NDEBUG
     auto printID = [&](SILBasicBlock *block) {
       block->printID(os);
@@ -2974,9 +2975,11 @@ class PartitionAnalysis {
                     }),
         allocator(), ptrSetFactory(allocator), function(fn), pofi(pofi),
         solved(false) {
-    // Initialize the entry block as needing an update, and having a partition
-    // that places all its non-sendable args in a single region
-    blockStates[fn->getEntryBlock()].needsUpdate = true;
+    // Mark all blocks as needing to be updated.
+    for (auto &block : *fn) {
+      blockStates[&block].needsUpdate = true;
+    }
+    // Set our entry partition to have the "entry partition".
     blockStates[fn->getEntryBlock()].entryPartition =
         translator.getEntryPartition();
   }
@@ -3001,57 +3004,32 @@ class PartitionAnalysis {
           continue;
         }
 
-        // mark this block as no longer needing an update
+        // Mark this block as no longer needing an update.
         blockState.needsUpdate = false;
 
-        // mark this block as reached by the analysis
-        blockState.reached = true;
-
-        // compute the new entry partition to this block
-        Partition newEntryPartition;
-        bool firstPred = true;
+        // Compute the new entry partition to this block.
+        Partition newEntryPartition = blockState.entryPartition;
 
         LLVM_DEBUG(llvm::dbgs() << "    Visiting Preds!\n");
 
-        // This loop computes the join of the exit partitions of all
+        // This loop computes the union of the exit partitions of all
         // predecessors of this block
         for (SILBasicBlock *predBlock : block->getPredecessorBlocks()) {
           BlockPartitionState &predState = blockStates[predBlock];
-          // ignore predecessors that haven't been reached by the analysis yet
-          if (!predState.reached)
-            continue;
-
-          if (firstPred) {
-            firstPred = false;
-            newEntryPartition = predState.exitPartition;
-            LLVM_DEBUG(llvm::dbgs() << "    First Pred. bb"
-                                    << predBlock->getDebugID() << ": ";
-                       newEntryPartition.print(llvm::dbgs()));
-            continue;
-          }
 
+          // Predecessors that have not been reached yet will have an empty pred
+          // state... so just merge them all. Also our initial value of
           LLVM_DEBUG(llvm::dbgs()
                          << "    Pred. bb" << predBlock->getDebugID() << ": ";
                      predState.exitPartition.print(llvm::dbgs()));
           newEntryPartition =
               Partition::join(newEntryPartition, predState.exitPartition);
         }
 
-        // If we found predecessor blocks, then attempt to use them to update
-        // the entry partition for this block, and abort this block's update if
-        // the entry partition was not updated.
-        if (!firstPred) {
-          // if the recomputed entry partition is the same as the current one,
-          // perform no update
-          if (Partition::equals(newEntryPartition, blockState.entryPartition)) {
-            LLVM_DEBUG(llvm::dbgs()
-                       << "    Entry partition is the same... skipping!\n");
-            continue;
-          }
-
-          // otherwise update the entry partition
-          blockState.entryPartition = newEntryPartition;
-        }
+        // Update the entry partition. We need to still try to
+        // recomputeExitFromEntry before we know if we made a difference to the
+        // exit partition after applying the instructions of the block.
+        blockState.entryPartition = newEntryPartition;
 
         // recompute this block's exit partition from its (updated) entry
         // partition, and if this changed the exit partition notify all
@@ -3169,6 +3147,8 @@ class PartitionAnalysis {
     LLVM_DEBUG(llvm::dbgs() << "Walking blocks for diagnostics.\n");
     for (auto [block, blockState] : blockStates) {
       LLVM_DEBUG(llvm::dbgs() << "|--> Block bb" << block.getDebugID() << "\n");
+      LLVM_DEBUG(llvm::dbgs() << "Entry Partition: ";
+                 blockState.getEntryPartition().print(llvm::dbgs()));
 
       // Grab its entry partition and setup an evaluator for the partition that
       // has callbacks that emit diagnsotics...
@@ -3180,6 +3160,9 @@ class PartitionAnalysis {
       for (auto &partitionOp : blockState.getPartitionOps()) {
         eval.apply(partitionOp);
       }
+
+      LLVM_DEBUG(llvm::dbgs() << "Exit Partition: ";
+                 workingPartition.print(llvm::dbgs()));
     }
 
     // Now that we have found all of our transferInsts/Requires emit errors.

test/Concurrency/Inputs/transfernonsendable_crashers_objc.h

@@ -0,0 +1,5 @@
+
+@interface MyNotificationCenter
+- (id)init;
+- (void)post;
+@end

test/Concurrency/transfernonsendable_cfg.sil

@@ -0,0 +1,54 @@
+// RUN: %target-sil-opt -transfer-non-sendable -enable-experimental-feature RegionBasedIsolation -strict-concurrency=complete %s -verify
+
+// PLEASE READ THIS!
+//
+// In this file, we test out how the checker handles dataflow in a CFG. We want
+// to validate that we properly propagate dataflow given instructions occur in
+// specific locations.
+
+sil_stage raw
+
+import Swift
+
+////////////////////////
+// MARK: Declarations //
+////////////////////////
+
+class NonSendableKlass {}
+
+sil @transferKlass : $@convention(thin) @async (@guaranteed NonSendableKlass) -> ()
+sil @useKlass : $@convention(thin) (@guaranteed NonSendableKlass) -> ()
+sil @constructKlass : $@convention(thin) () -> @owned NonSendableKlass
+
+/////////////////
+// MARK: Tests //
+/////////////////
+
+// The problem this validates that we properly propagate dataflow through blocks
+// if the initial block does not have any partition state.
+sil [ossa] @haveIntroducerInMiddleOfCFG : $@convention(thin) @async () -> () {
+bb0:
+  br bb0a
+
+bb0a:
+  %constructKlass = function_ref @constructKlass : $@convention(thin) () -> @owned NonSendableKlass
+  %klass = apply %constructKlass() : $@convention(thin) () -> @owned NonSendableKlass
+  cond_br undef, bb1, bb2
+
+bb1:
+  br bb3
+
+bb2:
+  br bb3
+
+bb3:
+  %transferKlass = function_ref @transferKlass : $@convention(thin) @async (@guaranteed NonSendableKlass) -> ()
+  apply [caller_isolation=nonisolated] [callee_isolation=global_actor] %transferKlass(%klass) : $@convention(thin) @async (@guaranteed NonSendableKlass) -> ()
+  // expected-warning @-1 {{passing argument of non-sendable type 'NonSendableKlass' from nonisolated context to global actor '<null>'-isolated context at this call site could yield a race with accesses later in this function}}
+  %useKlass = function_ref @useKlass : $@convention(thin) (@guaranteed NonSendableKlass) -> ()
+  apply %useKlass(%klass) : $@convention(thin) (@guaranteed NonSendableKlass) -> ()
+  // expected-note @-1 {{access here could race}}
+  destroy_value %klass : $NonSendableKlass
+  %9999 = tuple ()
+  return %9999 : $()
+}

test/Concurrency/transfernonsendable_crashers_objc.swift

@@ -0,0 +1,19 @@
+// RUN: %target-swift-frontend -emit-sil -strict-concurrency=complete -enable-experimental-feature RegionBasedIsolation -disable-availability-checking %s -o /dev/null -import-objc-header %S/Inputs/transfernonsendable_crashers_objc.h
+
+// REQUIRES: objc_interop
+
+import Foundation
+
+extension MyNotificationCenter {
+  static var value = MyNotificationCenter()
+}
+
+public func handleFile(at location: URL) throws {
+//  createDownloadDirectoryIfRequired()
+  let movedFileLocation = try moveTheme(from: location)
+  let unzippedFileLocation = try unzipFile(at: movedFileLocation)
+  MyNotificationCenter.value!.post()
+}
+
+private func moveTheme(from location: URL) throws -> URL { fatalError() }
+private func unzipFile(at location: URL) throws -> URL { fatalError() }