Make chunked(on:) include subject value in Element type (#142)

Author: natecook1000

Guides/Chunked.md

@@ -22,18 +22,20 @@ let chunks = numbers.chunked(by: { $0 <= $1 })
 
 The `chunk(on:)` method, by contrast, takes a projection of each element and
 separates chunks where the projection of two consecutive elements is not equal.
+The result includes both the projected value and the subsequence 
+that groups elements with that projected value:
 
 ```swift
 let names = ["David", "Kyle", "Karoy", "Nate"]
 let chunks = names.chunked(on: \.first!)
-// [["David"], ["Kyle", "Karoy"], ["Nate"]] 
+// [("D", ["David"]), ("K", ["Kyle", "Karoy"]), ("N", ["Nate"])] 
 ```
 
-The `chunks(ofCount:)` takes a `count` parameter (required to be > 0) and separates 
-the collection into `n` chunks of this given count. If the `count` parameter is 
-evenly divided by the count of the base `Collection` all the chunks will have 
-the count equals to the parameter. Otherwise, the last chunk will contain the 
-remaining elements.
+The `chunks(ofCount:)` method takes a `count` parameter (greater than zero) 
+and separates the collection into chunks of this given count.
+If the `count` parameter is evenly divided by the count of the base `Collection`,
+all the chunks will have a count equal to the parameter. 
+Otherwise, the last chunk will contain the remaining elements.
 
 ```swift
 let names = ["David", "Kyle", "Karoy", "Nate"]
@@ -44,17 +46,17 @@ let remaining = names.chunks(ofCount: 3)
 // equivalent to [["David", "Kyle", "Karoy"], ["Nate"]]
 ```
 
-The `chunks(ofCount:)` is the method of the [existing SE proposal][proposal]. 
-Unlike the `split` family of methods, the entire collection is included in the
-chunked result — joining the resulting chunks recreates the original collection. 
+The  `chunks(ofCount:)` is the subject of an [existing SE proposal][proposal].
+
+When "chunking" a collection, the entire collection is included in the result, 
+unlike the `split` family of methods, where separators are dropped.
+Joining the result of a chunking method call recreates the original collection. 
 
 ```swift
 c.elementsEqual(c.chunked(...).joined())
 // true
 ```
 
-Check the [proposal][proposal] detailed design section for more info. 
-
 [proposal]: https://github.com/apple/swift-evolution/pull/935
 
 ## Detailed Design
@@ -70,21 +72,21 @@ extension Collection {
 
     public func chunked<Subject: Equatable>(
         on projection: (Element) -> Subject
-    ) -> [SubSequence]
+    ) -> [(Subject, SubSequence)]
   }
 
   extension LazyCollectionProtocol {
     public func chunked(
         by belongInSameGroup: @escaping (Element, Element) -> Bool
-    ) -> Chunked<Elements>
+    ) -> ChunkedBy<Elements, Element>
 
     public func chunked<Subject: Equatable>(
         on projection: @escaping (Element) -> Subject
-    ) -> Chunked<Elements>
+    ) -> ChunkedOn<Elements, Subject>
 }
 ```
 
-The `Chunked` type is bidirectional when the wrapped collection is
+The `ChunkedBy` and `ChunkedOn` types are bidirectional when the wrapped collection is
 bidirectional.
 
 ### Complexity

Sources/Algorithms/Chunked.swift

@@ -10,8 +10,10 @@
 //===----------------------------------------------------------------------===//
 
 /// A collection wrapper that breaks a collection into chunks based on a
-/// predicate or projection.
-public struct Chunked<Base: Collection, Subject> {
+/// predicate.
+///
+/// Call `lazy.chunked(by:)` on a collection to create an instance of this type.
+public struct ChunkedBy<Base: Collection, Subject> {
   /// The collection that this instance provides a view onto.
   @usableFromInline
   internal let base: Base
@@ -45,7 +47,7 @@ public struct Chunked<Base: Collection, Subject> {
   }
 }
 
-extension Chunked: LazyCollectionProtocol {
+extension ChunkedBy: LazyCollectionProtocol {
   /// A position in a chunked collection.
   public struct Index: Comparable {
     /// The range corresponding to the chunk at this position.
@@ -106,9 +108,9 @@ extension Chunked: LazyCollectionProtocol {
   }
 }
 
-extension Chunked.Index: Hashable where Base.Index: Hashable {}
+extension ChunkedBy.Index: Hashable where Base.Index: Hashable {}
 
-extension Chunked: BidirectionalCollection
+extension ChunkedBy: BidirectionalCollection
   where Base: BidirectionalCollection
 {
   /// Returns the index in the base collection of the start of the chunk ending
@@ -131,11 +133,64 @@ extension Chunked: BidirectionalCollection
   }
 }
 
-@available(*, deprecated, renamed: "Chunked")
-public typealias LazyChunked<Base: Collection, Subject> = Chunked<Base, Subject>
+@available(*, deprecated, renamed: "ChunkedBy")
+public typealias LazyChunked<Base: Collection, Subject> = ChunkedBy<Base, Subject>
+
+@available(*, deprecated, renamed: "ChunkedBy")
+public typealias Chunked<Base: Collection, Subject> = ChunkedBy<Base, Subject>
+
+/// A collection wrapper that breaks a collection into chunks based on a
+/// predicate.
+///
+/// Call `lazy.chunked(on:)` on a collection to create an instance of this type.
+public struct ChunkedOn<Base: Collection, Subject> {
+  @usableFromInline
+  internal var chunked: ChunkedBy<Base, Subject>
+  
+  @inlinable
+  internal init(
+    base: Base,
+    projection: @escaping (Base.Element) -> Subject,
+    belongInSameGroup: @escaping (Subject, Subject) -> Bool
+  ) {
+    self.chunked = ChunkedBy(base: base, projection: projection, belongInSameGroup: belongInSameGroup)
+  }
+}
+
+extension ChunkedOn: LazyCollectionProtocol {
+  public typealias Index = ChunkedBy<Base, Subject>.Index
+  
+  @inlinable
+  public var startIndex: Index {
+    chunked.startIndex
+  }
+  
+  @inlinable
+  public var endIndex: Index {
+    chunked.endIndex
+  }
+  
+  @inlinable
+  public subscript(position: Index) -> (Subject, Base.SubSequence) {
+    let subsequence = chunked[position]
+    let subject = chunked.projection(subsequence.first!)
+    return (subject, subsequence)
+  }
+  
+  @inlinable
+  public func index(after i: Index) -> Index {
+    chunked.index(after: i)
+  }
+}
+
+extension ChunkedOn: BidirectionalCollection where Base: BidirectionalCollection {
+  public func index(before i: Index) -> Index {
+    chunked.index(before: i)
+  }
+}
 
 //===----------------------------------------------------------------------===//
-// lazy.chunked(by:)
+// lazy.chunked(by:) / lazy.chunked(on:)
 //===----------------------------------------------------------------------===//
 
 extension LazyCollectionProtocol {
@@ -146,8 +201,8 @@ extension LazyCollectionProtocol {
   @inlinable
   public func chunked(
     by belongInSameGroup: @escaping (Element, Element) -> Bool
-  ) -> Chunked<Elements, Element> {
-    Chunked(
+  ) -> ChunkedBy<Elements, Element> {
+    ChunkedBy(
       base: elements,
       projection: { $0 },
       belongInSameGroup: belongInSameGroup)
@@ -160,41 +215,38 @@ extension LazyCollectionProtocol {
   @inlinable
   public func chunked<Subject: Equatable>(
     on projection: @escaping (Element) -> Subject
-  ) -> Chunked<Elements, Subject> {
-    Chunked(
+  ) -> ChunkedOn<Elements, Subject> {
+    ChunkedOn(
       base: elements,
       projection: projection,
       belongInSameGroup: ==)
   }
 }
 
 //===----------------------------------------------------------------------===//
-// chunked(by:)
+// chunked(by:) / chunked(on:)
 //===----------------------------------------------------------------------===//
 
 extension Collection {
   /// Returns a collection of subsequences of this collection, chunked by
-  /// grouping elements that project to the same value according to the given
-  /// predicate.
+  /// the given predicate.
   ///
   /// - Complexity: O(*n*), where *n* is the length of this collection.
   @inlinable
-  internal func chunked<Subject>(
-    on projection: (Element) throws -> Subject,
-    by belongInSameGroup: (Subject, Subject) throws -> Bool
+  public func chunked(
+    by belongInSameGroup: (Element, Element) throws -> Bool
   ) rethrows -> [SubSequence] {
     guard !isEmpty else { return [] }
     var result: [SubSequence] = []
 
     var start = startIndex
-    var subject = try projection(self[start])
+    var current = self[start]
 
     for (index, element) in indexed().dropFirst() {
-      let nextSubject = try projection(element)
-      if try !belongInSameGroup(subject, nextSubject) {
+      if try !belongInSameGroup(current, element) {
         result.append(self[start..<index])
         start = index
-        subject = nextSubject
+        current = element
       }
     }
 
@@ -204,17 +256,6 @@ extension Collection {
 
     return result
   }
-  
-  /// Returns a collection of subsequences of this collection, chunked by
-  /// the given predicate.
-  ///
-  /// - Complexity: O(*n*), where *n* is the length of this collection.
-  @inlinable
-  public func chunked(
-    by belongInSameGroup: (Element, Element) throws -> Bool
-  ) rethrows -> [SubSequence] {
-    try chunked(on: { $0 }, by: belongInSameGroup)
-  }
 
   /// Returns a collection of subsequences of this collection, chunked by
   /// grouping elements that project to the same value.
@@ -223,8 +264,27 @@ extension Collection {
   @inlinable
   public func chunked<Subject: Equatable>(
     on projection: (Element) throws -> Subject
-  ) rethrows -> [SubSequence] {
-    try chunked(on: projection, by: ==)
+  ) rethrows -> [(Subject, SubSequence)] {
+    guard !isEmpty else { return [] }
+    var result: [(Subject, SubSequence)] = []
+    
+    var start = startIndex
+    var subject = try projection(self[start])
+    
+    for (index, element) in indexed().dropFirst() {
+      let nextSubject = try projection(element)
+      if subject != nextSubject {
+        result.append((subject, self[start..<index]))
+        start = index
+        subject = nextSubject
+      }
+    }
+    
+    if start != endIndex {
+      result.append((subject, self[start..<endIndex]))
+    }
+    
+    return result
   }
 }
 

Tests/SwiftAlgorithmsTests/ChunkedTests.swift

@@ -47,8 +47,12 @@ final class ChunkedTests: XCTestCase {
   func testSimple() {
     // Example
     let names = ["David", "Kyle", "Karoy", "Nate"]
-    let chunks = names.chunked(on: { $0.first })
-    XCTAssertEqualSequences([["David"], ["Kyle", "Karoy"], ["Nate"]], chunks)
+    let chunks = names.chunked(on: { $0.first! })
+    let expected: [(Character, ArraySlice<String>)] = [
+      ("D", ["David"]),
+      ("K", ["Kyle", "Karoy"]),
+      ("N", ["Nate"])]
+    XCTAssertEqualSequences(expected, chunks, by: ==)
 
     // Empty sequence
     XCTAssertEqual(0, names.prefix(0).chunked(on: { $0.first }).count)
@@ -59,28 +63,30 @@ final class ChunkedTests: XCTestCase {
   }
 
   func testChunkedOn() {
-    validateFruitChunks(fruits.chunked(on: { $0.first }))
+    validateFruitChunks(fruits.chunked(on: { $0.first }).map { $1 })
 
     let lazyChunks = fruits.lazy.chunked(on: { $0.first })
-    validateFruitChunks(lazyChunks)
+    validateFruitChunks(lazyChunks.map { $1 })
+    validateIndexTraversals(lazyChunks)
   }
 
   func testChunkedBy() {
     validateFruitChunks(fruits.chunked(by: { $0.first == $1.first }))
 
     let lazyChunks = fruits.lazy.chunked(by: { $0.first == $1.first })
     validateFruitChunks(lazyChunks)
+    validateIndexTraversals(lazyChunks)
   }
 
   func testChunkedLazy() {
     XCTAssertLazySequence(fruits.lazy.chunked(by: { $0.first == $1.first }))
     XCTAssertLazySequence(fruits.lazy.chunked(on: { $0.first }))
   }
 
-  
   //===----------------------------------------------------------------------===//
   // Tests for `chunks(ofCount:)`
   //===----------------------------------------------------------------------===//
+  
   func testChunksOfCount() {
     XCTAssertEqualSequences([Int]().chunks(ofCount: 1), [])
     XCTAssertEqualSequences([Int]().chunks(ofCount: 5), [])