[stdlib] Set, Dictionary: Clean up storage allocation

Author: lorentey

stdlib/public/core/DictionaryStorage.swift

@@ -359,21 +359,24 @@ final internal class _DictionaryStorage<Key: Hashable, Value>
 extension _DictionaryStorage {
   @usableFromInline
   @_effects(releasenone)
-  internal static func reallocate(
+  internal static func copy(
+    original: _RawDictionaryStorage
+  ) -> _DictionaryStorage {
+    return allocate(
+      scale: original._scale,
+      age: original._age,
+      seed: original._seed)
+  }
+
+  @usableFromInline
+  @_effects(releasenone)
+  static internal func resize(
     original: _RawDictionaryStorage,
-    capacity: Int
-  ) -> (storage: _DictionaryStorage, rehash: Bool) {
-    _sanityCheck(capacity >= original._count)
+    capacity: Int,
+    move: Bool
+  ) -> _DictionaryStorage {
     let scale = _HashTable.scale(forCapacity: capacity)
-    let rehash = (scale != original._scale)
-    if rehash {
-      return (.allocate(scale: scale, age: nil, seed: nil), rehash)
-    }
-    return (
-      .allocate(scale: scale, age: original._age, seed: original._seed),
-      rehash)
-    }
-    return (newStorage, rehash)
+    return allocate(scale: scale, age: nil, seed: nil)
   }
 
   @usableFromInline
@@ -392,7 +395,7 @@ extension _DictionaryStorage {
   ) -> _DictionaryStorage {
     let scale = _HashTable.scale(forCapacity: capacity)
     let age = _HashTable.age(for: cocoa.object)
-    return allocate(scale: scale, age: age)
+    return allocate(scale: scale, age: age, seed: nil)
   }
 #endif
 

stdlib/public/core/DictionaryVariant.swift

@@ -328,28 +328,22 @@ extension Dictionary._Variant {
     }
   }
 
