[stdlib] Optimize Set.intersection(_:)

Use a temporary bitset to speed up the `Sequence` variant by roughly a factor of ~4-6, and the set/set variant by a factor of ~1-4, depending on the ratio of overlapping elements.

Author: lorentey

stdlib/public/core/NativeSet.swift

@@ -674,6 +674,7 @@ extension _NativeSet {
       count -= 1
       if count == 0 { break }
     }
+    _internalInvariant(result.count == bitset.count)
     return result
   }
 
@@ -716,4 +717,49 @@ extension _NativeSet {
       return extractSubset(using: bitset, count: count)
     }
   }
+
+  @_alwaysEmitIntoClient
+  internal __consuming func intersection(
+    _ other: _NativeSet<Element>
+  ) -> _NativeSet<Element> {
+    // Prefer to iterate over the smaller set. However, we must be careful to
+    // only include elements from `self`, not `other`.
+    guard self.count <= other.count else {
+      return genericIntersection(other)
+    }
+    // Rather than directly creating a new set, mark common elements in a bitset
+    // first. This minimizes hashing, and ensures that we'll have an exact count
+    // for the result set, preventing rehashings during insertions.
+    return _UnsafeBitset.withTemporaryBitset(capacity: bucketCount) { bitset in
+      var count = 0
+      for bucket in hashTable {
+        if other.find(uncheckedElement(at: bucket)).found {
+          bitset.uncheckedInsert(bucket.offset)
+          count += 1
+        }
+      }
+      return extractSubset(using: bitset, count: count)
+    }
+  }
+
+  @_alwaysEmitIntoClient
+  internal __consuming func genericIntersection<S: Sequence>(
+    _ other: S
+  ) -> _NativeSet<Element>
+  where S.Element == Element {
+    // Rather than directly creating a new set, mark common elements in a bitset
+    // first. This minimizes hashing, and ensures that we'll have an exact count
+    // for the result set, preventing rehashings during insertions.
+    _UnsafeBitset.withTemporaryBitset(capacity: bucketCount) { bitset in
+      var count = 0
+      for element in other {
+        let (bucket, found) = find(element)
+        if found {
+          bitset.uncheckedInsert(bucket.offset)
+          count += 1
+        }
+      }
+      return extractSubset(using: bitset, count: count)
+    }
+  }
 }

stdlib/public/core/Set.swift

@@ -949,8 +949,10 @@ extension Set: SetAlgebra {
   @inlinable
   public __consuming func intersection<S: Sequence>(_ other: S) -> Set<Element>
   where S.Element == Element {
-    let otherSet = Set(other)
-    return intersection(otherSet)
+    if let other = other as? Set<Element> {
+      return self.intersection(other)
+    }
+    return Set(_native: _variant.convertedToNative.genericIntersection(other))
   }
 
   /// Removes the elements of the set that aren't also in the given sequence.
@@ -969,20 +971,10 @@ extension Set: SetAlgebra {
   @inlinable
   public mutating func formIntersection<S: Sequence>(_ other: S)
   where S.Element == Element {
-    // Because `intersect` needs to both modify and iterate over
-    // the left-hand side, the index may become invalidated during
-    // traversal so an intermediate set must be created.
-    //
-    // FIXME(performance): perform this operation at a lower level
-    // to avoid invalidating the index and avoiding a copy.
-    let result = self.intersection(other)
-
-    // The result can only have fewer or the same number of elements.
-    // If no elements were removed, don't perform a reassignment
-    // as this may cause an unnecessary uniquing COW.
-    if result.count != count {
-      self = result
-    }
+    // FIXME: This discards storage reserved with reserveCapacity.
+    // FIXME: Depending on the ratio of elements kept in the result, it may be
+    // faster to do the removals in place, in bulk.
+    self = self.intersection(other)
   }
 
   /// Returns a new set with the elements that are either in this set or in the
@@ -1233,15 +1225,7 @@ extension Set {
   /// - Returns: A new set.
   @inlinable
   public __consuming func intersection(_ other: Set<Element>) -> Set<Element> {
-    var newSet = Set<Element>()
-    let (smaller, larger) =
-      count < other.count ? (self, other) : (other, self)
-    for member in smaller {
-      if larger.contains(member) {
-        newSet.insert(member)
-      }
-    }
-    return newSet
+    Set(_native: _variant.intersection(other))
   }
 
   /// Removes the elements of the set that are also in the given sequence and

stdlib/public/core/SetVariant.swift

@@ -396,4 +396,36 @@ extension Set._Variant {
 #endif
     return try asNative.filter(isIncluded)
   }
+
+  @_alwaysEmitIntoClient
+  internal __consuming func intersection(
+    _ other: Set<Element>
+  ) -> _NativeSet<Element> {
+#if _runtime(_ObjC)
+    switch (self.isNative, other._variant.isNative) {
+    case (true, true):
+      return asNative.intersection(other._variant.asNative)
+    case (true, false):
+      return asNative.genericIntersection(other)
+    case (false, false):
+      return _NativeSet(asCocoa).genericIntersection(other)
+    case (false, true):
+      // Note: It is tempting to implement this as `that.intersection(this)`,
+      // but intersection isn't symmetric -- the result should only contain
+      // elements from `self`.
+      let that = other._variant.asNative
+      var result = _NativeSet<Element>()
+      for cocoaElement in asCocoa {
+        let nativeElement = _forceBridgeFromObjectiveC(
+          cocoaElement, Element.self)
+        if that.contains(nativeElement) {
+          result.insertNew(nativeElement, isUnique: true)
+        }
+      }
+      return result
+    }
+#else
+    return asNative.intersection(other._variant.asNative)
+#endif
+  }
 }