Merge pull request #59243 from rctcwyvrn/lily-dictionary-optimization

Optimize Dictionary.filter for native dictionaries

Author: lorentey

stdlib/public/core/Dictionary.swift

@@ -608,15 +608,19 @@ extension Dictionary {
   public __consuming func filter(
     _ isIncluded: (Element) throws -> Bool
   ) rethrows -> [Key: Value] {
-    // FIXME(performance): Try building a bitset of elements to keep, so that we
-    // eliminate rehashings during insertion.
-    var result = _NativeDictionary<Key, Value>()
-    for element in self {
-      if try isIncluded(element) {
-        result.insertNew(key: element.key, value: element.value)
+  #if _runtime(_ObjC)
+    guard _variant.isNative else {
+      // Slow path for bridged dictionaries
+      var result = _NativeDictionary<Key, Value>()
+      for element in self {
+        if try isIncluded(element) {
+          result.insertNew(key: element.key, value: element.value)
+        }
       }
+      return Dictionary(_native: result)
     }
-    return Dictionary(_native: result)
+  #endif
+    return Dictionary(_native: try _variant.asNative.filter(isIncluded))
   }
 }
 

stdlib/public/core/NativeDictionary.swift

@@ -599,6 +599,27 @@ extension _NativeDictionary {
     (_values + a.offset).moveInitialize(from: _values + b.offset, count: 1)
     (_values + b.offset).initialize(to: value)
   }
+  
+  @_alwaysEmitIntoClient
+  internal func extractDictionary(
+    using bitset: _UnsafeBitset, 
+    count: Int
+  ) -> _NativeDictionary<Key, Value> {
+    var count = count
+    if count == 0 { return _NativeDictionary<Key, Value>() }
+    if count == self.count { return self }
+    let result = _NativeDictionary<Key, Value>(capacity: count)
+    for offset in bitset {
+      let key = self.uncheckedKey(at: Bucket(offset: offset))
+      let value = self.uncheckedValue(at: Bucket(offset: offset))
+      result._unsafeInsertNew(key: key, value: value)
+      // The hash table can have set bits after the end of the bitmap.
+      // Ignore them.
+      count -= 1
+      if count == 0 { break }
+    }
+    return result
+  }
 }
 
 extension _NativeDictionary where Value: Equatable {
@@ -771,6 +792,26 @@ extension _NativeDictionary { // High-level operations
       }
     }
   }
+
+  @_alwaysEmitIntoClient
+  internal func filter(
+    _ isIncluded: (Element) throws -> Bool
+  ) rethrows -> _NativeDictionary<Key, Value> {
+    try _UnsafeBitset.withTemporaryBitset(
+      capacity: _storage._bucketCount
+    ) { bitset in
+      var count = 0
+      for bucket in hashTable {
+        if try isIncluded(
+          (uncheckedKey(at: bucket), uncheckedValue(at: bucket))
+        ) {
+          bitset.uncheckedInsert(bucket.offset)
+          count += 1
+        }
+      }
+      return extractDictionary(using: bitset, count: count)
+    }
+  }
 }
 
 extension _NativeDictionary: Sequence {

validation-test/stdlib/Dictionary.swift

@@ -2013,6 +2013,19 @@ DictionaryTestSuite.test("mapValues(_:)") {
   }
 }
 
+DictionaryTestSuite.test("filter(_:)") {
+  let d1 = [1: 1, 2: 2, 3: 100, 4: 4, 5: 100, 6: 6]
+  let d2 = d1.filter() {key, value in key == value}
+
+  expectEqual(d2.count, 4)
+  for (key, value) in d2 {
+    expectEqual(key, value)
+  }
+
+  expectNil(d2[3])
+  expectNil(d2[5])
+}
+
 DictionaryTestSuite.test("capacity/init(minimumCapacity:)") {
   let d0 = Dictionary<String, Int>(minimumCapacity: 0)
   expectGE(d0.capacity, 0)