Migrate _copyContents to be called from UnsafeMutableBufferPointer with checks for overrun.

Author: airspeedswift

stdlib/private/StdlibCollectionUnittest/CheckSequenceInstance.swift.gyb

@@ -86,16 +86,14 @@ public func checkSequence<
   _ = sequence._preprocessingPass { () -> Void in
     var count = 0
     for _ in sequence { count += 1 }
-    let buf = UnsafeMutablePointer<S.Iterator.Element>.allocate(capacity: count)
-    let end = sequence._copyContents(initializing: buf)
-    expectTrue(end == buf + count, "_copyContents returned the wrong value")
-    var j = expected.startIndex
-    for i in 0..<(end - buf) {
-      expectTrue(sameValue(expected[j], buf[i]))
-      j = expected.index(after: j)
-    }
-    buf.deinitialize(count: end - buf)
-    buf.deallocate(capacity: count)
+    let ptr = UnsafeMutablePointer<S.Iterator.Element>.allocate(capacity: count)
+    let buf = UnsafeMutableBufferPointer(start: ptr, count: count)
+    var remainders = sequence._copyContents(initializing: buf)
+    expectTrue(remainders == nil || remainders!.next() == nil, 
+      "_copyContents returned unwritten elements")
+    expectEqualSequence(expected, buf, ${trace}, sameValue: sameValue)
+    ptr.deinitialize(count: count)
+    ptr.deallocate(capacity: count)
   }
 
   // Test `_copyToContiguousArray()` if we can do so

stdlib/private/StdlibCollectionUnittest/LoggingWrappers.swift.gyb

@@ -215,7 +215,7 @@ public struct ${Self}<
 
   public typealias Iterator = LoggingIterator<Base.Iterator>
 
-  public func makeIterator() -> LoggingIterator<Base.Iterator> {
+  public func makeIterator() -> Iterator {
     Log.makeIterator[selfType] += 1
     return LoggingIterator(wrapping: base.makeIterator())
   }
@@ -321,10 +321,10 @@ public struct ${Self}<
 
   /// Copy a Sequence into an array.
   public func _copyContents(
-    initializing ptr: UnsafeMutablePointer<Base.Iterator.Element>
-  ) -> UnsafeMutablePointer<Base.Iterator.Element> {
+    initializing buf: UnsafeMutableBufferPointer<Iterator.Element>
+  ) -> Iterator? {
     Log._copyContents[selfType] += 1
-    return base._copyContents(initializing: ptr)
+    return base._copyContents(initializing: buf).map(Iterator.init)
   }
 
 %     if Kind in ['Collection', 'MutableCollection', 'RangeReplaceableCollection']:

stdlib/public/SDK/Foundation/Data.swift

@@ -551,7 +551,9 @@ public struct Data : ReferenceConvertible, Equatable, Hashable, RandomAccessColl
             }
 
             if replacementCount != 0 {
-                newElements._copyContents(initializing: bytes + start)
+                let buf = UnsafeMutableBufferPointer(start: bytes + start, count: numericCast(newElements.count))
+                // FIXME: is this guaranteed to return no residual elements?
+                let _ = newElements._copyContents(initializing: buf)
             }
         }
     }

stdlib/public/core/Arrays.swift.gyb

