[stdlib] Implement _copyContents on internal Array types

`_copyContents(initializing:)` is a core method of Sequence, and it is used surprisingly often to copy stuff out of sequences. Array’s internal types currently have explicit implementations of it that trap (to prevent a performance bug due to the default iterator-based implementation. This has proved a bad idea, as not all code paths that end up calling `_copyContents` have actually been expunged — so we replaced a performance bug with a catastrophic correctness bug. 😥

Rather than trying to play whack-a-mole with code paths that end up in `_copyContents`, replace the traps with (relatively) efficient implementations, based on the ancient `_copyContents(subRange:initializing)` methods that have already been there all this time.

This resolves https://bugs.swift.org/browse/SR-14663.

I expect specialization will make this fix deploy back to earlier OSes in most (but unfortunately not all) cases.

Author: lorentey

stdlib/public/core/Array.swift

@@ -1865,6 +1865,22 @@ extension Array {
   }
 }
 
+#if INTERNAL_CHECKS_ENABLED
+extension Array {
+  // This allows us to test the `_copyContents` implementation in
+  // `_ArrayBuffer`. (It's like `_copyToContiguousArray` but it always makes a
+  // copy.)
+  @_alwaysEmitIntoClient
+  public func _copyToNewArray() -> [Element] {
+    Array(unsafeUninitializedCapacity: self.count) { buffer, count in
+      var (it, c) = self._buffer._copyContents(initializing: buffer)
+      _precondition(it.next() == nil)
+      count = c
+    }
+  }
+}
+#endif
+
 #if _runtime(_ObjC)
 // We isolate the bridging of the Cocoa Array -> Swift Array here so that
 // in the future, we can eagerly bridge the Cocoa array. We need this function

stdlib/public/core/ArrayBuffer.swift

@@ -311,12 +311,20 @@ extension _ArrayBuffer {
     return UnsafeMutableRawPointer(result).assumingMemoryBound(to: Element.self)
   }
 
-  public __consuming func _copyContents(
+  @inlinable
+  internal __consuming func _copyContents(
     initializing buffer: UnsafeMutableBufferPointer<Element>
-  ) -> (Iterator,UnsafeMutableBufferPointer<Element>.Index) {
-    // This customization point is not implemented for internal types.
-    // Accidentally calling it would be a catastrophic performance bug.
-    fatalError("unsupported")
+  ) -> (Iterator, UnsafeMutableBufferPointer<Element>.Index) {
+    if _fastPath(_isNative) {
+      let (_, c) = _native._copyContents(initializing: buffer)
+      return (IndexingIterator(_elements: self, _position: c), c)
+    }
+    guard buffer.count > 0 else { return (makeIterator(), 0) }
+    let ptr = UnsafeMutableRawPointer(buffer.baseAddress)?
+      .assumingMemoryBound(to: AnyObject.self)
+    let (_, c) = _nonNative._copyContents(
+      initializing: UnsafeMutableBufferPointer(start: ptr, count: buffer.count))
+    return (IndexingIterator(_elements: self, _position: c), c)
   }
 
   /// Returns a `_SliceBuffer` containing the given sub-range of elements in

stdlib/public/core/ArraySlice.swift

@@ -1520,3 +1520,19 @@ extension ArraySlice {
 
 extension ArraySlice: Sendable, UnsafeSendable
   where Element: Sendable { }
+
+#if INTERNAL_CHECKS_ENABLED
+extension ArraySlice {
+  // This allows us to test the `_copyContents` implementation in
+  // `_SliceBuffer`. (It's like `_copyToContiguousArray` but it always makes a
+  // copy.)
+  @_alwaysEmitIntoClient
+  public func _copyToNewArray() -> [Element] {
+    Array(unsafeUninitializedCapacity: self.count) { buffer, count in
+      var (it, c) = self._buffer._copyContents(initializing: buffer)
+      _precondition(it.next() == nil)
+      count = c
+    }
+  }
+}
+#endif

stdlib/public/core/CocoaArray.swift

@@ -146,5 +146,16 @@ internal struct _CocoaArrayWrapper: RandomAccessCollection {
     }
     return result
   }
+
+  @_alwaysEmitIntoClient
+  internal __consuming func _copyContents(
+    initializing buffer: UnsafeMutableBufferPointer<Element>
+  ) -> (Iterator, UnsafeMutableBufferPointer<Element>.Index) {
+    guard buffer.count > 0 else { return (makeIterator(), 0) }
+    let start = buffer.baseAddress!
+    let c = Swift.min(self.count, buffer.count)
+    let end = _copyContents(subRange: 0 ..< c, initializing: start)
+    return (IndexingIterator(_elements: self, _position: c), c)
+  }
 }
 #endif

