Improve RangeSet union performance

The existing implementation of `RangeSet.formUnion` performs a naive
insertion of each range of one range set into the other. This
requires shuffling the array of ranges down with each insertion, which
can have quadratic performance.

This change uses the known invariants of the ranges array to instead
perform the merge in linear time. Each range in the resulting array
is determined by finding the next lowest bound between the two range
sets, and then searching for the first upper bound that isn't included
in the merged range set contents.

rdar://129296438

Author: natecook1000

stdlib/public/core/RangeSet.swift

@@ -251,10 +251,11 @@ extension RangeSet {
   /// Adds the contents of the given range set to this range set.
   ///
   /// - Parameter other: A range set to merge with this one.
+  ///
+  /// - Complexity: O(*m* + *n*), where *m* and *n* are the number of ranges in
+  ///   the range set.
   public mutating func formUnion(_ other: __owned RangeSet<Bound>) {
-    for range in other._ranges {
-      insert(contentsOf: range)
-    }
+    self = self.union(other)
   }
 
   /// Removes the contents of this range set that aren't also in the given
@@ -293,9 +294,7 @@ extension RangeSet {
   public __consuming func union(
     _ other: __owned RangeSet<Bound>
   ) -> RangeSet<Bound> {
-    var result = self
-    result.formUnion(other)
-    return result
+    return RangeSet(_ranges: _ranges._union(other._ranges))
   }
 
   /// Returns a new range set containing the contents of both this set and the

stdlib/public/core/RangeSetRanges.swift

@@ -265,6 +265,77 @@ extension RangeSet.Ranges {
 
     return Self(_ranges: result)
   }
+  
+  @usableFromInline
+  internal func _union(_ other: Self) -> Self {
+    // Empty cases
+    if other.isEmpty {
+      return self
+    } else if self.isEmpty {
+      return other
+    }
+    
+    // Instead of naively inserting the ranges of `other` into `self`,
+    // which can cause reshuffling with every insertion, this approach
+    // uses the guarantees that each array of ranges is non-overlapping and in
+    // increasing order to directly derive the union.
+    //
+    // Each range in the resulting range set is found by:
+    //
+    //    1. Finding the current lowest bound of the two range sets.
+    //    2. Searching for the first upper bound that is outside the merged
+    //       boundaries of the two range sets.
+    
+    // Use temporaries so that we can swap a/b, to simplify the logic below
+    var a = self._storage
+    var b = other._storage
+    var aIndex = a.startIndex
+    var bIndex = b.startIndex
+    
+    var result: [Range<Bound>] = []
+    while aIndex < a.endIndex, bIndex < b.endIndex {
+      // Make sure that `a` is the source of the lower bound and `b` is the
+      // potential source for the upper bound.
+      if b[bIndex].lowerBound < a[aIndex].lowerBound {
+        swap(&a, &b)
+        swap(&aIndex, &bIndex)
+      }
+      
+      var candidateRange = a[aIndex]
+      aIndex += 1
+      
+      // Look for the correct upper bound, which is the first upper bound that
+      // isn't contained in the next range of the "other" ranges array.
+      while bIndex < b.endIndex, candidateRange.upperBound >= b[bIndex].lowerBound {
+        if candidateRange.upperBound >= b[bIndex].upperBound {
+          // The range `b[bIndex]` is entirely contained by `candidateRange`,
+          // so we need to advance and look at the next range in `b`.
+          bIndex += 1
+        } else {
+          // The range `b[bIndex]` extends past `candidateRange`, so:
+          //
+          //   1. We grow `candidateRange` to the upper bound of `b[bIndex]`
+          //   2. We swap the two range arrays, so that we're looking for the
+          //      new upper bound in the other array.
+          candidateRange = candidateRange.lowerBound ..< b[bIndex].upperBound
+          bIndex += 1
+          swap(&a, &b)
+          swap(&aIndex, &bIndex)
+        }
+      }
+      
+      result.append(candidateRange)
+    }
+    
+    // Collect any remaining ranges without needing to merge.
+    if aIndex < a.endIndex {
+      result.append(contentsOf: a[aIndex...])
+    } else if bIndex < b.endIndex {
+      result.append(contentsOf: b[bIndex...])
+    }
+
+    return Self(_ranges: result)
+  }
 }
 
 @available(SwiftStdlib 6.0, *)
@@ -344,4 +415,4 @@ internal struct _Pair<Element>: RandomAccessCollection {
       }
     }
   }
-}
+}

validation-test/stdlib/RangeSet.swift

@@ -174,6 +174,55 @@ if #available(SwiftStdlib 6.0, *) {
     }
   }
 
+  RangeSetTests.test("union") {
+    func unionViaSet(
+      _ s1: RangeSet<Int>,
+      _ s2: RangeSet<Int>
+    ) -> RangeSet<Int> {
+      let set1 = Set(parent.indices[s1])
+      let set2 = Set(parent.indices[s2])
+      return RangeSet(set1.union(set2), within: parent)
+    }
+
+    func testUnion(
+      _ set1: RangeSet<Int>,
+      _ set2: RangeSet<Int>,
+      expect union: RangeSet<Int>
+    ) {
+      expectEqual(set1.union(set2), union)
+      expectEqual(set2.union(set1), union)
+
+      var set3 = set1
+      set3.formUnion(set2)
+      expectEqual(set3, union)
+
+      set3 = set2
+      set3.formUnion(set1)
+      expectEqual(set3, union)
+    }
+    
+    // Simple tests
+    testUnion([0..<5, 9..<14],
+              [1..<3, 4..<6, 8..<12],
+              expect: [0..<6, 8..<14])
+    
+    testUnion([10..<20, 50..<60],
+              [15..<55, 58..<65],
+              expect: [10..<65])
+
+    // Test with upper bound / lower bound equality
+    testUnion([10..<20, 30..<40],
+              [15..<30, 40..<50],
+              expect: [10..<50])
+    
+    for _ in 0..<100 {
+      let set1 = buildRandomRangeSet()
+      let set2 = buildRandomRangeSet()
+      testUnion(set1, set2,
+                expect: unionViaSet(set1, set2))
+    }
+  }
+  
   RangeSetTests.test("intersection") {
     func intersectionViaSet(
       _ s1: RangeSet<Int>,