[stdlib] Speed up Set.contains

- Simplify generated code; eliminate obvious inefficiencies.
- Revive guaranteedNative trick; removing the bridged branch does have a measurable impact for contains at least. (Apart from code size benefits)

Lookups in sets are still not quite as fast as on master for small sets (up to ~1k elements) when the element type has fast equality checks — I believe this is because the SIMD approach has worse latency than direct lookup, and the lookup chains aren’t long enough to take advantage of higher throughput. (The latency difference only matters when the set is fully cached, which is why the slowdown only occurs for small sets.)

Author: lorentey

stdlib/public/core/HashTable.swift

@@ -24,13 +24,15 @@ internal struct _HashTable {
   internal let buckets: UnsafeMutablePointer<Bucket>
 
   @usableFromInline
-  internal let bucketCount: Int
+  internal let bucketMask: Int
 
   @inlinable
   @inline(__always)
   internal init(buckets: UnsafeMutablePointer<Bucket>, bucketCount: Int) {
     self.buckets = buckets
-    self.bucketCount = bucketCount
+    // The bucket count is a positive power of two, so subtracting 1 will
+    // never overflow and get us a nice mask.
+    self.bucketMask = bucketCount &- 1
   }
 
   @inlinable
@@ -41,11 +43,9 @@ internal struct _HashTable {
   }
 
   @inlinable
-  internal var bucketMask: Int {
+  internal var bucketCount: Int {
     @inline(__always) get {
-      // The bucket count is a positive power of two, so subtracting 1 will
-      // never overflow and get us a nice mask.
-      return bucketCount &- 1
+      return bucketMask &+ 1
     }
   }
 }
@@ -67,7 +67,7 @@ extension _HashTable {
 
   /// The inverse of the maximum hash table load factor.
   private static var maxLoadFactorInverse: Double {
-    @inline(__always) get { return 100 / 90 }
+    @inline(__always) get { return 100 / 75 }
   }
 
   internal static func capacity(forScale scale: Int) -> Int {
@@ -107,14 +107,6 @@ extension _HashTable {
   @usableFromInline
   @_fixed_layout
   internal struct Entry {
-    @inlinable
-    internal static var occupiedFlag: UInt8 {
-      @inline(__always) get { return 0x80 }
-    }
-    @inlinable
-    internal static var payloadMask: UInt8 {
-      @inline(__always) get { return 0x7F }
-    }
     @inlinable
     internal static var unoccupied: Entry {
       @inline(__always) get { return Entry(_value: 0) }
@@ -132,16 +124,15 @@ extension _HashTable {
     @inlinable
     @inline(__always)
     internal init(payload: UInt8) {
-      _sanityCheck(payload & ~Entry.payloadMask == 0)
-      self.init(_value: Entry.occupiedFlag | payload)
+      _sanityCheck(payload != 0)
+      self.init(_value: payload)
     }
 
     @inlinable
     @inline(__always)
     internal init(forHashValue hashValue: Int) {
-      // Use the highest seven bits of the hash value.
-      let payload = UInt(bitPattern: hashValue) &>> (Int.bitWidth &- 7)
-      self.init(payload: UInt8(truncatingIfNeeded: payload))
+      let payload = UInt(bitPattern: hashValue) &>> (Int.bitWidth &- 8)
+      self.init(payload: Swift.max(1, UInt8(truncatingIfNeeded: payload)))
     }
 
     @inlinable
@@ -152,12 +143,23 @@ extension _HashTable {
     @inlinable
     internal var payload: UInt8 {
       @inline(__always) get {
-        return _value & _HashTable.Entry.payloadMask
+        return _value
       }
     }
   }
 }
 
+extension _HashTable.Entry {
+  @inlinable
+  var pattern: UInt64 {
+    @inline(__always) get {
+      // Fill a 64-bit integer with 8 copies of this entry.
+      let p = UInt64(truncatingIfNeeded: _value)
+      return p &* 0x01010101_01010101
+    }
+  }
+}
+
 extension _HashTable.Entry: Equatable {
   @inlinable
   @inline(__always)
@@ -270,7 +272,7 @@ extension _HashTable.Bucket {
     @inline(__always)
     get {
       // Holes are always at the end, so it's enough to check the highest byte.
-      return _value &>> (UInt64.bitWidth &- UInt8.bitWidth) != 0
+      return _value.leadingZeroBitCount < UInt8.bitWidth
     }
   }
 
@@ -295,7 +297,7 @@ extension _HashTable.Bucket {
     @inline(__always)
     get {
       // Holes are zero bytes at the most significant places.
-      let shift = UInt64.bitWidth - (_value.leadingZeroBitCount & ~7)
+      let shift = UInt64.bitWidth &- (_value.leadingZeroBitCount & ~7)
       return _HashTable.Slot(shift: shift)
     }
   }
@@ -363,14 +365,9 @@ extension _HashTable.Bucket {
 
   /// Returns a sequence of Slots matching the given entry.
   @inlinable
-  internal func slots(matching entry: _HashTable.Entry) -> SlotSet {
-    // Fill a 64-bit integer with 8 copies of the entry we're looking for.
-    var p = UInt64(entry._value)
-    p |= p &<< 8
-    p |= p &<< 16
-    p |= p &<< 32
-    // Xoring `p` to `self._value` turns matching bytes into zeroes.
-    p ^= self._value
+  @inline(__always)
+  internal func _slots(matching pattern: UInt64) -> SlotSet {
+    let p = self._value ^ pattern
     // The problem now reduces to finding zero bytes in `p`.  For every 8-bit
     // integer `b`, `x = ((b & 0x7F) + 0x7F) | b` has bit 7 set iff `b !=
     // 0`. Further, `~(x | 0x7F)` leaves bit 7 set to one iff `b == 0` and
@@ -384,6 +381,13 @@ extension _HashTable.Bucket {
     // bits to the start of their corresponding bytes.
     return SlotSet(_shifts: y &>> 7)
   }
+
+  /// Returns a sequence of Slots matching the given entry.
+  @inlinable
+  @inline(__always)
+  internal func slots(matching entry: _HashTable.Entry) -> SlotSet {
+    return _slots(matching: entry.pattern)
+  }
 }
 
 extension _HashTable {
@@ -483,7 +487,6 @@ extension _HashTable: Collection {
   }
 
   @inlinable
-  @inline(__always)
   internal func formIndex(after index: inout Index) {
     index.offset += 1
     if index.bucket >= bucketCount { return }
@@ -560,7 +563,7 @@ extension _HashTable {
     @usableFromInline
     let _hashTable: _HashTable
     @usableFromInline
-    let _entry: Entry
+    let _pattern: UInt64
     @usableFromInline
     var _bucket: Int
     @usableFromInline
@@ -569,14 +572,15 @@ extension _HashTable {
     @inlinable
     internal init(hashTable: _HashTable, hashValue: Int) {
       self._hashTable = hashTable
-      self._entry = Entry(forHashValue: hashValue)
       self._bucket = hashTable._idealBucket(forHashValue: hashValue)
-      self._matches = hashTable.buckets[_bucket].slots(matching: _entry)
+      let b = hashTable.buckets[_bucket]
+      self._pattern = Entry(forHashValue: hashValue).pattern
+      self._matches = b._slots(matching: _pattern)
     }
 
     @inlinable
     internal mutating func next() -> (index: Index, found: Bool) {
-      repeat {
+      while true {
         if let slot = _matches.next() {
           return (Index(bucket: _bucket, slot: slot), true)
         }
@@ -586,8 +590,8 @@ extension _HashTable {
           return (Index(bucket: _bucket, slot: hole), false)
         }
         _bucket = _hashTable._succ(_bucket)
-        _matches = _hashTable.buckets[_bucket].slots(matching: _entry)
-      } while true
+        _matches = _hashTable.buckets[_bucket]._slots(matching: _pattern)
+      }
     }
   }
 
@@ -598,7 +602,37 @@ extension _HashTable {
 }
 
 extension _HashTable {
-  @usableFromInline
+  @inlinable
+  @inline(__always)
+  func contains<Element: Equatable>(
+    hashValue: Int,
+    element: Element,
+    elements: UnsafePointer<Element>
+  ) -> Bool {
+    var bucket = _idealBucket(forHashValue: hashValue)
+    var b = self.buckets[bucket]
+    let pattern = Entry(forHashValue: hashValue).pattern
+    var matches = b._slots(matching: pattern)._shifts
+    while true {
+      if _fastPath(matches != 0) {
+        let shift = matches.trailingZeroBitCount
+        let index = Index(bucket: bucket, slotOffset: shift &>> 3)
+        if elements[index.offset] == element { return true }
+        matches &= matches &- 1
+      } else if b.isFull {
+        bucket = _succ(bucket)
+        b = self.buckets[bucket]
+        matches = b._slots(matching: pattern)._shifts
+      } else {
+        return false
+      }
+    }
+  }
+}
+
+extension _HashTable {
+  @inlinable
+  @inline(__always)
   @_effects(releasenone)
   internal func copyContents(of other: _HashTable) {
     _sanityCheck(bucketCount == other.bucketCount)
@@ -607,7 +641,8 @@ extension _HashTable {
 
   /// Insert a new entry with the specified hash value into the table.
   /// The entry must not already exist in the table -- duplicates are ignored.
-  @usableFromInline
+  @inlinable
+  @inline(__always)
   @_effects(releasenone)
   internal func insertNew(hashValue: Int) -> Index {
     var bucket = _idealBucket(forHashValue: hashValue)
@@ -621,14 +656,16 @@ extension _HashTable {
   /// Insert a new entry for an element with the specified hash value at
   /// `bucket`. The bucket must have been returned by `lookupFirst` or
   /// `lookupNext` for the same hash value, with `found == false`.
-  @usableFromInline
+  @inlinable
   @inline(__always)
   @_effects(releasenone)
   internal func insert(hashValue: Int, at index: Index) {
     _sanityCheck(!isOccupied(index))
     self[index] = Entry(forHashValue: hashValue)
   }
 
+  @inlinable
+  @inline(__always)
   internal func removeAll() {
     buckets.assign(repeating: Bucket(0), count: bucketCount)
   }

stdlib/public/core/Set.swift

@@ -1513,6 +1513,7 @@ internal class _SwiftRawSetStorage: _SwiftNativeNSSet {
 
 /// The storage class for the singleton empty set.
 /// The single instance of this class is created by the runtime.
+@_fixed_layout
 @usableFromInline
 internal class _SwiftEmptySetStorage: _SwiftRawSetStorage {
   override internal init(_doNotCallMe: ()) {
@@ -1803,7 +1804,7 @@ extension _NativeSet {
     }
   }
 
-  @inlinable
+  @usableFromInline @_transparent
   internal var elements: UnsafeMutablePointer<Element> {
     @inline(__always)
     get {
@@ -1818,8 +1819,7 @@ extension _NativeSet {
     }
   }
 
-  @inlinable
-  @inline(__always)
+  @usableFromInline @_transparent
   internal func uncheckedElement(at index: Index) -> Element {
     _sanityCheck(hashTable.isOccupied(index))
     return elements[index.offset]
@@ -1954,7 +1954,10 @@ extension _NativeSet: _SetBuffer {
     //   // Fast path that avoids computing the hash of the key.
     //   return false
     // }
-    return find(member).found
+    return hashTable.contains(
+      hashValue: hashValue(for: member),
+      element: member,
+      elements: elements)
   }
 
   @inlinable
@@ -2963,9 +2966,17 @@ extension Set._Variant: _SetBuffer {
     }
   }
 
+  @usableFromInline @_transparent
+  internal var guaranteedNative: Bool {
+    return _canBeClass(Element.self) == 0
+  }
+
   @inlinable
   @inline(__always)
   internal func contains(_ member: Element) -> Bool {
+    if guaranteedNative {
+      return asNative.contains(member)
+    }
     switch self {
     case .native:
       return asNative.contains(member)
@@ -3766,20 +3777,26 @@ public typealias SetIterator<Element: Hashable> = Set<Element>.Iterator
 extension Set {
   // FIXME: Remove
   public // @testable performance metrics
-  var _stats: (maxLookups: Int, averageLookups: Double, maxCollisions: Int)? {
+  var _stats: (maxLookups: Int, averageLookups: Double, maxCollisions: Int, averageCollisions: Double)? {
     guard case .native(let native) = _variant else { return nil }
-    guard native.count > 0 else { return (0, 0, 0) }
+    guard native.count > 0 else { return (0, 0, 0, 0) }
     defer { _fixLifetime(self) }
     var maxLookups = 0
     var sumLookups = 0
     var maxCollisions = 0
+    var sumCollisions = 0
     for i in native.indices {
       let stats = native.stats(ofElementAt: i)
       maxLookups = Swift.max(maxLookups, stats.displacement + 1)
       sumLookups += stats.displacement + 1
       maxCollisions = Swift.max(maxCollisions, stats.collisions)
+      sumCollisions += stats.collisions
     }
-    return (maxLookups, Double(sumLookups) / Double(count), maxCollisions)
+    return (
+      maxLookups,
+      Double(sumLookups) / Double(count),
+      maxCollisions,
+      Double(sumCollisions) / Double(count))
   }
 }
 
@@ -3826,7 +3843,7 @@ extension _NativeSet {
       let hashValue = self.hashValue(for: element)
       let stats = self.stats(ofElementAt: i)
       let h = String(UInt(bitPattern: hashValue), radix: 16)
-      let hl = String(hashTable[i].payload << 1, radix: 16)
+      let hl = String(hashTable[i].payload, radix: 16)
       result += "  <\(i.offset)> "
       result += "delta: \(stats.displacement) (\(stats.collisions) coll) "
       result += "hash: \(hl)/\(h): "