@@ -1279,7 +1279,7 @@ extension ${Self} : RangeReplaceableCollection, _ArrayProtocol {
 
   /// Adds the elements of a sequence to the end of the array.
   ///
-  /// Use this method to append the elements of a sequence to the end of an
+  /// Use this method to append the elements of a sequence to the end of this
   /// array. This example appends the elements of a `Range<Int>` instance
   /// to an array of integers.
   ///
@@ -1293,39 +1293,8 @@ extension ${Self} : RangeReplaceableCollection, _ArrayProtocol {
   /// - Complexity: O(*n*), where *n* is the length of the resulting array.
   public mutating func append<S : Sequence>(contentsOf newElements: S)
     where S.Iterator.Element == Element {
-    let oldCount = self.count
-    let capacity = self.capacity
-    let newCount = oldCount + newElements.underestimatedCount
 
-    if newCount > capacity {
-      self.reserveCapacity(
-        Swift.max(newCount, _growArrayCapacity(capacity)))
-      }
-    _buffer._arrayAppendSequence(newElements)
-  }
-
-  // An overload of `append(contentsOf:)` that uses the += that is optimized for
-  // collections.
-  // FIXME(ABI)#13 (Performance): remove this entrypoint.  The overload for `Sequence` should be
-  // made optimal for this case, too.
-  /// Adds the elements of a collection to the end of the array.
-  ///
-  /// Use this method to append the elements of a collection to the end of this
-  /// array. This example appends the elements of a `Range<Int>` instance
-  /// to an array of integers.
-  ///
-  ///     var numbers = [1, 2, 3, 4, 5]
-  ///     numbers.append(contentsOf: 10...15)
-  ///     print(numbers)
-  ///     // Prints "[1, 2, 3, 4, 5, 10, 11, 12, 13, 14, 15]"
-  ///
-  /// - Parameter newElements: The elements to append to the array.
-  ///
-  /// - Complexity: O(*n*), where *n* is the length of the resulting array.
-  public mutating func append<C : Collection>(contentsOf newElements: C)
-    where C.Iterator.Element == Element {
-
-    let newElementsCount = numericCast(newElements.count) as Int
+    let newElementsCount = newElements.underestimatedCount
 
     let oldCount = self.count
     let capacity = self.capacity
@@ -1338,8 +1307,13 @@ extension ${Self} : RangeReplaceableCollection, _ArrayProtocol {
       Swift.max(newCount, _growArrayCapacity(capacity))
       : newCount)
 
-    (self._buffer.firstElementAddress + oldCount).initialize(from: newElements)
-    self._buffer.count = newCount
+    let buf = UnsafeMutableBufferPointer(start: _buffer.firstElementAddress + oldCount, count: newElementsCount)
+    _buffer.count = newCount    
+    if let remainder = buf.initialize(from: newElements) {
+      // there were elements that didn't fit in the existing buffer,
+      // append them in slow sequence-only mode
+      _buffer._arrayAppendSequence(IteratorSequence(remainder))
+    }
   }
 
 %if Self == 'ArraySlice':
@@ -1624,23 +1598,22 @@ extension ${Self} {
     return try body(&inoutBufferPointer)
   }
 
-  @discardableResult
   public func _copyContents(
-    initializing ptr: UnsafeMutablePointer<Element>
-  ) -> UnsafeMutablePointer<Element> {
-    if let s = self._baseAddressIfContiguous {
-      let count = self.count
-      ptr.initialize(from: s, count: count)
+    initializing buf: UnsafeMutableBufferPointer<Element>
+  ) -> Iterator? {
+    // FIXME: decide if this really is a fatalError
+    guard var p = buf.baseAddress 
+      else { fatalError("Attempt to copy contents into nil buffer pointer") }
+    if let s = _baseAddressIfContiguous {
+      p.initialize(from: s, count: self.count)
       _fixLifetime(self._owner)
-      return ptr + count
     } else {
-      var p = ptr
       for x in self {
         p.initialize(to: x)
         p += 1
       }
-      return p
     }
+    return nil
   }
 }
 %end
@@ -1790,30 +1763,6 @@ public func += <
 >(lhs: inout ${Self}<S.Iterator.Element>, rhs: S) {
   lhs.append(contentsOf: rhs)
 }
-
-// FIXME(ABI)#17 : remove this entrypoint.  The functionality should be provided by
-// a `+=` operator on `RangeReplaceableCollection`.
-/// Appends the elements of a collection to ${a_Self}.
-///
-/// Use this operator to append the elements of a collection to the end of
-/// ${a_Self} with same `Element` type. This example appends
-/// the elements of a `Range<Int>` instance to an array of integers.
-///
-///     var numbers = [1, 2, 3, 4, 5]
-///     numbers += 10...15
-///     print(numbers)
-///     // Prints "[1, 2, 3, 4, 5, 10, 11, 12, 13, 14, 15]"
-///
-/// - Parameters:
-///   - lhs: The array to append to.
-///   - rhs: A collection.
-///
-/// - Complexity: O(*n*), where *n* is the length of the resulting array.
-public func += <
-  C : Collection
->(lhs: inout ${Self}<C.Iterator.Element>, rhs: C) {
-  lhs.append(contentsOf: rhs)
-}
 % end
 
 //===--- generic helpers --------------------------------------------------===//

stdlib/public/core/ContiguousArrayBuffer.swift

@@ -477,9 +477,11 @@ internal func += <Element, C : Collection>(
   let oldCount = lhs.count
   let newCount = oldCount + numericCast(rhs.count)
 
+  let buf: UnsafeMutableBufferPointer<Element>
+  
   if _fastPath(newCount <= lhs.capacity) {
+    buf = UnsafeMutableBufferPointer(start: lhs.firstElementAddress + oldCount, count: numericCast(rhs.count))
     lhs.count = newCount
-    (lhs.firstElementAddress + oldCount).initialize(from: rhs)
   }
   else {
     var newLHS = _ContiguousArrayBuffer<Element>(
@@ -490,7 +492,13 @@ internal func += <Element, C : Collection>(
       from: lhs.firstElementAddress, count: oldCount)
     lhs.count = 0
     swap(&lhs, &newLHS)
-    (lhs.firstElementAddress + oldCount).initialize(from: rhs)
+    buf = UnsafeMutableBufferPointer(start: lhs.firstElementAddress + oldCount, count: numericCast(rhs.count))
+  }
+
+  if let remainder = buf.initialize(from: rhs) {
+    // there were elements that didn't fit in the existing buffer,
+    // append them in slow sequence-only mode
+    //FIXME: handle this possibility
   }
 }
 

stdlib/public/core/ExistentialCollection.swift.gyb

@@ -192,8 +192,8 @@ internal class _AnyRandomAccessCollectionBox<Element>
     _abstract()
   }
 
-  internal func __copyContents(initializing ptr: UnsafeMutablePointer<Element>)
-    -> UnsafeMutablePointer<Element> {
+  internal func __copyContents(initializing buf: UnsafeMutableBufferPointer<Element>)
+    -> AnyIterator<Element>? {
     _abstract()
   }
 
@@ -387,9 +387,9 @@ internal final class _${Kind}Box<S : ${Kind}> : _Any${Kind}Box<S.Iterator.Elemen
   internal override func __copyToContiguousArray() -> ContiguousArray<Element> {
     return _base._copyToContiguousArray()
   }
-  internal override func __copyContents(initializing ptr: UnsafeMutablePointer<Element>)
-    -> UnsafeMutablePointer<Element> {
-    return _base._copyContents(initializing: ptr)
+  internal override func __copyContents(initializing buf: UnsafeMutableBufferPointer<Element>)
+    -> AnyIterator<Element>? {
+    return _base._copyContents(initializing: buf).map(AnyIterator.init)
   }
   internal override func _drop(
     while predicate: (Element) throws -> Bool
@@ -588,6 +588,8 @@ public struct AnySequence<Element> : Sequence {
     self.init(_ClosureBasedSequence(makeUnderlyingIterator))
   }
 
+  public typealias Iterator = AnyIterator<Element>
+
   internal init(_box: _AnySequenceBox<Element>) {
     self._box = _box
   }
@@ -602,7 +604,7 @@ extension Any${Kind} {
 %   else:
   /// Returns an iterator over the elements of this collection.
 %   end
-  public func makeIterator() -> AnyIterator<Element> {
+  public func makeIterator() -> Iterator {
     return _box._makeIterator()
   }
 
@@ -683,9 +685,9 @@ extension Any${Kind} {
     return self._box.__copyToContiguousArray()
   }
 
-  public func _copyContents(initializing ptr: UnsafeMutablePointer<Element>)
-    -> UnsafeMutablePointer<Element> {
-    return _box.__copyContents(initializing: ptr)
+  public func _copyContents(initializing buf: UnsafeMutableBufferPointer<Iterator.Element>)
+    -> Iterator? {
+    return _box.__copyContents(initializing: buf).map(AnyIterator.init)
   }
 }
 % end
@@ -808,6 +810,8 @@ public struct ${Self}<Element>
 //  public typealias Indices
 //    = Default${Traversal.replace('Forward', '')}Indices<${Self}>
 
+  public typealias Iterator = AnyIterator<Element>
+
   internal init(_box: _${Self}Box<Element>) {
     self._box = _box
   }

stdlib/public/core/LazyCollection.swift.gyb

@@ -75,11 +75,13 @@ public struct ${Self}<Base : ${TraversalCollection}> : LazyCollectionProtocol {
 /// Forward implementations to the base collection, to pick up any
 /// optimizations it might implement.
 extension ${Self} : Sequence {
+  
+  public typealias Iterator = Base.Iterator
 
   /// Returns an iterator over the elements of this sequence.
   ///
   /// - Complexity: O(1).
-  public func makeIterator() -> Base.Iterator {
+  public func makeIterator() -> Iterator {
     return _base.makeIterator()
   }
 
@@ -94,11 +96,10 @@ extension ${Self} : Sequence {
     return _base._copyToContiguousArray()
   }
 
-  @discardableResult
   public func _copyContents(
-    initializing ptr: UnsafeMutablePointer<Base.Iterator.Element>
-  ) -> UnsafeMutablePointer<Base.Iterator.Element> {
-    return _base._copyContents(initializing: ptr)
+    initializing buf: UnsafeMutableBufferPointer<Iterator.Element>
+  ) -> Iterator? {
+    return _base._copyContents(initializing: buf)
   }
 
   public func _customContainsEquatableElement(

stdlib/public/core/Sequence.swift

@@ -605,11 +605,11 @@ public protocol Sequence {
   /// in the same order.
   func _copyToContiguousArray() -> ContiguousArray<Iterator.Element>
 
-  /// Copy a Sequence into an array, returning one past the last
-  /// element initialized.
-  @discardableResult
-  func _copyContents(initializing ptr: UnsafeMutablePointer<Iterator.Element>)
-    -> UnsafeMutablePointer<Iterator.Element>
+  /// Copy `self` into an unsafe buffer, returning a partially-consumed 
+  /// iterator with any elements that didn't fit remaining.
+  func _copyContents(
+    initializing ptr: UnsafeMutableBufferPointer<Iterator.Element>
+  ) -> Iterator?
 }
 
 /// A default makeIterator() function for `IteratorProtocol` instances that
@@ -1317,19 +1317,23 @@ extension Sequence {
 }
 
 extension Sequence {
-  @discardableResult
   public func _copyContents(
-    initializing ptr: UnsafeMutablePointer<Iterator.Element>
-  ) -> UnsafeMutablePointer<Iterator.Element> {
-    var p = UnsafeMutablePointer<Iterator.Element>(ptr)
-    for x in IteratorSequence(self.makeIterator()) {
+    initializing buf: UnsafeMutableBufferPointer<Iterator.Element>
+  ) -> Iterator? {
+    var it = self.makeIterator()
+    // FIXME: decide if a nil buffer base pointer is valid or an error
+    guard let base = buf.baseAddress else { return it }
+    for p in base..<(base + buf.count) {
+      // FIXME: decide if underflow is an error that should trap
+      guard let x = it.next() else { break }
       p.initialize(to: x)
-      p += 1
     }
-    return p
+    return it
   }
 }
 
+// FIXME(ABI) #167(Recursive Protocol Constraints): IteratorSequence
+// shouldn't be necessary, iterators should be sequences.
 // Pending <rdar://problem/14011860> and <rdar://problem/14396120>,
 // pass an IteratorProtocol through IteratorSequence to give it "Sequence-ness"
 /// A sequence built around an iterator of type `Base`.

stdlib/public/core/SequenceWrapper.swift

@@ -119,8 +119,8 @@ extension Sequence
   /// element initialized.
   @discardableResult
   public func _copyContents(
-    initializing ptr: UnsafeMutablePointer<Base.Iterator.Element>
-  ) -> UnsafeMutablePointer<Base.Iterator.Element> {
-    return _base._copyContents(initializing: ptr)
+    initializing buf: UnsafeMutableBufferPointer<Base.Iterator.Element>
+  ) -> Base.Iterator? {
+    return _base._copyContents(initializing: buf)
   }
 }

stdlib/public/core/UnsafeBufferPointer.swift.gyb

@@ -224,6 +224,24 @@ extension Unsafe${Mutable}BufferPointer : CustomDebugStringConvertible {
 }
 %end
 
+
+extension UnsafeMutableBufferPointer {
+  /// Initializes memory in the buffer starting with the elements of `source`.
+  /// Returns an iterator to any elements of `source` that didn't fit in the 
+  /// buffer.
+  ///
+  /// - Precondition: The memory at `self..<self + count` is
+  ///   uninitialized.
+  ///
+  /// - Postcondition: The `Pointee`s at `self..<self + count` are
+  ///   initialized.
+  public func initialize<S: Sequence>(from source: S) -> S.Iterator?
+    where S.Iterator.Element == Iterator.Element {
+    return source._copyContents(initializing: self)
+  }
+}
+
+
 @available(*, unavailable, renamed: "UnsafeBufferPointerIterator")
 public struct UnsafeBufferPointerGenerator<Element> {}