[stdlib] Optimize Set.subtracting(_:)

Use a temporary bitset to avoid hashing elements more than once, and to prevent rehashings during the creation of the result set.

This leads to a speedup of about 0-4x, depending on the number of elements removed.

Author: lorentey

stdlib/public/core/Bitset.swift

@@ -369,3 +369,19 @@ extension _UnsafeBitset {
     }
   }
 }
+
+extension _UnsafeBitset {
+  @_alwaysEmitIntoClient
+  @inline(__always)
+  internal static func withTemporaryCopy<R>(
+    of original: _UnsafeBitset,
+    body: (_UnsafeBitset) throws -> R
+  ) rethrows -> R {
+    try _withTemporaryUninitializedBitset(
+      wordCount: original.wordCount
+    ) { bitset in
+      bitset.words.initialize(from: original.words, count: original.wordCount)
+      return try body(bitset)
+    }
+  }
+}

stdlib/public/core/HashTable.swift

@@ -52,6 +52,17 @@ internal struct _HashTable {
       return _UnsafeBitset.wordCount(forCapacity: bucketCount)
     }
   }
+
+  /// Return a bitset representation of the occupied buckets in this table.
+  ///
+  /// Note that if we have only a single partial word in the hash table's
+  /// bitset, then its out-of-bounds bits are guaranteed to be all set. These
+  /// filler bits are there to speed up finding holes -- they don't correspond
+  /// to occupied buckets in the table.
+  @_alwaysEmitIntoClient
+  internal var bitset: _UnsafeBitset {
+    _UnsafeBitset(words: words, wordCount: wordCount)
+  }
 }
 
 extension _HashTable {

stdlib/public/core/NativeSet.swift

@@ -657,4 +657,38 @@ extension _NativeSet {
       return true
     }
   }
+
+  @_alwaysEmitIntoClient
+  internal __consuming func subtracting<S: Sequence>(_ other: S) -> _NativeSet
+  where S.Element == Element {
+    guard count > 0 else { return _NativeSet() }
+    // Rather than directly creating a new set, calculate the difference in a
+    // bitset first. This ensures we hash each element (in both sets) only once,
+    // and that we'll have an exact count for the result set, preventing
+    // rehashings during insertions.
+    return _UnsafeBitset.withTemporaryCopy(of: hashTable.bitset) { difference in
+      var remainingCount = self.count
+      for element in other {
+        let (bucket, found) = find(element)
+        if found {
+          if difference.uncheckedRemove(bucket.offset) {
+            remainingCount -= 1
+            if remainingCount == 0 { return _NativeSet() }
+          }
+        }
+      }
+      _internalInvariant(remainingCount > 0)
+      if remainingCount == self.count { return self }
+      let result = _NativeSet(capacity: remainingCount)
+      for offset in difference {
+        result._unsafeInsertNew(
+          self.uncheckedElement(at: Bucket(offset: offset)))
+        // The hash table can have set bits after the end of the bitmap.
+        // Ignore them.
+        remainingCount -= 1
+        if remainingCount == 0 { break }
+      }
+      return result
+    }
+  }
 }

stdlib/public/core/Set.swift

@@ -903,9 +903,7 @@ extension Set: SetAlgebra {
     _ other: S
   ) -> Set<Element>
   where S.Element == Element {
-    var newSet = self
-    newSet.subtract(other)
-    return newSet
+    return Set(_native: _variant.convertedToNative.subtracting(other))
   }
 
   /// Removes the elements of the given sequence from the set.