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

Author: airspeedswift

stdlib/private/StdlibCollectionUnittest/CheckSequenceInstance.swift.gyb

@@ -86,16 +86,16 @@ 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,writtenUpTo) = sequence._copyContents(initializing: buf)
+    expectTrue(remainders.next() == nil, 
+      "_copyContents returned unwritten elements")
+    expectTrue(writtenUpTo == buf.endIndex, 
+      "_copyContents failed to use entire buffer")
+    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,11 @@ public struct ${Self}<
 
   /// Copy a Sequence into an array.
   public func _copyContents(
-    initializing ptr: UnsafeMutablePointer<Base.Iterator.Element>
-  ) -> UnsafeMutablePointer<Base.Iterator.Element> {
+    initializing buffer: UnsafeMutableBufferPointer<Iterator.Element>
+  ) -> (Iterator,UnsafeMutableBufferPointer<Iterator.Element>.Index) {
     Log._copyContents[selfType] += 1
-    return base._copyContents(initializing: ptr)
+    let (it,idx) = base._copyContents(initializing: buffer)
+    return (Iterator(wrapping: it),idx)
   }
 
 %     if Kind in ['Collection', 'MutableCollection', 'RangeReplaceableCollection']:

stdlib/public/SDK/Foundation/Data.swift

@@ -1425,7 +1425,7 @@ public struct Data : ReferenceConvertible, Equatable, Hashable, RandomAccessColl
                 // In the future, if we keep the malloced pointer and count inside this struct/ref instead of deferring to NSData, we may be able to do this more efficiently.
                 self.count = resultCount
             }
-            
+
             let shift = resultCount - currentCount
             let start = subrange.lowerBound
 
@@ -1437,7 +1437,11 @@ public struct Data : ReferenceConvertible, Equatable, Hashable, RandomAccessColl
                 }
 
                 if replacementCount != 0 {
-                    newElements._copyContents(initializing: bytes + start)
+                    let buf = UnsafeMutableBufferPointer(start: bytes + start, 
+                                                         count: replacementCount)
+                    var (it,idx) = newElements._copyContents(initializing: buf)
+                    precondition(it.next() == nil && idx == buf.endIndex,
+                      "newElements iterator returned different count to newElements.count")
                 }
             }
     }

stdlib/public/core/Arrays.swift.gyb

@@ -1344,7 +1344,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.
   ///
@@ -1358,53 +1358,40 @@ 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
+    let oldCapacity = self.capacity
     let newCount = oldCount + newElementsCount
 
     // Ensure uniqueness, mutability, and sufficient storage.  Note that
     // for consistency, we need unique self even if newElements is empty.
     self.reserveCapacity(
-      newCount > capacity ?
-      Swift.max(newCount, _growArrayCapacity(capacity))
+      newCount > oldCapacity ?
+      Swift.max(newCount, _growArrayCapacity(oldCapacity))
       : newCount)
 
-    (self._buffer.firstElementAddress + oldCount).initialize(from: newElements)
-    self._buffer.count = newCount
+    let startNewElements = _buffer.firstElementAddress + oldCount
+    let buf = UnsafeMutableBufferPointer(
+                start: startNewElements, 
+                count: self.capacity - oldCount)
+
+    let (remainder,writtenUpTo) = buf.initialize(from: newElements)
+    
+    // trap on underflow from the sequence's underestimate:
+    let writtenCount = buf.distance(from: buf.startIndex, to: writtenUpTo)
+    _precondition(newElementsCount <= writtenCount, 
+      "newElements.underestimatedCount was an overestimate")
+    // can't check for overflow as sequences can underestimate
+
+    _buffer.count += writtenCount
+
+    if writtenUpTo == buf.endIndex {
+      // there may be elements that didn't fit in the existing buffer,
+      // append them in slow sequence-only mode
+      _buffer._arrayAppendSequence(IteratorSequence(remainder))
+    }
   }
 
 %if Self == 'ArraySlice':
@@ -1697,23 +1684,34 @@ 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 buffer: UnsafeMutableBufferPointer<Iterator.Element>
+  ) -> (Iterator,UnsafeMutableBufferPointer<Iterator.Element>.Index) {
+
+    guard !self.isEmpty else { return (makeIterator(),buffer.startIndex) }
+
+    // It is not OK for there to be no pointer/not enough space, as this is
+    // a precondition and Array never lies about its count.
+    guard var p = buffer.baseAddress
+      else { _preconditionFailure("Attempt to copy contents into nil buffer pointer") }
+    _precondition(self.count <= buffer.count, 
+      "Insufficient space allocated to copy array contents")
+
+    if let s = _baseAddressIfContiguous {
+      p.initialize(from: s, count: self.count)
+      // Need a _fixLifetime bracketing the _baseAddressIfContiguous getter
+      // and all uses of the pointer it returns:
       _fixLifetime(self._owner)
-      return ptr + count
     } else {
-      var p = ptr
       for x in self {
         p.initialize(to: x)
         p += 1
       }
-      return p
     }
+
+    var it = IndexingIterator(_elements: self)
+    it._position = endIndex
+    return (it,buffer.index(buffer.startIndex, offsetBy: self.count))
   }
 }
 %end
@@ -1863,30 +1861,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,8 +492,14 @@ 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))
   }
+
+  var (remainders,writtenUpTo) = buf.initialize(from: rhs)
+
+  // ensure that exactly rhs.count elements were written
+  _precondition(remainders.next() == nil, "rhs underreported its count")
+  _precondition(writtenUpTo == buf.endIndex, "rhs overreported its count")    
 }
 
 extension _ContiguousArrayBuffer : RandomAccessCollection {

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>,UnsafeMutableBufferPointer<Element>.Index) {
     _abstract()
   }
 
@@ -387,9 +387,10 @@ 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>,UnsafeMutableBufferPointer<Element>.Index) {
+    let (it,idx) = _base._copyContents(initializing: buf)
+    return (AnyIterator(it),idx)
   }
   internal override func _drop(
     while predicate: (Element) throws -> Bool
@@ -588,6 +589,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 +605,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 +686,10 @@ 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>)
+  -> (AnyIterator<Element>,UnsafeMutableBufferPointer<Element>.Index) {
+    let (it,idx) = _box.__copyContents(initializing: buf)
+    return (AnyIterator(it),idx)
   }
 }
 % end
@@ -808,6 +812,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,UnsafeMutableBufferPointer<Iterator.Element>.Index) {
+    return _base._copyContents(initializing: buf)
   }
 
   public func _customContainsEquatableElement(