stdlib/public/core/ContiguousArrayBuffer.swift

@@ -644,12 +644,17 @@ internal struct _ContiguousArrayBuffer<Element>: _ArrayBufferProtocol {
     return target + initializedCount
   }
 
-  public __consuming func _copyContents(
+  @inlinable
+  internal __consuming func _copyContents(
     initializing buffer: UnsafeMutableBufferPointer<Element>
-  ) -> (Iterator,UnsafeMutableBufferPointer<Element>.Index) {
-    // This customization point is not implemented for internal types.
-    // Accidentally calling it would be a catastrophic performance bug.
-    fatalError("unsupported")
+  ) -> (Iterator, UnsafeMutableBufferPointer<Element>.Index) {
+    guard buffer.count > 0 else { return (makeIterator(), 0) }
+    let c = Swift.min(self.count, buffer.count)
+    buffer.baseAddress!.initialize(
+      from: firstElementAddress,
+      count: c)
+    _fixLifetime(owner)
+    return (IndexingIterator(_elements: self, _position: c), c)
   }
 
   /// Returns a `_SliceBuffer` containing the given `bounds` of values

stdlib/public/core/SliceBuffer.swift

@@ -245,12 +245,18 @@ internal struct _SliceBuffer<Element>
     return target + c
   }
 
