diff --git a/include/swift/SILOptimizer/Utils/PartitionUtils.h b/include/swift/SILOptimizer/Utils/PartitionUtils.h
index bd399e12c6f5b..915ee952967cc 100644
--- a/include/swift/SILOptimizer/Utils/PartitionUtils.h
+++ b/include/swift/SILOptimizer/Utils/PartitionUtils.h
@@ -225,6 +225,11 @@ static void horizontalUpdate(std::map<Element, Region> &map, Element key,
 }
 
 class Partition {
+public:
+  /// A class defined in PartitionUtils unittest used to grab state from
+  /// Partition without exposing it to other users.
+  struct PartitionTester;
+
 private:
   // 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 transferred
@@ -351,6 +356,21 @@ class Partition {
     return p;
   }
 
+  static Partition separateRegions(ArrayRef<Element> indices) {
+    Partition p;
+    if (indices.empty())
+      return p;
+
+    Region max_index =
+        Region(*std::max_element(indices.begin(), indices.end()));
+    p.fresh_label = Region(max_index + 1);
+    for (Element index : indices) {
+      p.labels.insert_or_assign(index, Region(index));
+    }
+    assert(p.is_canonical_correct());
+    return p;
+  }
+
   // linear time - Test two partititons for equality by first putting them
   // in canonical form then comparing for exact equality.
   static bool equals(Partition &fst, Partition &snd) {
@@ -370,30 +390,63 @@ class Partition {
   // 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
   // region in either operand.
-  static Partition join(Partition &fst, Partition &snd) {
-    // 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.insert_or_assign(i, fst.labels.at(i));
-        snd_reduced.labels.insert_or_assign(i, snd.labels.at(i));
-      }
+  static Partition join(const Partition &fst, const Partition &snd) {
+    // First copy and canonicalize our inputs.
+    Partition fst_reduced = fst;
+    Partition snd_reduced = snd;
+
     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.isTransferred())
-        // if snd says that the region has been transferred, mark it transferred
-        // in fst
-        horizontalUpdate(fst_reduced.labels, i, Region::transferred());
-      else
-        fst_reduced.merge(i, Element(snd_label));
+    // For each element in snd_reduced...
+    for (const auto &[sndEltNumber, sndRegionNumber] : snd_reduced.labels) {
+      // For values that are both in fst_reduced and snd_reduced, we need to
+      // merge their regions.
+      if (fst_reduced.labels.count(sndEltNumber)) {
+        if (sndRegionNumber.isTransferred()) {
+          // If snd says that the region has been transferred, mark it
+          // transferred in fst.
+          horizontalUpdate(fst_reduced.labels, sndEltNumber,
+                           Region::transferred());
+          continue;
+        }
+
+        // Otherwise merge. This maintains canonicality.
+        fst_reduced.merge(sndEltNumber, Element(sndRegionNumber));
+        continue;
+      }
+
+      // Otherwise, we have an element in snd that is not in fst. First see if
+      // our region is transferred. In such a case, just add this element as
+      // transferred.
+      if (sndRegionNumber.isTransferred()) {
+        fst_reduced.labels.insert({sndEltNumber, Region::transferred()});
+        continue;
+      }
+
+      // Then check if the representative element number for this element in snd
+      // is in fst. In that case, we know that we visited it before we visited
+      // this elt number (since we are processing in order) so what ever is
+      // mapped to that number in snd must be the correct number for this
+      // element as well since this number is guaranteed to be greater than our
+      // representative and the number mapped to our representative in fst must
+      // be <= our representative.
+      auto iter = fst_reduced.labels.find(Element(sndRegionNumber));
+      if (iter != fst_reduced.labels.end()) {
+        fst_reduced.labels.insert({sndEltNumber, iter->second});
+        if (fst_reduced.fresh_label < Region(sndEltNumber))
+          fst_reduced.fresh_label = Region(sndEltNumber + 1);
+        continue;
+      }
+
+      // Otherwise, we have an element that is not in fst and its representative
+      // is not in fst. This means that we must be our representative in snd
+      // since we should have visited our representative earlier if we were not
+      // due to our traversal being in order. Thus just add this to fst_reduced.
+      assert(sndEltNumber == Element(sndRegionNumber));
+      fst_reduced.labels.insert({sndEltNumber, sndRegionNumber});
+      if (fst_reduced.fresh_label < sndRegionNumber)
+        fst_reduced.fresh_label = Region(sndEltNumber + 1);
     }
 
     LLVM_DEBUG(llvm::dbgs() << "JOIN PEFORMED: \nFST: ";
@@ -581,6 +634,9 @@ class Partition {
     }
     os << "]\n";
   }
+
+  /// Routine used in unittests
+  Region getRegion(Element elt) const { return labels.at(elt); }
 };
 
 } // namespace swift
diff --git a/test/Concurrency/sendnonsendable_basic.swift b/test/Concurrency/sendnonsendable_basic.swift
index 277797ffe7bf1..c00adcdda4d24 100644
--- a/test/Concurrency/sendnonsendable_basic.swift
+++ b/test/Concurrency/sendnonsendable_basic.swift
@@ -263,10 +263,10 @@ extension Actor {
       }
     }
 
