[stdlib] Dictionary: Add unsafe bulk-loading initializer

Interface inspired by this Array pitch:
https://forums.swift.org/t/array-initializer-with-access-to-uninitialized-buffer/13689

Author: lorentey

stdlib/public/core/DictionaryBuilder.swift

@@ -42,3 +42,85 @@ struct _DictionaryBuilder<Key: Hashable, Value> {
     return Dictionary(_native: _target)
   }
 }
+
+extension Dictionary {
+  /// Creates a new dictionary with the specified capacity, then calls the given
+  /// closure to initialize its contents.
+  ///
+  /// Foundation uses this initializer to bridge the contents of an NSDictionary
+  /// instance without allocating a pair of intermediary buffers.  Pass the
+  /// required capacity and a closure that can intialize the dictionary's
+  /// elements. The closure must return `c`, the number of initialized elements
+  /// in both buffers, such that the elements in the range `0..<c` are
+  /// initialized and the elements in the range `c..<capacity` are
+  /// uninitialized. The resulting dictionary has a `count` equal to `c`.
+  ///
+  /// The closure must not add any duplicate keys to the keys buffer.  The
+  /// buffers passed to the closure are only valid for the duration of the call.
+  /// After the closure returns, this initializer moves all initialized elements
+  /// into their correct buckets.
+  ///
+  /// - Parameters:
+  ///   - capacity: The capacity of the new dictionary.
+  ///   - body: A closure that can initialize the dictionary's elements. This
+  ///     closure must return the count of the initialized elements, starting at
+  ///     the beginning of the buffer.
+  @inlinable
+  public // SPI(Foundation)
+  init(
+    _unsafeUninitializedCapacity capacity: Int,
+    initializingWith initializer: (
+      _ keys: UnsafeMutableBufferPointer<Key>,
+      _ values: UnsafeMutableBufferPointer<Value>,
+      _ initializedCount: inout Int
+    ) -> Void
+  ) {
+    let native = _NativeDictionary<Key, Value>(capacity: capacity)
+    let keys = UnsafeMutableBufferPointer<Key>(
+      start: native._keys,
+      count: capacity)
+    let values = UnsafeMutableBufferPointer<Value>(
+      start: native._values,
+      count: capacity)
+    var count = 0
+    initializer(keys, values, &count)
+    _precondition(count >= 0 && count <= capacity)
+
+    // Hash elements into the correct buckets.
+    var bucket = _HashTable.Bucket(offset: 0)
+    while bucket.offset < count {
+      if native.hashTable._isOccupied(bucket) {
+        // We've moved an element here in a previous iteration.
+        bucket.offset += 1
+        continue
+      }
+      let hashValue = native.hashValue(for: native._keys[bucket.offset])
+      let target: _HashTable.Bucket
+      if _isDebugAssertConfiguration() {
+        let (b, found) = native.find(
+          native._keys[bucket.offset],
+          hashValue: hashValue)
+        guard !found else {
+          _preconditionFailure("Duplicate keys found")
+        }
+        native.hashTable.insert(b)
+        target = b
+      } else {
+        target = native.hashTable.insertNew(hashValue: hashValue)
+      }
+      if target < bucket || target.offset >= count {
+        // Target bucket is uninitialized
+        native.moveEntry(from: bucket, to: target)
+        bucket.offset += 1
+      } else if target == bucket {
+        // Already in place
+        bucket.offset += 1
+      } else {
+        // Swap into place, then try again with the swapped-in value.
+        native.swapEntry(target, with: bucket)
+      }
+    }
+    native._storage._count = count
+    self = Dictionary(_native: native)
+  }
+}

stdlib/public/core/NativeDictionary.swift

@@ -619,11 +619,22 @@ extension _NativeDictionary: _HashTableDelegate {
   @inlinable
   @inline(__always)
   internal func moveEntry(from source: Bucket, to target: Bucket) {
+    _internalInvariant(hashTable.isValid(source))
+    _internalInvariant(hashTable.isValid(target))
     (_keys + target.offset)
       .moveInitialize(from: _keys + source.offset, count: 1)
     (_values + target.offset)
       .moveInitialize(from: _values + source.offset, count: 1)
   }
+
+  @inlinable
+  @inline(__always)
+  internal func swapEntry(_ left: Bucket, with right: Bucket) {
+    _internalInvariant(hashTable.isValid(left))
+    _internalInvariant(hashTable.isValid(right))
+    swap(&_keys[left.offset], &_keys[right.offset])
+    swap(&_values[left.offset], &_values[right.offset])
+  }
 }
 
 extension _NativeDictionary { // Deletion

validation-test/stdlib/Dictionary.swift

@@ -5547,6 +5547,32 @@ DictionaryTestSuite.test("IndexValidation.RemoveAt.AfterGrow") {
   d.remove(at: i)
 }
 
+DictionaryTestSuite.test("BulkLoadingInitializer") {
+  for c in [0, 1, 2, 3, 5, 10, 25, 150] {
+    let d1 = Dictionary<TestKeyTy, TestEquatableValueTy>(
+      _unsafeUninitializedCapacity: c
+    ) { keys, values, count in
+      let k = keys.baseAddress!
+      let v = values.baseAddress!
+      for i in 0 ..< c {
+        (k + i).initialize(to: TestKeyTy(i))
+        (v + i).initialize(to: TestEquatableValueTy(i))
+        count += 1
+      }
+    }
+
+    let d2 = Dictionary(
+      uniqueKeysWithValues: (0..<c).map {
+        (TestKeyTy($0), TestEquatableValueTy($0))
+      })
+
+    for i in 0 ..< c {
+      expectEqual(TestEquatableValueTy(i), d1[TestKeyTy(i)])
+    }
+    expectEqual(d1, d2)
+  }
+}
+
 DictionaryTestSuite.setUp {
 #if _runtime(_ObjC)
   // Exercise ARC's autoreleased return value optimization in Foundation.