Evenly divide a collection into chunks (#96)

Adds `evenlyChunked(in:)` as a `Collection` method that divides a
collection into a specific number of chunks as evenly as possible.

Co-authored-by: Tim Vermeulen <[email protected]>

Author: timvermeulen

Guides/Chunked.md

@@ -3,16 +3,16 @@
 [[Source](https://github.com/apple/swift-algorithms/blob/main/Sources/Algorithms/Chunked.swift) | 
  [Tests](https://github.com/apple/swift-algorithms/blob/main/Tests/SwiftAlgorithmsTests/ChunkedTests.swift)]
 
-Break a collection into subsequences where consecutive elements pass a binary
-predicate, or where all elements in each chunk project to the same value.
+Break a collection into nonoverlapping subsequences:
 
-Also includes a `chunks(ofCount:)` that breaks a collection into subsequences 
-of a given `count`.
+* `chunked(by:)` forms chunks of consecutive elements that pass a binary predicate,
+* `chunked(on:)` forms chunks of consecutive elements that project to equal values,
+* `chunks(ofCount:)` forms chunks of a given size, and
+* `evenlyChunked(in:)` forms a given number of equally-sized chunks.
 
-There are two variations of the `chunked` method: `chunked(by:)` and
-`chunked(on:)`. `chunked(by:)` uses a binary predicate to test consecutive
-elements, separating chunks where the predicate returns `false`. For example,
-you can chunk a collection into ascending sequences using this method:
+`chunked(by:)` uses a binary predicate to test consecutive elements, separating
+chunks where the predicate returns `false`. For example, you can chunk a
+collection into ascending sequences using this method:
 
 ```swift
 let numbers = [10, 20, 30, 10, 40, 40, 10, 20]
@@ -31,11 +31,10 @@ let chunks = names.chunked(on: \.first!)
 // [("D", ["David"]), ("K", ["Kyle", "Karoy"]), ("N", ["Nate"])] 
 ```
 
-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.
+The `chunks(ofCount:)` method takes a `count` parameter (required to be > 0) and
+separates the collection into chunks of this given count. If the length of the
+collection is a multiple of the `count` parameter, all chunks will have the
+a count equal to the parameter. Otherwise, the last chunk will contain the remaining elements.
 
 ```swift
 let names = ["David", "Kyle", "Karoy", "Nate"]
@@ -46,7 +45,21 @@ let remaining = names.chunks(ofCount: 3)
 // equivalent to [["David", "Kyle", "Karoy"], ["Nate"]]
 ```
 
-The `chunks(ofCount:)` is the subject of an [existing SE proposal][proposal].
+The `chunks(ofCount:)` method was previously [proposed](proposal) for inclusion
+in the standard library.
+
+The `evenlyChunked(in:)` method takes a `count` parameter and divides the
+collection into `count` number of equally-sized chunks. If the length of the
+collection is not a multiple of the `count` parameter, the chunks at the start
+will be longer than the chunks at the end.
+
+```swift
+let evenChunks = (0..<15).evenlyChunked(in: 3)
+// equivalent to [0..<5, 5..<10, 10..<15]
+
+let nearlyEvenChunks = (0..<15).evenlyChunked(in: 4)
+// equivalent to [0..<4, 4..<8, 8..<12, 12..<15]
+```
 
 When "chunking" a collection, the entire collection is included in the result,
 unlike the `split` family of methods, where separators are dropped.
@@ -61,38 +74,40 @@ c.elementsEqual(c.chunked(...).joined())
 
 ## Detailed Design
 
-The three methods are added as extension to `Collection`. `chunked(by:)` and
-`chunked(on:)` are eager by default, both with a matching version that return a
-lazy wrapper added to `LazySequenceProtocol`.
+The four methods are added to `Collection`, with matching versions of
+`chunked(by:)` and `chunked(on:)` that return a lazy wrapper added to
+`LazyCollectionProtocol`.
 
 ```swift
 extension Collection {
-  public func chunked(
-      by belongInSameGroup: (Element, Element) -> Bool
-  ) -> [SubSequence]
+    public func chunked(
+        by belongInSameGroup: (Element, Element) -> Bool
+    ) -> [SubSequence]
 
-  public func chunked<Subject: Equatable>(
+    public func chunked<Subject: Equatable>(
       on projection: (Element) -> Subject
-  ) -> [(Subject, SubSequence)]
-  
-  public func chunks(ofCount count: Int) -> ChunksOfCountCollection<Self>
+    ) -> [SubSequence]
+      
+    public func chunks(ofCount count: Int) -> ChunkedByCount<Self>
+      
+    public func evenlyChunked(in count: Int) -> EvenChunks<Self>
 }
 
-extension LazySequenceProtocol where Self: Collection, Elements: Collection {
-  public func chunked(
-      by belongInSameGroup: @escaping (Element, Element) -> Bool
-  ) -> ChunkedByCollection<Elements, Element>
+extension LazyCollectionProtocol {
+    public func chunked(
+        by belongInSameGroup: @escaping (Element, Element) -> Bool
+    ) -> ChunkedByCollection<Elements>
 
-  public func chunked<Subject: Equatable>(
-      on projection: @escaping (Element) -> Subject
-  ) -> ChunkedOnCollection<Elements, Subject>
+    public func chunked<Subject: Equatable>(
+        on projection: @escaping (Element) -> Subject
+    ) -> ChunkedOnCollection<Elements, Subject>
 }
 ```
 
-The `ChunkedByCollection`, `ChunkedOnCollection`, and `ChunksOfCountCollection`
-types are bidirectional when the wrapped collection is bidirectional.
-`ChunksOfCountCollection` also conforms to `LazySequenceProtocol` when the base
-collection conforms.
+Each of the "chunked" collection types are bidirectional when the wrapped
+collection is bidirectional. `ChunksOfCountCollection` and 
+`EvenChunksCollection` also conform to `RandomAccessCollection` and 
+`LazySequenceProtocol` when their base collections conform.
 
 ### Complexity
 
@@ -120,5 +135,5 @@ into potentially overlapping subsequences.
 **Ruby:** Ruby’s `Enumerable` class defines `chunk_while` and `chunk`, which map
 to the proposed `chunked(by:)` and `chunked(on:)` methods.
 
-**Rust:** Rust defines a variety of size-based `chunks` methods, but doesn’t
-include any with the functionality described here.
+**Rust:** Rust defines a variety of size-based `chunks` methods, of which the
+standard version corresponds to the `chunks(ofCount:)` method defined here.

Sources/Algorithms/Chunked.swift

@@ -198,6 +198,210 @@ extension ChunkedOnCollection: BidirectionalCollection
 
 extension ChunkedOnCollection: LazyCollectionProtocol {}
 
+/// A collection wrapper that evenly breaks a collection into a given number of
+/// chunks.
+public struct EvenChunksCollection<Base: Collection> {
+  /// The base collection.
+  @usableFromInline
+  internal let base: Base
+  
+  /// The number of equal chunks the base collection is divided into.
+  @usableFromInline
+  internal let numberOfChunks: Int
+  
+  /// The count of the base collection.
+  @usableFromInline
+  internal let baseCount: Int
+  
+  /// The upper bound of the first chunk.
+  @usableFromInline
+  internal var firstUpperBound: Base.Index
+  
+  @inlinable
+  internal init(base: Base, numberOfChunks: Int) {
+    self.base = base
+    self.numberOfChunks = numberOfChunks
+    self.baseCount = base.count
+    self.firstUpperBound = base.startIndex
+    
+    if numberOfChunks > 0 {
+      firstUpperBound = endOfChunk(startingAt: base.startIndex, offset: 0)
+    }
+  }
+}
+
+extension EvenChunksCollection {
+  /// Returns the number of chunks with size `smallChunkSize + 1` at the start
+  /// of this collection.
+  @inlinable
+  internal var numberOfLargeChunks: Int {
+    baseCount % numberOfChunks
+  }
+  
+  /// Returns the size of a chunk at a given offset.
+  @inlinable
+  internal func sizeOfChunk(offset: Int) -> Int {
+    let isLargeChunk = offset < numberOfLargeChunks
+    return baseCount / numberOfChunks + (isLargeChunk ? 1 : 0)
+  }
+  
+  /// Returns the index in the base collection of the end of the chunk starting
+  /// at the given index.
+  @inlinable
+  internal func endOfChunk(startingAt start: Base.Index, offset: Int) -> Base.Index {
+    base.index(start, offsetBy: sizeOfChunk(offset: offset))
+  }
+  
+  /// Returns the index in the base collection of the start of the chunk ending
+  /// at the given index.
+  @inlinable
+  internal func startOfChunk(endingAt end: Base.Index, offset: Int) -> Base.Index {
+    base.index(end, offsetBy: -sizeOfChunk(offset: offset))
+  }
+  
+  /// Returns the index that corresponds to the chunk that starts at the given
+  /// base index.
+  @inlinable
+  internal func indexOfChunk(startingAt start: Base.Index, offset: Int) -> Index {
+    guard offset != numberOfChunks else { return endIndex }
+    let end = endOfChunk(startingAt: start, offset: offset)
+    return Index(start..<end, offset: offset)
+  }
+  
+  /// Returns the index that corresponds to the chunk that ends at the given
+  /// base index.
+  @inlinable
+  internal func indexOfChunk(endingAt end: Base.Index, offset: Int) -> Index {
+    let start = startOfChunk(endingAt: end, offset: offset)
+    return Index(start..<end, offset: offset)
+  }
+}
+
+extension EvenChunksCollection: Collection {
+  public struct Index: Comparable {
+    /// The range corresponding to the chunk at this position.
+    @usableFromInline
+    internal var baseRange: Range<Base.Index>
+    
+    /// The offset corresponding to the chunk at this position. The first chunk
+    /// has offset `0` and all other chunks have an offset `1` greater than the
+    /// previous.
+    @usableFromInline
+    internal var offset: Int
+    
+    @inlinable
+    internal init(_ baseRange: Range<Base.Index>, offset: Int) {
+      self.baseRange = baseRange
+      self.offset = offset
+    }
+    
+    @inlinable
+    public static func == (lhs: Self, rhs: Self) -> Bool {
+      lhs.offset == rhs.offset
+    }
+    
+    @inlinable
+    public static func < (lhs: Self, rhs: Self) -> Bool {
+      lhs.offset < rhs.offset
+    }
+  }
+  
+  public typealias Element = Base.SubSequence
+
+  @inlinable
+  public var startIndex: Index {
+    Index(base.startIndex..<firstUpperBound, offset: 0)
+  }
+  
+  @inlinable
+  public var endIndex: Index {
+    Index(base.endIndex..<base.endIndex, offset: numberOfChunks)
+  }
+  
+  @inlinable
+  public func index(after i: Index) -> Index {
+    precondition(i != endIndex, "Can't advance past endIndex")
+    let start = i.baseRange.upperBound
+    return indexOfChunk(startingAt: start, offset: i.offset + 1)
+  }
+  
+  @inlinable
+  public subscript(position: Index) -> Element {
+    precondition(position != endIndex)
+    return base[position.baseRange]
+  }
+  
+  @inlinable
+  public func index(_ i: Index, offsetBy distance: Int) -> Index {
+    /// Returns the base distance between two `EvenChunksCollection` indices
+    /// from the end of one to the start of the other, when given their offsets.
+    func baseDistance(from offsetA: Int, to offsetB: Int) -> Int {
+      let smallChunkSize = baseCount / numberOfChunks
+      let numberOfChunks = (offsetB - offsetA) - 1
+      
+      let largeChunksEnd = Swift.min(self.numberOfLargeChunks, offsetB)
+      let largeChunksStart = Swift.min(self.numberOfLargeChunks, offsetA + 1)
+      let numberOfLargeChunks = largeChunksEnd - largeChunksStart
+      
+      return smallChunkSize * numberOfChunks + numberOfLargeChunks
+    }
+    
+    if distance == 0 {
+      return i
+    } else if distance > 0 {
+      let offset = i.offset + distance
+      let baseOffset = baseDistance(from: i.offset, to: offset)
+      let start = base.index(i.baseRange.upperBound, offsetBy: baseOffset)
+      return indexOfChunk(startingAt: start, offset: offset)
+    } else {
+      let offset = i.offset + distance
+      let baseOffset = baseDistance(from: offset, to: i.offset)
+      let end = base.index(i.baseRange.lowerBound, offsetBy: -baseOffset)
+      return indexOfChunk(endingAt: end, offset: offset)
+    }
+  }
+  
+  @inlinable
+  public func index(_ i: Index, offsetBy distance: Int, limitedBy limit: Index) -> Index? {
+    if distance >= 0 {
+      if (0..<distance).contains(self.distance(from: i, to: limit)) {
+        return nil
+      }
+    } else {
+      if (0..<(-distance)).contains(self.distance(from: limit, to: i)) {
+        return nil
+      }
+    }
+    return index(i, offsetBy: distance)
+  }
+  
+  @inlinable
+  public func distance(from start: Index, to end: Index) -> Int {
+    end.offset - start.offset
+  }
+}
+
+extension EvenChunksCollection.Index: Hashable where Base.Index: Hashable {}
+
+extension EvenChunksCollection: BidirectionalCollection
+  where Base: BidirectionalCollection
+{
+  @inlinable
+  public func index(before i: Index) -> Index {
+    precondition(i != startIndex, "Can't advance before startIndex")
+    return indexOfChunk(endingAt: i.baseRange.lowerBound, offset: i.offset - 1)
+  }
+}
+
+extension EvenChunksCollection: RandomAccessCollection
+  where Base: RandomAccessCollection {}
+
+extension EvenChunksCollection: LazySequenceProtocol
+  where Base: LazySequenceProtocol {}
+
+extension EvenChunksCollection: LazyCollectionProtocol
+  where Base: LazyCollectionProtocol {}
+
 //===----------------------------------------------------------------------===//
 // lazy.chunked(by:) / lazy.chunked(on:)
 //===----------------------------------------------------------------------===//
@@ -583,5 +787,40 @@ extension Collection {
 
 extension ChunksOfCountCollection.Index: Hashable where Base.Index: Hashable {}
 
-extension ChunksOfCountCollection: LazySequenceProtocol, LazyCollectionProtocol
+extension ChunksOfCountCollection: LazySequenceProtocol
   where Base: LazySequenceProtocol {}
+
+extension ChunksOfCountCollection: LazyCollectionProtocol
+  where Base: LazyCollectionProtocol {}
+
+//===----------------------------------------------------------------------===//
+// evenlyChunked(in:)
+//===----------------------------------------------------------------------===//
+
+extension Collection {
+  /// Returns a collection of `count` evenly divided subsequences of this
+  /// collection.
+  ///
+  /// This method divides the collection into a given number of equally sized
+  /// chunks. If the length of the collection is not divisible by `count`, the
+  /// chunks at the start will be longer than the chunks at the end, like in
+  /// this example:
+  ///
+  ///     for chunk in "Hello, world!".evenlyChunked(in: 5) {
+  ///         print(chunk)
+  ///     }
+  ///     // "Hel"
+  ///     // "lo,"
+  ///     // " wo"
+  ///     // "rl"
+  ///     // "d!"
+  ///
+  /// - Complexity: O(1) if the collection conforms to `RandomAccessCollection`,
+  ///   otherwise O(*n*), where *n* is the length of the collection.
+  @inlinable
+  public func evenlyChunked(in count: Int) -> EvenChunksCollection<Self> {
+    precondition(count >= 0, "Can't divide into a negative number of chunks")
+    precondition(count > 0 || isEmpty, "Can't divide a non-empty collection into 0 chunks")
+    return EvenChunksCollection(base: self, numberOfChunks: count)
+  }
+}