diff --git a/stdlib/public/core/Array.swift b/stdlib/public/core/Array.swift index 7985db4ad810d..0a4de7c9ddfe7 100644 --- a/stdlib/public/core/Array.swift +++ b/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, _ initializedCount: inout Int) throws -> Void ) rethrows { var firstElementAddress: UnsafeMutablePointer (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( - 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..( + capacity: Int, + _ body: ( + _ buffer: inout UnsafeMutableBufferPointer, + _ 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 { diff --git a/test/stdlib/Inputs/CommonArrayTests.gyb b/test/stdlib/Inputs/CommonArrayTests.gyb index ee8f2bf2550ee..315d4f7720d6c 100644 --- a/test/stdlib/Inputs/CommonArrayTests.gyb +++ b/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} { - let result = ${ArrayType}(_unsafeUninitializedCapacity: self.count) { + let result = ${ArrayType}(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(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() + 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(_unsafeUninitializedCapacity: 10) { buffer, c in + var a = Array() + 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) }