Add an array initializer with access to uninitialized storage (#17774)

* Add an Array initializer for using an uninitialized capacity.

* Add tests for reserveCapacity & init(_unsafeUninitializedCapacity:...:)

Author: natecook1000

stdlib/public/core/Array.swift

@@ -1353,6 +1353,55 @@ extension Array {
 }
 
 extension Array {
+  /// Creates an array with the specified capacity, then calls the given
+  /// closure with a buffer covering the array's uninitialized memory.
+  ///
+  /// Inside the closure, set the `initializedCount` parameter to the number of
+  /// elements that are initialized by the closure. The memory in the range
+  /// `buffer[0..<initializedCount]` must be initialized at the end of the
+  /// closure's execution, and the memory in the range
+  /// `buffer[initializedCount...]` must be uninitialized.
+  ///
+  /// - Note: While the resulting array may have a capacity larger than the
+  ///   requested amount, the buffer passed to the closure will cover exactly
+  ///   the requested number of elements.
+  ///
+  /// - Parameters:
+  ///   - _unsafeUninitializedCapacity: The number of elements to allocate
+  ///     space for in the new array.
+  ///   - initializer: A closure that initializes elements and sets the count
+  ///     of the new array.
+  ///     - Parameters:
+  ///       - buffer: A buffer covering uninitialized memory with room for the
+  ///         specified number of of elements.
+  ///       - initializedCount: The count of initialized elements in the array,
+  ///         which begins as zero. Set `initializedCount` to the number of
+  ///         elements you initialize.
+  @inlinable
+  public init(
+    _unsafeUninitializedCapacity: Int,
+    initializingWith initializer: (
+      _ buffer: inout UnsafeMutableBufferPointer<Element>,
+      _ initializedCount: inout Int) throws -> Void
+  ) rethrows {
+    var firstElementAddress: UnsafeMutablePointer<Element>
+    (self, firstElementAddress) =
+      Array._allocateUninitialized(_unsafeUninitializedCapacity)
+    
+    var initializedCount = 0
+    defer {
+      // Update self.count even if initializer throws an error.
+      _precondition(
+        initializedCount <= _unsafeUninitializedCapacity,
+        "Initialized count set to greater than specified capacity."
+      )
+      self._buffer.count = initializedCount
+    }
+    var buffer = UnsafeMutableBufferPointer<Element>(
+      start: firstElementAddress, count: _unsafeUninitializedCapacity)
+    try initializer(&buffer, &initializedCount)
+  }
+  
   /// Calls a closure with a pointer to the array's contiguous storage.
   ///
   /// Often, the optimizer can eliminate bounds checks within an array

test/stdlib/Inputs/CommonArrayTests.gyb

@@ -340,6 +340,102 @@ ${Suite}.test("${ArrayType}/withUnsafeMutableBytes")
   expectEqual(10, b[0])
 }
 
+// FIXME: Implement these changes for ArraySlice and ContiguousArray
+
+%if ArrayType == 'Array':
+
+//===----------------------------------------------------------------------===//
+// reserveCapacity semantics
+//===----------------------------------------------------------------------===//
+
+${Suite}.test("${ArrayType}/reserveCapacity") {
+  var a = ${ArrayType}<Int>()
+  for capacity in stride(from: 0, to: 1000, by: 53) {
+    a.reserveCapacity(capacity)
+    expectLE(capacity, a.capacity)
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// init(_unsafeUninitializedCapacity:initializingWith:)
+//===----------------------------------------------------------------------===//
+
+extension Collection {
+  func stablyPartitioned(
+    by belongsInFirstPartition: (Element) -> Bool
+  ) -> ${ArrayType}<Element> {
+    let result = ${ArrayType}<Element>(_unsafeUninitializedCapacity: self.count) {
+      buffer, initializedCount in
+      var lowIndex = 0
+      var highIndex = buffer.count
+      for element in self {
+        if belongsInFirstPartition(element) {
+          buffer[lowIndex] = element
+          lowIndex += 1
+        } else {
+          highIndex -= 1
+          buffer[highIndex] = element
+        }
+      }
+      
+      buffer[highIndex...].reverse()
+      initializedCount = buffer.count
+    }
+    return result
+  }
+}
+
+${Suite}.test("${ArrayType}/init(_unsafeUninitializedCapacity:...:)") {
+  var a = ${ArrayType}(0..<300)
+  let p = a.stablyPartitioned(by: { $0 % 2 == 0 })
+  expectEqualSequence(
+    [stride(from: 0, to: 300, by: 2), stride(from: 1, to: 300, by: 2)].joined(),
+    p
+  )
+}
+
+${Suite}.test("${ArrayType}/init(_unsafeUninitializedCapacity:...:)/throwing") {
+  final class InstanceCountedClass {
+    static var instanceCounter = 0
+
+    init() { InstanceCountedClass.instanceCounter += 1 }
+    deinit { InstanceCountedClass.instanceCounter -= 1 }
+  }
+  enum E: Error { case error }
+
+  do {
+    var a = Array<InstanceCountedClass>(_unsafeUninitializedCapacity: 10) { buffer, c in
+      let p = buffer.baseAddress!
+      for i in 0..<5 {
+        (p + i).initialize(to: InstanceCountedClass())
+      }
+      c = 5
+    }
+    expectEqual(5, a.count)
+    expectEqual(5, InstanceCountedClass.instanceCounter)
+
+    a = []
+    expectEqual(0, InstanceCountedClass.instanceCounter)
+
+    a = try Array(_unsafeUninitializedCapacity: 10) { buffer, c in
+      let p = buffer.baseAddress!
+      for i in 0..<5 {
+        (p + i).initialize(to: InstanceCountedClass())
+      }
+      c = 5
+      throw E.error
+    }
+
+    // The throw above should prevent reaching here, which should mean the
+    // instances created in the closure should get deallocated before the final
+    // expectation outside the do/catch block.
+    expectTrue(false)
+  } catch {}
+  expectEqual(0, InstanceCountedClass.instanceCounter)
+}
+
+%end
+
 //===---
 // Check that iterators traverse a snapshot of the collection.
 //===---