[stdlib] Optimize Set.filter(_:)

This works the same way as `Set.subtracting<S>(_:)`, and has similar performance benefits.

Author: lorentey

stdlib/public/core/HashTable.swift

@@ -42,7 +42,7 @@ internal struct _HashTable {
   @inlinable
   internal var bucketCount: Int {
     @inline(__always) get {
-      return bucketMask &+ 1
+      return _assumeNonNegative(bucketMask &+ 1)
     }
   }
 

stdlib/public/core/NativeSet.swift

@@ -658,6 +658,25 @@ extension _NativeSet {
     }
   }
 
+  @_alwaysEmitIntoClient
+  internal __consuming func extractSubset(
+    using bitset: _UnsafeBitset,
+    count: Int
+  ) -> _NativeSet {
+    var count = count
+    if count == 0 { return _NativeSet() }
+    if count == self.count { return self }
+    let result = _NativeSet(capacity: count)
+    for offset in bitset {
+      result._unsafeInsertNew(self.uncheckedElement(at: Bucket(offset: offset)))
+      // The hash table can have set bits after the end of the bitmap.
+      // Ignore them.
+      count -= 1
+      if count == 0 { break }
+    }
+    return result
+  }
+
   @_alwaysEmitIntoClient
   internal __consuming func subtracting<S: Sequence>(_ other: S) -> _NativeSet
   where S.Element == Element {
@@ -677,18 +696,24 @@ extension _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 }
+      _internalInvariant(difference.count > 0)
+      return extractSubset(using: difference, count: remainingCount)
+    }
+  }
+
+  @_alwaysEmitIntoClient
+  internal __consuming func filter(
+    _ isIncluded: (Element) throws -> Bool
+  ) rethrows -> _NativeSet<Element> {
+    try _UnsafeBitset.withTemporaryBitset(capacity: bucketCount) { bitset in
+      var count = 0
+      for bucket in hashTable {
+        if try isIncluded(uncheckedElement(at: bucket)) {
+          bitset.uncheckedInsert(bucket.offset)
+          count += 1
+        }
       }
-      return result
+      return extractSubset(using: bitset, count: count)
     }
   }
 }

stdlib/public/core/Set.swift

@@ -295,14 +295,7 @@ extension Set {
   public __consuming func filter(
     _ isIncluded: (Element) throws -> Bool
   ) rethrows -> Set {
-    // FIXME(performance): Eliminate rehashes by using a bitmap.
-    var result = Set()
-    for element in self {
-      if try isIncluded(element) {
-        result.insert(element)
-      }
-    }
-    return result
+    return try Set(_native: _variant.filter(isIncluded))
   }
 }
 

stdlib/public/core/SetVariant.swift

@@ -376,3 +376,24 @@ extension Set._Variant {
   }
 }
 
+extension Set._Variant {
+  @_alwaysEmitIntoClient
+  internal __consuming func filter(
+    _ isIncluded: (Element) throws -> Bool
+  ) rethrows -> _NativeSet<Element> {
+#if _runtime(_ObjC)
+    guard isNative else {
+      var result = _NativeSet<Element>()
+      for cocoaElement in asCocoa {
+        let nativeElement = _forceBridgeFromObjectiveC(
+          cocoaElement, Element.self)
+        if try isIncluded(nativeElement) {
+          result.insertNew(nativeElement, isUnique: true)
+        }
+      }
+      return result
+    }
+#endif
+    return try asNative.filter(isIncluded)
+  }
+}