-    // We emit the first error that we see. So we do not emit the self error.
     await transferToMain(closure) // expected-sns-note {{access here could race}}
-    // expected-complete-warning @-1 {{passing argument of non-sendable type '() -> ()' into main actor-isolated context may introduce data races}}
-    // expected-complete-note @-2 {{a function type must be marked '@Sendable' to conform to 'Sendable'}}
+    // expected-sns-warning @-1 {{call site passes `self` or a non-sendable argument of this function to another thread, potentially yielding a race with the caller}}
+    // expected-complete-warning @-2 {{passing argument of non-sendable type '() -> ()' into main actor-isolated context may introduce data races}}
+    // expected-complete-note @-3 {{a function type must be marked '@Sendable' to conform to 'Sendable'}}
   }
 
   // In this case, we reinit along both paths, but only one has an actor derived
@@ -286,9 +286,7 @@ extension Actor {
       closure = {}
     }
 
-    // TODO: We do not error here yet since we are performing the control flow
-    // union incorrectly.
-    await transferToMain(closure)
+    await transferToMain(closure) // expected-sns-warning {{call site passes `self` or a non-sendable argument of this function to another thread, potentially yielding a race with the caller}}
     // expected-complete-warning @-1 {{passing argument of non-sendable type '() -> ()' into main actor-isolated context may introduce data races}}
     // expected-complete-note @-2 {{a function type must be marked '@Sendable' to conform to 'Sendable'}}
   }
diff --git a/unittests/CMakeLists.txt b/unittests/CMakeLists.txt
index f50837589506a..9342efbbc9ed3 100644
--- a/unittests/CMakeLists.txt
+++ b/unittests/CMakeLists.txt
@@ -16,6 +16,7 @@ if(SWIFT_INCLUDE_TOOLS)
   add_subdirectory(Remangler)
   add_subdirectory(Sema)
   add_subdirectory(SIL)
+  add_subdirectory(SILOptimizer)
   add_subdirectory(SwiftDemangle)
 
   add_subdirectory(Threading)