-  public __consuming func _copyContents(
+  @inlinable
+  internal __consuming func _copyContents(
     initializing buffer: UnsafeMutableBufferPointer<Element>
-  ) -> (Iterator,UnsafeMutableBufferPointer<Element>.Index) {
-    // This customization point is not implemented for internal types.
-    // Accidentally calling it would be a catastrophic performance bug.
-    fatalError("unsupported")
+  ) -> (Iterator, UnsafeMutableBufferPointer<Element>.Index) {
+    _invariantCheck()
+    guard buffer.count > 0 else { return (makeIterator(), 0) }
+    let c = Swift.min(self.count, buffer.count)
+    buffer.baseAddress!.initialize(
+      from: firstElementAddress,
+      count: c)
+    _fixLifetime(owner)
+    return (IndexingIterator(_elements: self, _position: startIndex + c), c)
   }
 
   /// True, if the array is native and does not need a deferred type check.

test/stdlib/ArrayBridge.swift.gyb

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Copyright (c) 2014 - 2021 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
@@ -471,6 +471,8 @@ tests.test("testMutableArray") {
 }
 
 tests.test("rdar://problem/27905230") {
+  // Casting an NSArray to Array<Any> would trap because of an erroneous
+  // precondition.
   let dict = RDar27905230.mutableDictionaryOfMutableLists()!
   let arr = dict["list"]!
   expectEqual(arr[0] as! NSNull, NSNull())
@@ -482,4 +484,69 @@ tests.test("rdar://problem/27905230") {
   expectEqual(arr[5] as! Date, Date(timeIntervalSince1970: 0))
 }
 
+tests.test("verbatimBridged/Base/withUnsafeBufferPointer") {
+  let a = NSArray(array: [Base(0), Base(1), Base(2), Base(3)])
+  let b = a as! [Base]
+  let success: Bool = b.withUnsafeBufferPointer { buffer in
+    expectEqual(buffer.count, 4)
+    guard buffer.count == 4 else { return false }
+    expectEqual(buffer[0].value, 0)
+    expectEqual(buffer[1].value, 1)
+    expectEqual(buffer[2].value, 2)
+    expectEqual(buffer[3].value, 3)
+    return true
+  }
+  expectTrue(success)
+}
+
+// https://bugs.swift.org/browse/SR-14663
+tests.test("verbatimBridged/AnyObject/withUnsafeBufferPointer") {
+  let a = NSArray(array: [Base(0), Base(1), Base(2), Base(3)])
+  let b = a as [AnyObject]
+  let success: Bool = b.withUnsafeBufferPointer { buffer in
+    expectEqual(buffer.count, 4)
+    guard buffer.count == 4 else { return false }
+    expectEqual((buffer[0] as? Base)?.value, 0)
+    expectEqual((buffer[1] as? Base)?.value, 1)
+    expectEqual((buffer[2] as? Base)?.value, 2)
+    expectEqual((buffer[3] as? Base)?.value, 3)
+    return true
+  }
+  expectTrue(success)
+}
+
+tests.test("verbatimBridged/Base/withUnsafeMutableBufferPointer") {
+  let a = NSArray(array: [Base(0), Base(1), Base(2), Base(3)])
+  var b = a as! [Base]
+  let success: Bool = b.withUnsafeMutableBufferPointer { buffer in
+    expectEqual(buffer.count, 4)
+    guard buffer.count == 4 else { return false }
+    expectEqual(buffer[0].value, 0)
+    expectEqual(buffer[1].value, 1)
+    expectEqual(buffer[2].value, 2)
+    expectEqual(buffer[3].value, 3)
+    buffer[0] = Base(4)
+    return true
+  }
+  expectTrue(success)
+  expectEqual(b[0].value, 4)
+}
+
+tests.test("verbatimBridged/AnyObject/withUnsafeMutableBufferPointer") {
+  let a = NSArray(array: [Base(0), Base(1), Base(2), Base(3)])
+  var b = a as [AnyObject]
+  let success: Bool = b.withUnsafeMutableBufferPointer { buffer in
+    expectEqual(buffer.count, 4)
+    guard buffer.count == 4 else { return false }
+    expectEqual((buffer[0] as? Base)?.value, 0)
+    expectEqual((buffer[1] as? Base)?.value, 1)
+    expectEqual((buffer[2] as? Base)?.value, 2)
+    expectEqual((buffer[3] as? Base)?.value, 3)
+    buffer[0] = Base(4)
+    return true
+  }
+  expectTrue(success)
+  expectEqual((b[0] as? Base)?.value, 4)
+}
+
 runAllTests()

test/stdlib/ArrayBuffer_CopyContents.swift

@@ -0,0 +1,79 @@
+//===--- ArrayBuffer_CopyContents.swift -----------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2021 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+// REQUIRES: swift_stdlib_asserts
+
+import Foundation
+import StdlibUnittest
+
+let suite = TestSuite("ArrayBuffer_CopyContents")
+defer { runAllTests() }
+
+
+var trackedCount = 0
+var nextBaseSerialNumber = 0
+
+/// A type that will be bridged verbatim to Objective-C
+class Thing: NSObject {
+  var value: Int
+  var serialNumber: Int
+
+  func foo() { }
+
+  required init(_ value: Int) {
+    trackedCount += 1
+    nextBaseSerialNumber += 1
+    serialNumber = nextBaseSerialNumber
+    self.value = value
+  }
+
+  deinit {
+    assert(serialNumber > 0, "double destruction!")
+    trackedCount -= 1
+    serialNumber = -serialNumber
+  }
+
+  override func isEqual(_ other: Any?) -> Bool {
+    return (other as? Thing)?.value == self.value
+  }
+
+  override var hash: Int { value }
+}
+
+
+suite.test("nativeArray/_copyContents") {
+  let array = [Thing(0), Thing(1), Thing(2), Thing(3)]
+  expectEqualSequence(array._copyToNewArray(), array)
+}
+
+suite.test("nativeArraySlice/_copyContents") {
+  let array = (0 ..< 100).map { Thing($0) }
+  expectEqualSequence(
+    array[20 ..< 30]._copyToNewArray(),
+    (20 ..< 30).map { Thing($0) })
+}
+
+suite.test("bridgedArray/_copyContents") {
+  let array = NSArray(array: (0 ..< 5).map { Thing($0) }) as! [Thing]
+  expectEqualSequence(
+    array._copyToNewArray(),
+    (0 ..< 5).map { Thing($0) })
+}
+
+suite.test("bridgedArraySlice/_copyContents") {
+  let array = NSArray(array: (0 ..< 100).map { Thing($0) }) as! [Thing]
+  expectEqualSequence(
+    array[20 ..< 30]._copyToNewArray(),
+    (20 ..< 30).map { Thing($0) })
+}