This post will dive into a PR I made to the Swift project that improves the performance of IndexSet.union. Instruments is used to profile performance. If you are not familiar with Instruments (part of Xcode) check out this Instruments Overview.

Previous union Implementation

The union of two IndexSet objects is the union of the integer value collections the IndexSet objects contain. The existing implementation used two loops to insert the ranges from self and other into a new IndexSet:

public func union(_ other: IndexSet) -> IndexSet {
    // This algorithm is naïve but it works. We could avoid calling insert in some cases.

    var result = IndexSet()
    for r in self.rangeView {
        result.insert(integersIn: r)
    }
    
    for r in other.rangeView {
         result.insert(integersIn: r)
    }

    return result
}

Identifying The Performance Improvement

Looking for for loops and while loops with comments can be very insightful when evaluating a new codebase for performance opportunities. The previous author left the following comment:

// This algorithm is naïve but it works. We could avoid calling insert in some cases.

Thinking about a union, all elements from both sets will be contained in the final result. Thus, it may be worth investigating how to eliminate one of the two for loops from the existing implementation.

Faster union Implementation

A copy of the sparser IndexSet (the index set with a larger number of internal ranges) can be used instead of an existing loop for the best performance boost. Since Swift has copy-on-write mechanics, the sparser IndexSet can be assigned to a new variable and Swift will take care of the copy as needed to maintain correctness:

public func union(_ other: IndexSet) -> IndexSet {
    var result: IndexSet
    var dense: IndexSet

    // Prepare to make a copy of the more sparse IndexSet to prefer copy over repeated inserts
    if self.rangeView.count > other.rangeView.count {
        result = self
        dense = other
    } else {
        result = other
        dense = self
    }

    // Insert each range from the less sparse IndexSet
    dense.rangeView.forEach {
        result.insert(integersIn: $0)
    }

    return result
}

Eliminating a loop brings significant performance improvements, particularly since ranges are internally managed in an IndexSet by NSIndexSet with an expensive expensive add(in:) function.

In cases where the union is taken of sets with different range counts, for example taking the union of a set of 200 ranges with a set of 1,000 ranges, a 300% speedup is observed. In cases where each IndexSet contains the same number of values, a 30% speedup is observed.