diff --git a/unittests/SILOptimizer/CMakeLists.txt b/unittests/SILOptimizer/CMakeLists.txt
index f4852e80c7cb6..616236c62d726 100644
--- a/unittests/SILOptimizer/CMakeLists.txt
+++ b/unittests/SILOptimizer/CMakeLists.txt
@@ -5,4 +5,7 @@ add_swift_unittest(SwiftSILOptimizerTests
 target_link_libraries(SwiftSILOptimizerTests
    PRIVATE
    swiftSILOptimizer
+   swiftIRGen
+   swiftAST
+   swiftFrontend
 )
diff --git a/unittests/SILOptimizer/PartitionUtilsTest.cpp b/unittests/SILOptimizer/PartitionUtilsTest.cpp
index 47a15dde94554..3ecf14c0f3a10 100644
--- a/unittests/SILOptimizer/PartitionUtilsTest.cpp
+++ b/unittests/SILOptimizer/PartitionUtilsTest.cpp
@@ -14,7 +14,32 @@
 
 #include "gtest/gtest.h"
 
+#include <array>
+
 using namespace swift;
+using namespace swift::PartitionPrimitives;
+
+//===----------------------------------------------------------------------===//
+//                                 Utilities
+//===----------------------------------------------------------------------===//
+
+struct Partition::PartitionTester {
+  const Partition &p;
+
+  PartitionTester(const Partition &p) : p(p) {}
+
+  signed getRegion(unsigned elt) const {
+    return signed(p.labels.at(Element(elt)));
+  }
+};
+
+namespace {
+using PartitionTester = Partition::PartitionTester;
+} // namespace
+
+//===----------------------------------------------------------------------===//
+//                                   Tests
+//===----------------------------------------------------------------------===//
 
 // this test tests that if a series of merges is split between two partitions
 // p1 and p2, but also applied in its entirety to p3, then joining p1 and p2
@@ -100,6 +125,235 @@ TEST(PartitionUtilsTest, TestMergeAndJoin) {
   apply_to_p1_and_p3(PartitionOp::Merge(Element(7), Element(8)));
 }
 
+TEST(PartitionUtilsTest, Join1) {
+  Element data1[] = {Element(0), Element(1), Element(2),
+                     Element(3), Element(4), Element(5)};
+  Partition p1 = Partition::separateRegions(llvm::makeArrayRef(data1));
+
+  p1.apply(PartitionOp::Assign(Element(0), Element(0)));
+  p1.apply(PartitionOp::Assign(Element(1), Element(0)));
+  p1.apply(PartitionOp::Assign(Element(2), Element(2)));
+  p1.apply(PartitionOp::Assign(Element(3), Element(3)));
+  p1.apply(PartitionOp::Assign(Element(4), Element(3)));
+  p1.apply(PartitionOp::Assign(Element(5), Element(2)));
+
+  Partition p2 = Partition::separateRegions(llvm::makeArrayRef(data1));
+  p2.apply(PartitionOp::Assign(Element(0), Element(0)));
+  p2.apply(PartitionOp::Assign(Element(1), Element(0)));
+  p2.apply(PartitionOp::Assign(Element(2), Element(2)));
+  p2.apply(PartitionOp::Assign(Element(3), Element(3)));
+  p2.apply(PartitionOp::Assign(Element(4), Element(3)));
+  p2.apply(PartitionOp::Assign(Element(5), Element(5)));
+
+  auto result = Partition::join(p1, p2);
+  PartitionTester tester(result);
+  ASSERT_EQ(tester.getRegion(0), 0);
+  ASSERT_EQ(tester.getRegion(1), 0);
+  ASSERT_EQ(tester.getRegion(2), 2);
+  ASSERT_EQ(tester.getRegion(3), 3);
+  ASSERT_EQ(tester.getRegion(4), 3);
+  ASSERT_EQ(tester.getRegion(5), 2);
+}
+
+TEST(PartitionUtilsTest, Join2) {
+  Element data1[] = {Element(0), Element(1), Element(2),
+                     Element(3), Element(4), Element(5)};
+  Partition p1 = Partition::separateRegions(llvm::makeArrayRef(data1));
+
+  p1.apply(PartitionOp::Assign(Element(0), Element(0)));
+  p1.apply(PartitionOp::Assign(Element(1), Element(0)));
+  p1.apply(PartitionOp::Assign(Element(2), Element(2)));
+  p1.apply(PartitionOp::Assign(Element(3), Element(3)));
+  p1.apply(PartitionOp::Assign(Element(4), Element(3)));
+  p1.apply(PartitionOp::Assign(Element(5), Element(2)));
+
+  Element data2[] = {Element(4), Element(5), Element(6),
+                     Element(7), Element(8), Element(9)};
+  Partition p2 = Partition::separateRegions(llvm::makeArrayRef(data2));
+  p2.apply(PartitionOp::Assign(Element(4), Element(4)));
+  p2.apply(PartitionOp::Assign(Element(5), Element(5)));
+  p2.apply(PartitionOp::Assign(Element(6), Element(4)));
+  p2.apply(PartitionOp::Assign(Element(7), Element(7)));
+  p2.apply(PartitionOp::Assign(Element(8), Element(7)));
+  p2.apply(PartitionOp::Assign(Element(9), Element(4)));
+
+  auto result = Partition::join(p1, p2);
+  PartitionTester tester(result);
+  ASSERT_EQ(tester.getRegion(0), 0);
+  ASSERT_EQ(tester.getRegion(1), 0);
+  ASSERT_EQ(tester.getRegion(2), 2);
+  ASSERT_EQ(tester.getRegion(3), 3);
+  ASSERT_EQ(tester.getRegion(4), 3);
+  ASSERT_EQ(tester.getRegion(5), 2);
+  ASSERT_EQ(tester.getRegion(6), 3);
+  ASSERT_EQ(tester.getRegion(7), 7);
+  ASSERT_EQ(tester.getRegion(8), 7);
+  ASSERT_EQ(tester.getRegion(9), 3);
+}
+
+TEST(PartitionUtilsTest, Join2Reversed) {
+  Element data1[] = {Element(0), Element(1), Element(2),
+                     Element(3), Element(4), Element(5)};
+  Partition p1 = Partition::separateRegions(llvm::makeArrayRef(data1));
+
+  p1.apply(PartitionOp::Assign(Element(0), Element(0)));
+  p1.apply(PartitionOp::Assign(Element(1), Element(0)));
+  p1.apply(PartitionOp::Assign(Element(2), Element(2)));
+  p1.apply(PartitionOp::Assign(Element(3), Element(3)));
+  p1.apply(PartitionOp::Assign(Element(4), Element(3)));
+  p1.apply(PartitionOp::Assign(Element(5), Element(2)));
+
+  Element data2[] = {Element(4), Element(5), Element(6),
+                     Element(7), Element(8), Element(9)};
+  Partition p2 = Partition::separateRegions(llvm::makeArrayRef(data2));
+  p2.apply(PartitionOp::Assign(Element(4), Element(4)));
+  p2.apply(PartitionOp::Assign(Element(5), Element(5)));
+  p2.apply(PartitionOp::Assign(Element(6), Element(4)));
+  p2.apply(PartitionOp::Assign(Element(7), Element(7)));
+  p2.apply(PartitionOp::Assign(Element(8), Element(7)));
+  p2.apply(PartitionOp::Assign(Element(9), Element(4)));
+
+  auto result = Partition::join(p2, p1);
+  PartitionTester tester(result);
+  ASSERT_EQ(tester.getRegion(0), 0);
+  ASSERT_EQ(tester.getRegion(1), 0);
+  ASSERT_EQ(tester.getRegion(2), 2);
+  ASSERT_EQ(tester.getRegion(3), 3);
+  ASSERT_EQ(tester.getRegion(4), 3);
+  ASSERT_EQ(tester.getRegion(5), 2);
+  ASSERT_EQ(tester.getRegion(6), 3);
+  ASSERT_EQ(tester.getRegion(7), 7);
+  ASSERT_EQ(tester.getRegion(8), 7);
+  ASSERT_EQ(tester.getRegion(9), 3);
+}
+
+TEST(PartitionUtilsTest, JoinLarge) {
+  Element data1[] = {
+      Element(0),  Element(1),  Element(2),  Element(3),  Element(4),
+      Element(5),  Element(6),  Element(7),  Element(8),  Element(9),
+      Element(10), Element(11), Element(12), Element(13), Element(14),
+      Element(15), Element(16), Element(17), Element(18), Element(19),
+      Element(20), Element(21), Element(22), Element(23), Element(24),
+      Element(25), Element(26), Element(27), Element(28), Element(29)};
+  Partition p1 = Partition::separateRegions(llvm::makeArrayRef(data1));
+  p1.apply(PartitionOp::Assign(Element(0), Element(29)));
+  p1.apply(PartitionOp::Assign(Element(1), Element(17)));
+  p1.apply(PartitionOp::Assign(Element(2), Element(0)));
+  p1.apply(PartitionOp::Assign(Element(3), Element(12)));
+  p1.apply(PartitionOp::Assign(Element(4), Element(13)));
+  p1.apply(PartitionOp::Assign(Element(5), Element(9)));
+  p1.apply(PartitionOp::Assign(Element(6), Element(15)));
+  p1.apply(PartitionOp::Assign(Element(7), Element(27)));
+  p1.apply(PartitionOp::Assign(Element(8), Element(3)));
+  p1.apply(PartitionOp::Assign(Element(9), Element(3)));
+  p1.apply(PartitionOp::Assign(Element(10), Element(3)));
+  p1.apply(PartitionOp::Assign(Element(11), Element(21)));
+  p1.apply(PartitionOp::Assign(Element(12), Element(14)));
+  p1.apply(PartitionOp::Assign(Element(13), Element(25)));
+  p1.apply(PartitionOp::Assign(Element(14), Element(1)));
+  p1.apply(PartitionOp::Assign(Element(15), Element(25)));
+  p1.apply(PartitionOp::Assign(Element(16), Element(12)));
+  p1.apply(PartitionOp::Assign(Element(17), Element(3)));
+  p1.apply(PartitionOp::Assign(Element(18), Element(25)));
+  p1.apply(PartitionOp::Assign(Element(19), Element(13)));
+  p1.apply(PartitionOp::Assign(Element(20), Element(19)));
+  p1.apply(PartitionOp::Assign(Element(21), Element(7)));
+  p1.apply(PartitionOp::Assign(Element(22), Element(19)));
+  p1.apply(PartitionOp::Assign(Element(23), Element(27)));
+  p1.apply(PartitionOp::Assign(Element(24), Element(1)));
+  p1.apply(PartitionOp::Assign(Element(25), Element(9)));
+  p1.apply(PartitionOp::Assign(Element(26), Element(18)));
+  p1.apply(PartitionOp::Assign(Element(27), Element(29)));
+  p1.apply(PartitionOp::Assign(Element(28), Element(28)));
+  p1.apply(PartitionOp::Assign(Element(29), Element(13)));
+
+  Element data2[] = {
+      Element(15), Element(16), Element(17), Element(18), Element(19),
+      Element(20), Element(21), Element(22), Element(23), Element(24),
+      Element(25), Element(26), Element(27), Element(28), Element(29),
+      Element(30), Element(31), Element(32), Element(33), Element(34),
+      Element(35), Element(36), Element(37), Element(38), Element(39),
+      Element(40), Element(41), Element(42), Element(43), Element(44)};
+  Partition p2 = Partition::separateRegions(llvm::makeArrayRef(data2));
+  p2.apply(PartitionOp::Assign(Element(15), Element(31)));
+  p2.apply(PartitionOp::Assign(Element(16), Element(34)));
+  p2.apply(PartitionOp::Assign(Element(17), Element(35)));
+  p2.apply(PartitionOp::Assign(Element(18), Element(41)));
+  p2.apply(PartitionOp::Assign(Element(19), Element(15)));
+  p2.apply(PartitionOp::Assign(Element(20), Element(32)));
+  p2.apply(PartitionOp::Assign(Element(21), Element(17)));
+  p2.apply(PartitionOp::Assign(Element(22), Element(31)));
+  p2.apply(PartitionOp::Assign(Element(23), Element(21)));
+  p2.apply(PartitionOp::Assign(Element(24), Element(33)));
+  p2.apply(PartitionOp::Assign(Element(25), Element(25)));
+  p2.apply(PartitionOp::Assign(Element(26), Element(31)));
+  p2.apply(PartitionOp::Assign(Element(27), Element(16)));
+  p2.apply(PartitionOp::Assign(Element(28), Element(35)));
+  p2.apply(PartitionOp::Assign(Element(29), Element(40)));
+  p2.apply(PartitionOp::Assign(Element(30), Element(33)));
+  p2.apply(PartitionOp::Assign(Element(31), Element(34)));
+  p2.apply(PartitionOp::Assign(Element(32), Element(22)));
+  p2.apply(PartitionOp::Assign(Element(33), Element(42)));
+  p2.apply(PartitionOp::Assign(Element(34), Element(37)));
+  p2.apply(PartitionOp::Assign(Element(35), Element(34)));
+  p2.apply(PartitionOp::Assign(Element(36), Element(18)));
+  p2.apply(PartitionOp::Assign(Element(37), Element(32)));
+  p2.apply(PartitionOp::Assign(Element(38), Element(22)));
+  p2.apply(PartitionOp::Assign(Element(39), Element(44)));
+  p2.apply(PartitionOp::Assign(Element(40), Element(20)));
+  p2.apply(PartitionOp::Assign(Element(41), Element(37)));
+  p2.apply(PartitionOp::Assign(Element(43), Element(29)));
+  p2.apply(PartitionOp::Assign(Element(44), Element(25)));
+
+  auto result = Partition::join(p1, p2);
+  PartitionTester tester(result);
+  ASSERT_EQ(tester.getRegion(0), 0);
+  ASSERT_EQ(tester.getRegion(1), 1);
+  ASSERT_EQ(tester.getRegion(2), 0);
+  ASSERT_EQ(tester.getRegion(3), 3);
+  ASSERT_EQ(tester.getRegion(4), 4);
+  ASSERT_EQ(tester.getRegion(5), 5);
+  ASSERT_EQ(tester.getRegion(6), 6);
+  ASSERT_EQ(tester.getRegion(7), 3);
+  ASSERT_EQ(tester.getRegion(8), 3);
+  ASSERT_EQ(tester.getRegion(9), 3);
+  ASSERT_EQ(tester.getRegion(10), 3);
+  ASSERT_EQ(tester.getRegion(11), 11);
+  ASSERT_EQ(tester.getRegion(12), 0);
+  ASSERT_EQ(tester.getRegion(13), 13);
+  ASSERT_EQ(tester.getRegion(14), 1);
+  ASSERT_EQ(tester.getRegion(15), 13);
+  ASSERT_EQ(tester.getRegion(16), 0);
+  ASSERT_EQ(tester.getRegion(17), 3);
+  ASSERT_EQ(tester.getRegion(18), 13);
+  ASSERT_EQ(tester.getRegion(19), 13);
+  ASSERT_EQ(tester.getRegion(20), 13);
+  ASSERT_EQ(tester.getRegion(21), 3);
+  ASSERT_EQ(tester.getRegion(22), 13);
+  ASSERT_EQ(tester.getRegion(23), 3);
+  ASSERT_EQ(tester.getRegion(24), 1);
+  ASSERT_EQ(tester.getRegion(25), 3);
+  ASSERT_EQ(tester.getRegion(26), 13);
+  ASSERT_EQ(tester.getRegion(27), 0);
+  ASSERT_EQ(tester.getRegion(28), 3);
+  ASSERT_EQ(tester.getRegion(29), 13);
+  ASSERT_EQ(tester.getRegion(30), 1);
+  ASSERT_EQ(tester.getRegion(31), 0);
+  ASSERT_EQ(tester.getRegion(32), 13);
+  ASSERT_EQ(tester.getRegion(33), 33);
+  ASSERT_EQ(tester.getRegion(34), 34);
+  ASSERT_EQ(tester.getRegion(35), 34);
+  ASSERT_EQ(tester.getRegion(36), 13);
+  ASSERT_EQ(tester.getRegion(37), 13);
+  ASSERT_EQ(tester.getRegion(38), 13);
+  ASSERT_EQ(tester.getRegion(39), 39);
+  ASSERT_EQ(tester.getRegion(40), 13);
+  ASSERT_EQ(tester.getRegion(41), 13);
+  ASSERT_EQ(tester.getRegion(42), 33);
+  ASSERT_EQ(tester.getRegion(43), 13);
+  ASSERT_EQ(tester.getRegion(44), 3);
+}
+
 // This test tests the semantics of assignment
 TEST(PartitionUtilsTest, TestAssign) {
   Partition p1;