[stdlib] Update partition(by:) tests

Adds a dispatch test for partition(by:), since it's now a protocol
requirement. Also adds a new logging collection wrapper that only
logs when _withUnsafeMutableBufferPointerIfSupported is called --
any calls to this method from dispatched methods are uncounted by
the standard logging wrappers.

Author: natecook1000

stdlib/private/StdlibCollectionUnittest/CheckMutableCollectionType.swift.gyb

@@ -12,7 +12,7 @@
 
 import StdlibUnittest
 
-// These tests are shared between partition() and sort().
+// These tests are shared between partition(by:) and sort().
 public struct PartitionExhaustiveTest {
   public let sequence: [Int]
   public let loc: SourceLoc
@@ -67,6 +67,18 @@ internal func _mapInPlace<C : MutableCollection>(
   }
 }
 
+internal func makeBufferAccessLoggingMutableCollection<
+  C : MutableCollection & BidirectionalCollection
+>(wrapping c: C) -> BufferAccessLoggingMutableBidirectionalCollection<C> {
+  return BufferAccessLoggingMutableBidirectionalCollection(wrapping: c)
+}
+
+internal func makeBufferAccessLoggingMutableCollection<
+  C : MutableCollection & RandomAccessCollection
+>(wrapping c: C) -> BufferAccessLoggingMutableBidirectionalCollection<C> {
+  return BufferAccessLoggingMutableBidirectionalCollection(wrapping: c)
+}
+
 extension TestSuite {
   public func addMutableCollectionTests<
     C : MutableCollection,
@@ -575,7 +587,7 @@ self.test("\(testNamePrefix).sorted/${'Predicate' if predicate else 'WhereElemen
 % end
 
 //===----------------------------------------------------------------------===//
-// partition()
+// partition(by:)
 //===----------------------------------------------------------------------===//
 
 func checkPartition(
@@ -584,7 +596,6 @@ func checkPartition(
   lessImpl: ((Int, Int) -> Bool),
   verifyOrder: Bool
 ) {
-  let extract = extractValue
   let elements: [OpaqueValue<Int>] =
     zip(sequence, 0..<sequence.count).map {
       OpaqueValue($0, identity: $1)
@@ -594,23 +605,23 @@ func checkPartition(
   let closureLifetimeTracker = LifetimeTracked(0)
   let pivot = c.partition(by: { val in
     _blackHole(closureLifetimeTracker)
-    return !lessImpl(extract(val).value, pivotValue)
+    return !lessImpl(extractValue(val).value, pivotValue)
   })
 
   // Check that we didn't lose any elements.
-  let identities = c.map { extract($0).identity }
+  let identities = c.map { extractValue($0).identity }
   expectEqualsUnordered(0..<sequence.count, identities)
 
   if verifyOrder {
     // All the elements in the first partition are less than the pivot
     // value.
     for i in c[c.startIndex..<pivot].indices {
-      expectLT(extract(c[i]).value, pivotValue)
+      expectLT(extractValue(c[i]).value, pivotValue)
     }
     // All the elements in the second partition are greater or equal to
     // the pivot value.
     for i in c[pivot..<c.endIndex].indices {
-      expectGE(extract(c[i]).value, pivotValue)
+      expectGE(extractValue(c[i]).value, pivotValue)
     }
   }
 }
@@ -776,6 +787,37 @@ self.test("\(testNamePrefix).reverse()") {
   }
 }
 
+//===----------------------------------------------------------------------===//
+// partition(by:)
+//===----------------------------------------------------------------------===//
+
+self.test("\(testNamePrefix).partition/DispatchesThrough_withUnsafeMutableBufferPointerIfSupported") {
+  let sequence = [ 5, 4, 3, 2, 1 ]
+  let elements: [OpaqueValue<Int>] =
+    zip(sequence, 0..<sequence.count).map {
+      OpaqueValue($0, identity: $1)
+    }
+  let c = makeWrappedCollection(elements)
+  var lc = makeBufferAccessLoggingMutableCollection(wrapping: c)
+
+  let closureLifetimeTracker = LifetimeTracked(0)
+  let first = c.first
+  let pivot = lc.partition(by: { val in
+    _blackHole(closureLifetimeTracker)
+    return !(extractValue(val).value < extractValue(first!).value)
+  })
+
+  expectEqual(
+    1, lc.log._withUnsafeMutableBufferPointerIfSupported[lc.dynamicType])
+  expectEqual(
+    withUnsafeMutableBufferPointerIsSupported ? 1 : 0,
+    lc.log._withUnsafeMutableBufferPointerIfSupportedNonNilReturns[lc.dynamicType])
+
+  expectEqual(4, lc.distance(from: lc.startIndex, to: pivot))
+  expectEqualsUnordered([1, 2, 3, 4], lc.prefix(upTo: pivot).map { extractValue($0).value })
+  expectEqualsUnordered([5], lc.suffix(from: pivot).map { extractValue($0).value })
+}
+
 //===----------------------------------------------------------------------===//
 
   } // addMutableBidirectionalCollectionTests
@@ -956,39 +998,6 @@ self.test("\(testNamePrefix).sort/${'Predicate' if predicate else 'WhereElementI
 
 % end
 
-//===----------------------------------------------------------------------===//
-// partition()
-//===----------------------------------------------------------------------===//
-
-// FIXME(tests): Move to addMutableCollectionTests?
-self.test("\(testNamePrefix).partition/DispatchesThrough_withUnsafeMutableBufferPointerIfSupported") {
-  let sequence = [ 5, 4, 3, 2, 1 ]
-  let extract = extractValue
-  let elements: [OpaqueValue<Int>] =
-    zip(sequence, 0..<sequence.count).map {
-      OpaqueValue($0, identity: $1)
-    }
-  let c = makeWrappedCollection(elements)
-
-  var lc = LoggingMutableRandomAccessCollection(wrapping: c)
-
-  let closureLifetimeTracker = LifetimeTracked(0)
-  let first = c.first
-  let pivot = lc.partition(by: { val in
-    _blackHole(closureLifetimeTracker)
-    return !(extract(val).value < extract(first!).value)
-  })
-
-  expectEqual(
-    1, lc.log._withUnsafeMutableBufferPointerIfSupported[lc.dynamicType])
-  expectEqual(
-    withUnsafeMutableBufferPointerIsSupported ? 1 : 0,
-    lc.log._withUnsafeMutableBufferPointerIfSupportedNonNilReturns[lc.dynamicType])
-
-  expectEqual(4, lc.distance(from: lc.startIndex, to: pivot))
-  expectEqualSequence([ 1, 4, 3, 2, 5 ], lc.map { extract($0).value })
-}
-
 //===----------------------------------------------------------------------===//
 
   } // addMutableRandomAccessCollectionTests

stdlib/private/StdlibCollectionUnittest/LoggingWrappers.swift.gyb

@@ -140,6 +140,7 @@ public class MutableCollectionLog : CollectionLog {
   }
   public static var subscriptIndexSet = TypeIndexed(0)
   public static var subscriptRangeSet = TypeIndexed(0)
+  public static var partitionBy = TypeIndexed(0)
   public static var _withUnsafeMutableBufferPointerIfSupported = TypeIndexed(0)
   public static var _withUnsafeMutableBufferPointerIfSupportedNonNilReturns =
     TypeIndexed(0)
@@ -440,6 +441,13 @@ public struct ${Self}<
 %     end
 
 %     if Kind == 'MutableCollection':
+  public mutating func partition(
+    by belongsInSecondPartition: @noescape (Iterator.Element) throws -> Bool
+  ) rethrows -> Index {
+    Log.partitionBy[selfType] += 1
+    return try base.partition(by: belongsInSecondPartition)
+  }
+
   public mutating func _withUnsafeMutableBufferPointerIfSupported<R>(
     _ body: @noescape (UnsafeMutablePointer<Iterator.Element>, Int) throws -> R
   ) rethrows -> R? {
@@ -565,6 +573,95 @@ public struct ${Self}<
 %   end
 % end
 
+//===----------------------------------------------------------------------===//
+// Collections that count calls to `_withUnsafeMutableBufferPointerIfSupported`
+//===----------------------------------------------------------------------===//
+
+% for Traversal in TRAVERSALS:
+%   Self = 'BufferAccessLoggingMutable' + collectionForTraversal(Traversal)
+/// Interposes between `_withUnsafeMutableBufferPointerIfSupported` method calls
+/// to increment a counter. Calls to this method from within dispatched methods
+/// are uncounted by the standard logging collection wrapper.
+public struct ${Self}<
+  Base : MutableCollection & ${collectionForTraversal(Traversal)}
+> : MutableCollection, ${collectionForTraversal(Traversal)}, LoggingType {
+
+  public var base: Base
+
+  public typealias Log = MutableCollectionLog
+
+  public typealias SubSequence = Base.SubSequence
+
+  public typealias Iterator = Base.Iterator
+
+  public init(wrapping base: Base) {
+    self.base = base
+  }
+
+  public func makeIterator() -> Iterator {
+    return base.makeIterator()
+  }
+
+  public typealias Index = Base.Index
+
+  public var startIndex: Index {
+    return base.startIndex
+  }
+
+  public var endIndex: Index {
+    return base.endIndex
+  }
+
+  public subscript(position: Index) -> Base.Iterator.Element {
+    get {
+      return base[position]
+    }
+    set {
+      base[position] = newValue
+    }
+  }
+
+  public subscript(bounds: Range<Index>) -> SubSequence {
+    get {
+      return base[bounds]
+    }
+    set {
+      base[bounds] = newValue
+    }
+  }
+
+  public func index(after i: Index) -> Index {
+    return base.index(after: i)
+  }
+
+%     if Traversal in ['Bidirectional', 'RandomAccess']:
+  public func index(before i: Index) -> Index {
+    return base.index(before: i)
+  }
+%     end
+
+  public func index(_ i: Index, offsetBy n: Base.IndexDistance) -> Index {
+    return base.index(i, offsetBy: n)
+  }
+
+  public func distance(from start: Index, to end: Index) -> Base.IndexDistance {
+    return base.distance(from: start, to: end)
+  }
+
+  public mutating func _withUnsafeMutableBufferPointerIfSupported<R>(
+    _ body: @noescape (UnsafeMutablePointer<Iterator.Element>, Int) throws -> R
+  ) rethrows -> R? {
+    print("Log._withUnsafeMutableBufferPointerIfSupported[selfType] += 1")
+    Log._withUnsafeMutableBufferPointerIfSupported[selfType] += 1
+    let result = try base._withUnsafeMutableBufferPointerIfSupported(body)
+    if result != nil {
+      Log._withUnsafeMutableBufferPointerIfSupportedNonNilReturns[selfType] += 1
+    }
+    return result
+  }
+}
+% end
+
 //===----------------------------------------------------------------------===//
 // Custom assertions
 //===----------------------------------------------------------------------===//

validation-test/stdlib/CollectionType.swift.gyb

@@ -778,6 +778,16 @@ CollectionTypeTests.test("Collection/suffix(from:)/dispatch") {
 // MutableCollection
 //===----------------------------------------------------------------------===//
 
+//===----------------------------------------------------------------------===//
+// partition(by:)
+//===----------------------------------------------------------------------===//
+
+CollectionTypeTests.test("partition(by:)/dispatch") {
+  var tester = MutableCollectionLog.dispatchTester([OpaqueValue(1)])
+  _ = tester.partition(by: { _ in true })
+  expectCustomizable(tester, tester.log.partitionBy)
+}
+
 //===----------------------------------------------------------------------===//
 // _withUnsafeMutableBufferPointerIfSupported()
 //===----------------------------------------------------------------------===//