[stdlib] Make unsafe array initializer public (#23134)

[stdlib] Make unsafe array initializer public

This implements SE-0245. The public versions of this initializer call
into the existing, underscored version, which avoids the need for
availability constraints.

Author: natecook1000

CHANGELOG.md

@@ -99,6 +99,11 @@ Swift 5.1
   foo(0) // prints "Any" in Swift < 5.1, "T" in Swift 5.1
   ```
 
+* [SE-0245][]:
+
+  `Array` and `ContiguousArray` now have `init(unsafeUninitializedCapacity:initializingWith:)`,
+  which provides access to the array's uninitialized storage.
+
 **Add new entries to the top of this section, not here!**
 
 Swift 5.0
@@ -7506,6 +7511,7 @@ Swift 1.0
 [SE-0228]: <https://github.com/apple/swift-evolution/blob/master/proposals/0228-fix-expressiblebystringinterpolation.md>
 [SE-0230]: <https://github.com/apple/swift-evolution/blob/master/proposals/0230-flatten-optional-try.md>
 [SE-0235]: <https://github.com/apple/swift-evolution/blob/master/proposals/0235-add-result.md>
+[SE-0245]: <https://github.com/apple/swift-evolution/blob/master/proposals/0245-array-uninitialized-initializer.md>
 
 [SR-106]: <https://bugs.swift.org/browse/SR-106>
 [SR-419]: <https://bugs.swift.org/browse/SR-419>

stdlib/public/core/Array.swift

@@ -1366,21 +1366,53 @@ extension Array {
 }
 
 extension Array {
+  /// Implementation for Array(unsafeUninitializedCapacity:initializingWith:)
+  /// and ContiguousArray(unsafeUninitializedCapacity:initializingWith:)
+  @inlinable
+  internal 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
+    var buffer = UnsafeMutableBufferPointer<Element>(
+      start: firstElementAddress, count: _unsafeUninitializedCapacity)
+    defer {
+      // Update self.count even if initializer throws an error.
+      _precondition(
+        initializedCount <= _unsafeUninitializedCapacity,
+        "Initialized count set to greater than specified capacity."
+      )
+      _precondition(
+        buffer.baseAddress == firstElementAddress,
+        "Can't reassign buffer in Array(unsafeUninitializedCapacity:initializingWith:)"
+      )
+      self._buffer.count = initializedCount
+    }
+    try initializer(&buffer, &initializedCount)
+  }
+
   /// 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.
+  /// `buffer[initializedCount...]` must be uninitialized. This postcondition
+  /// must hold even if the `initializer` closure throws an error.
   ///
   /// - 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
+  ///   - 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.
@@ -1390,31 +1422,18 @@ extension Array {
   ///       - initializedCount: The count of initialized elements in the array,
   ///         which begins as zero. Set `initializedCount` to the number of
   ///         elements you initialize.
-  @inlinable
+  @_alwaysEmitIntoClient @inlinable
   public init(
-    _unsafeUninitializedCapacity: Int,
+    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)
+    self = try Array(
+      _unsafeUninitializedCapacity: unsafeUninitializedCapacity,
+      initializingWith: initializer)
   }
-  
+
   /// Calls a closure with a pointer to the array's contiguous storage.
   ///
   /// Often, the optimizer can eliminate bounds checks within an array

stdlib/public/core/ContiguousArray.swift

@@ -985,6 +985,43 @@ extension ContiguousArray {
 }
 
 extension ContiguousArray {
+  /// 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. This postcondition
+  /// must hold even if the `initializer` closure throws an error.
+  ///
+  /// - 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.
+  @_alwaysEmitIntoClient @inlinable
+  public init(
+    unsafeUninitializedCapacity: Int,
+    initializingWith initializer: (
+      _ buffer: inout UnsafeMutableBufferPointer<Element>,
+      _ initializedCount: inout Int) throws -> Void
+  ) rethrows {
+    self = try ContiguousArray(Array(
+      _unsafeUninitializedCapacity: unsafeUninitializedCapacity,
+      initializingWith: initializer))
+  }
+
   /// 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

@@ -462,10 +462,6 @@ ${Suite}.test("${ArrayType}/withUnsafeMutableBytes")
   expectEqual(10, b[0])
 }
 
-// FIXME: Implement these changes for ArraySlice and ContiguousArray
-
-%if ArrayType == 'Array':
-
 //===----------------------------------------------------------------------===//
 // reserveCapacity semantics
 //===----------------------------------------------------------------------===//
@@ -478,36 +474,39 @@ ${Suite}.test("${ArrayType}/reserveCapacity") {
   }
 }
 
+%if ArrayType in ['Array', 'ContiguousArray']:
+
 //===----------------------------------------------------------------------===//
-// init(_unsafeUninitializedCapacity:initializingWith:)
+// init(unsafeUninitializedCapacity:initializingWith:)
 //===----------------------------------------------------------------------===//
 
 extension Collection {
   func stablyPartitioned(
     by belongsInFirstPartition: (Element) -> Bool
   ) -> ${ArrayType}<Element> {
-    let result = ${ArrayType}<Element>(_unsafeUninitializedCapacity: self.count) {
+    let result = ${ArrayType}<Element>(unsafeUninitializedCapacity: self.count) {
       buffer, initializedCount in
-      var lowIndex = 0
-      var highIndex = buffer.count
+      var low = buffer.baseAddress!
+      var high = low + buffer.count
       for element in self {
         if belongsInFirstPartition(element) {
-          buffer[lowIndex] = element
-          lowIndex += 1
+          low.initialize(to: element)
+          low += 1
         } else {
-          highIndex -= 1
-          buffer[highIndex] = element
+          high -= 1
+          high.initialize(to: element)
         }
       }
-      
+    
+      let highIndex = high - buffer.baseAddress!
       buffer[highIndex...].reverse()
       initializedCount = buffer.count
     }
     return result
   }
 }
 
-${Suite}.test("${ArrayType}/init(_unsafeUninitializedCapacity:...:)") {
+${Suite}.test("${ArrayType}/init(unsafeUninitializedCapacity:...:)") {
   var a = ${ArrayType}(0..<300)
   let p = a.stablyPartitioned(by: { $0 % 2 == 0 })
   expectEqualSequence(
@@ -516,7 +515,7 @@ ${Suite}.test("${ArrayType}/init(_unsafeUninitializedCapacity:...:)") {
   )
 }
 
-${Suite}.test("${ArrayType}/init(_unsafeUninitializedCapacity:...:)/throwing") {
+${Suite}.test("${ArrayType}/init(unsafeUninitializedCapacity:...:)/throwing") {
   final class InstanceCountedClass {
     static var instanceCounter = 0
 
@@ -526,7 +525,8 @@ ${Suite}.test("${ArrayType}/init(_unsafeUninitializedCapacity:...:)/throwing") {
   enum E: Error { case error }
 
   do {
-    var a = Array<InstanceCountedClass>(_unsafeUninitializedCapacity: 10) { buffer, c in
+    var a = ${ArrayType}<InstanceCountedClass>(unsafeUninitializedCapacity: 10) {
+        buffer, c in
       let p = buffer.baseAddress!
       for i in 0..<5 {
         (p + i).initialize(to: InstanceCountedClass())
@@ -539,7 +539,7 @@ ${Suite}.test("${ArrayType}/init(_unsafeUninitializedCapacity:...:)/throwing") {
     a = []
     expectEqual(0, InstanceCountedClass.instanceCounter)
 
-    a = try Array(_unsafeUninitializedCapacity: 10) { buffer, c in
+    a = try ${ArrayType}(unsafeUninitializedCapacity: 10) { buffer, c in
       let p = buffer.baseAddress!
       for i in 0..<5 {
         (p + i).initialize(to: InstanceCountedClass())
@@ -548,14 +548,37 @@ ${Suite}.test("${ArrayType}/init(_unsafeUninitializedCapacity:...:)/throwing") {
       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)
+    expectUnreachable()
   } catch {}
   expectEqual(0, InstanceCountedClass.instanceCounter)
 }
 
+${Suite}.test("${ArrayType}/init(unsafeUninitializedCapacity:...:)/noCopy") {
+  var storageAddress: UnsafeMutablePointer<Int>?
+  var array = ${ArrayType}<Int>(unsafeUninitializedCapacity: 20) { buffer, _ in
+    storageAddress = buffer.baseAddress
+  }
+  array.withUnsafeMutableBufferPointer { buffer in
+    expectEqual(storageAddress, buffer.baseAddress)
+  }
+}
+
+${Suite}.test("${ArrayType}/init(unsafeUninitializedCapacity:...:)/validCount") {
+  expectCrashLater()
+  let array = ${ArrayType}<Int>(unsafeUninitializedCapacity: 10) { buffer, c in
+    for i in 0..<10 { buffer[i] = i }
+    c = 20
+  }
+}
+
+${Suite}.test("${ArrayType}/init(unsafeUninitializedCapacity:...:)/reassignBuffer") {
+  expectCrashLater()
+  let otherBuffer = UnsafeMutableBufferPointer<Int>.allocate(capacity: 1)
+  let array = ${ArrayType}<Int>(unsafeUninitializedCapacity: 10) { buffer, _ in
+    buffer = otherBuffer
+  }
+}
+
 %end
 
 //===---