-  /// Ensure uniquely held native storage, while preserving the given index.
-  /// (If the variant had bridged storage, then the returned index will be the
-  /// corresponding native representation. Otherwise it's kept the same.)
   @inlinable
   @inline(__always)
   internal mutating func ensureUniqueNative() -> _NativeDictionary<Key, Value> {
-    switch self {
-    case .native:
-      let isUnique = isUniquelyReferenced()
-      if !isUnique {
-        let rehashed = asNative.copy(capacity: asNative.capacity)
-        _sanityCheck(!rehashed)
-      }
-      return asNative
 #if _runtime(_ObjC)
-    case .cocoa(let cocoa):
+    if case .cocoa(let cocoa) = self {
       cocoaPath()
       let native = _NativeDictionary<Key, Value>(cocoa)
       self = .native(native)
       return native
+    }
 #endif
+    let isUnique = isUniquelyReferenced()
+    if !isUnique {
+      asNative.copy()
     }
+    return asNative
   }
 
   @inlinable

stdlib/public/core/HashTable.swift

@@ -92,7 +92,7 @@ extension _HashTable {
 
   internal static func hashSeed(
     for object: AnyObject,
-    scale: Int
+    scale: Int8
   ) -> Int {
 #if false // FIXME: Enable per-instance seeding
     // We generate a new hash seed whenever a new hash table is allocated and
@@ -106,7 +106,7 @@ extension _HashTable {
       // When we're using deterministic hashing, the scale value as the seed is
       // still allowed, and it covers most cases. (Unfortunately some operations
       // that merge two similar-sized hash tables will still be quadratic.)
-      return scale
+      return Int(scale)
     }
     // Use the object address as the hash seed. This is cheaper than
     // SystemRandomNumberGenerator, while it has the same practical effect.
@@ -117,7 +117,7 @@ extension _HashTable {
     return unsafeBitCast(object, to: Int.self)
 #else
     // Use per-capacity seeding for now.
-    return scale
+    return Int(scale)
 #endif
   }
 }

stdlib/public/core/NativeDictionary.swift

@@ -182,8 +182,11 @@ extension _NativeDictionary { // ensureUnique
   @inlinable
   internal mutating func resize(capacity: Int) {
     let capacity = Swift.max(capacity, self.capacity)
-    let result = _NativeDictionary(
-      _DictionaryStorage<Key, Value>.allocate(capacity: capacity))
+    let newStorage = _DictionaryStorage<Key, Value>.resize(
+      original: _storage,
+      capacity: capacity,
+      move: true)
+    let result = _NativeDictionary(newStorage)
     if count > 0 {
       for bucket in hashTable {
         let key = (_keys + bucket.offset).move()
@@ -199,31 +202,40 @@ extension _NativeDictionary { // ensureUnique
   }
 
   @inlinable
-  internal mutating func copy(capacity: Int) -> Bool {
+  internal mutating func copyAndResize(capacity: Int) {
     let capacity = Swift.max(capacity, self.capacity)
-    let (newStorage, rehash) = _DictionaryStorage<Key, Value>.reallocate(
+    let newStorage = _DictionaryStorage<Key, Value>.resize(
       original: _storage,
-      capacity: capacity)
+      capacity: capacity,
+      move: false)
     let result = _NativeDictionary(newStorage)
     if count > 0 {
-      if rehash {
-        for bucket in hashTable {
-          result._unsafeInsertNew(
-            key: self.uncheckedKey(at: bucket),
-            value: self.uncheckedValue(at: bucket))
-        }
-      } else {
-        result.hashTable.copyContents(of: hashTable)
-        result._storage._count = self.count
-        for bucket in hashTable {
-          let key = uncheckedKey(at: bucket)
-          let value = uncheckedValue(at: bucket)
-          result.uncheckedInitialize(at: bucket, toKey: key, value: value)
-        }
+      for bucket in hashTable {
+        result._unsafeInsertNew(
+          key: self.uncheckedKey(at: bucket),
+          value: self.uncheckedValue(at: bucket))
+      }
+    }
+    _storage = result._storage
+  }
+
+  @inlinable
+  internal mutating func copy() {
+    let newStorage = _DictionaryStorage<Key, Value>.copy(original: _storage)
+    _sanityCheck(newStorage._scale == _storage._scale)
+    _sanityCheck(newStorage._age == _storage._age)
+    _sanityCheck(newStorage._seed == _storage._seed)
+    let result = _NativeDictionary(newStorage)
+    if count > 0 {
+      result.hashTable.copyContents(of: hashTable)
+      result._storage._count = self.count
+      for bucket in hashTable {
+        let key = uncheckedKey(at: bucket)
+        let value = uncheckedValue(at: bucket)
+        result.uncheckedInitialize(at: bucket, toKey: key, value: value)
       }
     }
     _storage = result._storage
-    return rehash
   }
 
   /// Ensure storage of self is uniquely held and can hold at least `capacity`
@@ -234,10 +246,15 @@ extension _NativeDictionary { // ensureUnique
     if _fastPath(capacity <= self.capacity && isUnique) {
       return false
     }
-    guard isUnique else {
-      return copy(capacity: capacity)
+    if isUnique {
+      resize(capacity: capacity)
+      return true
+    }
+    if capacity <= self.capacity {
+      copy()
+      return false
     }
-    resize(capacity: capacity)
+    copyAndResize(capacity: capacity)
     return true
   }
 
@@ -571,10 +588,7 @@ extension _NativeDictionary { // High-level operations
   internal func mapValues<T>(
     _ transform: (Value) throws -> T
   ) rethrows -> _NativeDictionary<Key, T> {
-    let (resultStorage, rehash) = _DictionaryStorage<Key, T>.reallocate(
-      original: _storage,
-      capacity: capacity)
-    _sanityCheck(!rehash)
+    let resultStorage = _DictionaryStorage<Key, T>.copy(original: _storage)
     _sanityCheck(resultStorage._seed == _storage._seed)
     let result = _NativeDictionary<Key, T>(resultStorage)
     // Because the current and new buffer have the same scale and seed, we can

stdlib/public/core/NativeSet.swift

@@ -154,7 +154,10 @@ extension _NativeSet { // ensureUnique
   @inlinable
   internal mutating func resize(capacity: Int) {
     let capacity = Swift.max(capacity, self.capacity)
-    let result = _NativeSet(_SetStorage<Element>.allocate(capacity: capacity))
+    let result = _NativeSet(_SetStorage<Element>.resize(
+        original: _storage,
+        capacity: capacity,
+        move: true))
     if count > 0 {
       for bucket in hashTable {
         let element = (self._elements + bucket.offset).move()
@@ -169,28 +172,36 @@ extension _NativeSet { // ensureUnique
   }
 
   @inlinable
-  internal mutating func copy(capacity: Int) -> Bool {
+  internal mutating func copyAndResize(capacity: Int) {
     let capacity = Swift.max(capacity, self.capacity)
-    let (newStorage, rehash) = _SetStorage<Element>.reallocate(
-      original: _storage,
-      capacity: capacity)
+    let result = _NativeSet(_SetStorage<Element>.resize(
+        original: _storage,
+        capacity: capacity,
+        move: false))
+    if count > 0 {
+      for bucket in hashTable {
+        result._unsafeInsertNew(self.uncheckedElement(at: bucket))
+      }
+    }
+    _storage = result._storage
+  }
+
+  @inlinable
+  internal mutating func copy() {
+    let newStorage = _SetStorage<Element>.copy(original: _storage)
+    _sanityCheck(newStorage._scale == _storage._scale)
+    _sanityCheck(newStorage._age == _storage._age)
+    _sanityCheck(newStorage._seed == _storage._seed)
     let result = _NativeSet(newStorage)
     if count > 0 {
-      if rehash {
-        for bucket in hashTable {
-          result._unsafeInsertNew(self.uncheckedElement(at: bucket))
-        }
-      } else {
-        result.hashTable.copyContents(of: hashTable)
-        result._storage._count = self.count
-        for bucket in hashTable {
-          let element = uncheckedElement(at: bucket)
-          result.uncheckedInitialize(at: bucket, to: element)
-        }
+      result.hashTable.copyContents(of: hashTable)
+      result._storage._count = self.count
+      for bucket in hashTable {
+        let element = uncheckedElement(at: bucket)
+        result.uncheckedInitialize(at: bucket, to: element)
       }
     }
     _storage = result._storage
-    return rehash
   }
 
   /// Ensure storage of self is uniquely held and can hold at least `capacity`
@@ -201,10 +212,15 @@ extension _NativeSet { // ensureUnique
     if _fastPath(capacity <= self.capacity && isUnique) {
       return false
     }
-    guard isUnique else {
-      return copy(capacity: capacity)
+    if isUnique {
+      resize(capacity: capacity)
+      return true
+    }
+    if capacity <= self.capacity {
+      copy()
+      return false
     }
-    resize(capacity: capacity)
+    copyAndResize(capacity: capacity)
     return true
   }
 

stdlib/public/core/SetStorage.swift

@@ -285,19 +285,22 @@ final internal class _SetStorage<Element: Hashable>
 extension _SetStorage {
   @usableFromInline
   @_effects(releasenone)
-  internal static func reallocate(
+  internal static func copy(original: _RawSetStorage) -> _SetStorage {
+    return .allocate(
+      scale: original._scale,
+      age: original._age,
+      seed: original._seed)
+  }
+
+  @usableFromInline
+  @_effects(releasenone)
+  static internal func resize(
     original: _RawSetStorage,
-    capacity: Int
-  ) -> (storage: _SetStorage, rehash: Bool) {
-    _sanityCheck(capacity >= original._count)
+    capacity: Int,
+    move: Bool
+  ) -> _SetStorage {
     let scale = _HashTable.scale(forCapacity: capacity)
-    let rehash = (scale != original._scale)
-    if rehash {
-      return .allocate(scale: scale, age: nil, seed: nil)
-    }
-    return (
-      .allocate(scale: scale, age: original._age, seed: original._seed),
-      rehash)
+    return allocate(scale: scale, age: nil, seed: nil)
   }
 
   @usableFromInline