🙅🏼‍♂️ [stdlib] Add method for accessing full storage capacity of array

stdlib/public/core/Array.swift

@@ -1373,8 +1373,8 @@ extension Array {
   ///   the requested number of elements.
   ///
   /// - Parameters:
-  ///   - _unsafeUninitializedCapacity: The number of elements to allocate
-  ///     space for in the new array.
+  ///   - 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:
@@ -1385,26 +1385,26 @@ extension Array {
   ///         elements you initialize.
   @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)
+      Array._allocateUninitialized(unsafeUninitializedCapacity)
 
     var initializedCount = 0
     defer {
       // Update self.count even if initializer throws an error.
       _precondition(
-        initializedCount <= _unsafeUninitializedCapacity,
+        initializedCount <= unsafeUninitializedCapacity,
         "Initialized count set to greater than specified capacity."
       )
       self._buffer.count = initializedCount
     }
     var buffer = UnsafeMutableBufferPointer<Element>(
-      start: firstElementAddress, count: _unsafeUninitializedCapacity)
+      start: firstElementAddress, count: unsafeUninitializedCapacity)
     try initializer(&buffer, &initializedCount)
   }
 
@@ -1553,6 +1553,60 @@ extension Array {
     it._position = endIndex
     return (it,buffer.index(buffer.startIndex, offsetBy: self.count))
   }
+
+  /// Calls the given closure with a pointer to the full capacity of the
+  /// array's mutable contiguous storage.
+  ///
+  /// - Parameters:
+  ///   - capacity: The minimum capacity to reserve for the array. `capacity`
+  ///     must be greater than or equal to the array's current `count`.
+  ///   - body: A closure that can modify or deinitialize existing elements or
+  ///     initialize new elements.
+  ///   - Parameters:
+  ///     - buffer: An unsafe mutable buffer of the array's full storage,
+  ///       including any uninitialized capacity after the initialized
+  ///       elements. Only the elements in `buffer[0..<initializedCount]` are
+  ///       initialized. `buffer` covers the memory for exactly the number of
+  ///       elements specified in the `capacity` parameter.
+  ///     - initializedCount: The count of the array's initialized elements. If
+  ///       you initialize or deinitialize any elements inside `body`, update
+  ///       `initializedCount` with the new count for the array.
+  /// - Returns: The return value, if any, of the `body` closure parameter.
+  @inlinable
+  public mutating func withUnsafeMutableBufferPointerToStorage<R>(
+    capacity: Int,
+    _ body: (
+      _ buffer: inout UnsafeMutableBufferPointer<Element>,
+      _ initializedCount: inout Int
+    ) throws -> R
+  ) rethrows -> R {
+    // Ensure unique storage and requested capacity
+    _precondition(capacity >= self.count)
+    reserveCapacity(capacity)
+
+    var initializedCount = self.count
+
+    // Ensure that body can't invalidate the storage or its bounds by
+    // moving self into a temporary working array.
+    // See further notes above in withUnsafeMutableBuffer
+    var work = Array()
+    (work, self) = (self, work)
+
+    // Create an UnsafeBufferPointer over the full capacity of `work`
+    // that we can pass to `body`.
+    let pointer = work._buffer.firstElementAddress
+    var inoutBufferPointer = UnsafeMutableBufferPointer(
+      start: pointer, count: capacity)
+
+    // Put the working array back before returning and update the count.
+    defer {
+      (work, self) = (self, work)
+      _buffer.count = initializedCount
+    }
+
+    // Invoke the body.
+    return try body(&inoutBufferPointer, &initializedCount)
+  }
 }
 
 extension Array {

test/stdlib/Inputs/CommonArrayTests.gyb

@@ -357,14 +357,14 @@ ${Suite}.test("${ArrayType}/reserveCapacity") {
 }
 
 //===----------------------------------------------------------------------===//
-// 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
@@ -385,7 +385,15 @@ extension Collection {
   }
 }
 
-${Suite}.test("${ArrayType}/init(_unsafeUninitializedCapacity:...:)") {
+final class InstanceCountedClass {
+  static var instanceCounter = 0
+
+  init() { InstanceCountedClass.instanceCounter += 1 }
+  deinit { InstanceCountedClass.instanceCounter -= 1 }
+}
+enum E: Error { case error }
+
+${Suite}.test("${ArrayType}/init(unsafeUninitializedCapacity:...:)") {
   var a = ${ArrayType}(0..<300)
   let p = a.stablyPartitioned(by: { $0 % 2 == 0 })
   expectEqualSequence(
@@ -394,17 +402,59 @@ ${Suite}.test("${ArrayType}/init(_unsafeUninitializedCapacity:...:)") {
   )
 }
 
-${Suite}.test("${ArrayType}/init(_unsafeUninitializedCapacity:...:)/throwing") {
-  final class InstanceCountedClass {
-    static var instanceCounter = 0
+${Suite}.test("${ArrayType}/init(unsafeUninitializedCapacity:...:)/throwing") {
+  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)
 
-    init() { InstanceCountedClass.instanceCounter += 1 }
-    deinit { InstanceCountedClass.instanceCounter -= 1 }
+    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.
+    expectUnreachable()
+  } catch {}
+  expectEqual(0, InstanceCountedClass.instanceCounter)
+}
+
+${Suite}.test("${ArrayType}/withUnsafeMutableBufferPointerToStorage") {
+  var a = Array<Int>()
+  let result: Int = a.withUnsafeMutableBufferPointerToStorage(capacity: 10) {
+    buffer, initializedCount in
+    expectEqual(10, buffer.count)
+    expectEqual(0, initializedCount)
+
+    for i in 0..<5 {
+      buffer[i] = i
+    }
+    initializedCount = 5
+    return buffer[0..<5].reduce(0, +)
   }
-  enum E: Error { case error }
+  expectEqual(10, a.reduce(0, +))
+  expectEqual(5, a.count)
+}
 
+${Suite}.test("${ArrayType}/withUnsafeMutableBufferPointerToStorage/throwing") {
   do {
-    var a = Array<InstanceCountedClass>(_unsafeUninitializedCapacity: 10) { buffer, c in
+    var a = Array<InstanceCountedClass>()
+    a.withUnsafeMutableBufferPointerToStorage(capacity: 10) { buffer, c in
       let p = buffer.baseAddress!
       for i in 0..<5 {
         (p + i).initialize(to: InstanceCountedClass())
@@ -416,8 +466,7 @@ ${Suite}.test("${ArrayType}/init(_unsafeUninitializedCapacity:...:)/throwing") {
 
     a = []
     expectEqual(0, InstanceCountedClass.instanceCounter)
-
-    a = try Array(_unsafeUninitializedCapacity: 10) { buffer, c in
+    try a.withUnsafeMutableBufferPointerToStorage(capacity: 10) { buffer, c in
       let p = buffer.baseAddress!
       for i in 0..<5 {
         (p + i).initialize(to: InstanceCountedClass())
@@ -426,10 +475,10 @@ ${Suite}.test("${ArrayType}/init(_unsafeUninitializedCapacity:...:)/throwing") {
       throw E.error
     }
 
-    // The throw above should prevent reaching here, which should mean the
+    // The throw above should prevent reaching here, which means that 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)
 }