[stdlib] Optimize Set.subtracting even more

Author: lorentey

stdlib/public/core/NativeSet.swift

@@ -681,13 +681,31 @@ extension _NativeSet {
   internal __consuming func subtracting<S: Sequence>(_ other: S) -> _NativeSet
   where S.Element == Element {
     guard count > 0 else { return _NativeSet() }
+
+    // Find one item that we need to remove before creating a result set.
+    var it = other.makeIterator()
+    var bucket: Bucket? = nil
+    while let next = it.next() {
+      let (b, found) = find(next)
+      if found {
+        bucket = b
+        break
+      }
+    }
+    guard let bucket = bucket else { return self }
+
     // 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 removed = difference.uncheckedRemove(bucket.offset)
+      _internalInvariant(removed)
+      remainingCount -= 1
+
+      while let element = it.next() {
         let (bucket, found) = find(element)
         if found {
           if difference.uncheckedRemove(bucket.offset) {

stdlib/public/core/Set.swift

@@ -1158,6 +1158,16 @@ extension Set {
   /// - Returns: A new set.
   @inlinable
   public __consuming func subtracting(_ other: Set<Element>) -> Set<Element> {
+    // Heuristic: if `other` is small enough, it's better to make a copy of the
+    // set and remove each item one by one. (The best cutoff point depends on
+    // the `Element` type; the one below is an educated guess.) FIXME: Derive a
+    // better cutoff by benchmarking.
+    if other.count <= self.count / 8 {
+      var copy = self
+      copy._subtract(other)
+      return copy
+    }
+    // Otherwise do a regular subtraction using a temporary bitmap.
     return self._subtracting(other)
   }