Chaging implementation to store base slice range on the chunk index

Author: LucianoPAlmeida

Guides/Chunked.md

@@ -38,10 +38,10 @@ remaining elements.
 ```swift
 let names = ["David", "Kyle", "Karoy", "Nate"]
 let evenly = names.chunks(ofCount: 2)
-// [["David", "Kyle"], ["Karoy", "Nate"]] 
+// equivalent to [["David", "Kyle"], ["Karoy", "Nate"]] 
 
 let remaining = names.chunks(ofCount: 3)
-// [["David", "Kyle", "Karoy"], "Nate"]] 
+// equivalent to [["David", "Kyle", "Karoy"], ["Nate"]]
 ```
 
 The `chunks(ofCount:)` is the method of the [existing SE proposal][proposal]. 

Sources/Algorithms/Chunked.swift

@@ -254,98 +254,123 @@ extension Collection {
 /// A collection that presents the elements of its base collection
 /// in `SubSequence` chunks of any given size.
 ///
-/// A ChunkedCollection is a lazy view on the base Collection, but it does not implicitly confer
+/// A `ChunkedByCount` is a lazy view on the base Collection, but it does not implicitly confer
 /// laziness on algorithms applied to its result.  In other words, for ordinary collections `c`:
 ///
 /// * `c.chunks(of: 3)` does not create new storage
 /// * `c.chunks(of: 3).map(f)` maps eagerly and returns a new array
 /// * `c.lazy.chunks(of: 3).map(f)` maps lazily and returns a `LazyMapCollection`
-public struct ChunkedCollection<Base: Collection> {
+public struct ChunkedByCount<Base: Collection> {
 
   public typealias Element = Base.SubSequence
 
-  public let base: Base
   @usableFromInline
-  internal let _chunkCount: Int
+  internal let base: Base
+  
+  @usableFromInline
+  internal let chunkCount: Int
 
+  @usableFromInline
+  internal var computedStartIndex: Index
+
   ///  Creates a view instance that presents the elements of `base`
   ///  in `SubSequence` chunks of the given size.
   ///
-  /// - Complexity: O(1)
+  /// - Complexity: O(n)
   @inlinable
   internal init(_base: Base, _chunkCount: Int) {
     self.base = _base
-    self._chunkCount = _chunkCount
+    self.chunkCount = _chunkCount
+    
+    // Compute the start index upfront in order to make
+    // start index a O(1) lookup.
+    let baseEnd = _base.index(
+      _base.startIndex, offsetBy: _chunkCount,
+      limitedBy: _base.endIndex
+    ) ?? _base.endIndex
+    
+    self.computedStartIndex =
+      Index(_baseRange: _base.startIndex..<baseEnd)
   }
 }
 
-extension ChunkedCollection: Collection {
+extension ChunkedByCount: Collection {
   public struct Index {
-    public let base: Base.Index
+    @usableFromInline
+    internal let baseRange: Range<Base.Index>
 
     @usableFromInline
-    init(_base: Base.Index) {
-      self.base = _base
+    internal init(_baseRange: Range<Base.Index>) {
+      self.baseRange = _baseRange
     }
   }
-  
-  public var startIndex: Index { Index(_base: base.startIndex) }
-  public var endIndex: Index { Index(_base: base.endIndex) }
+
+  public var startIndex: Index { computedStartIndex }
+  public var endIndex: Index {
+    Index(_baseRange: base.endIndex..<base.endIndex)
+  }
 
   public subscript(i: Index) -> Element {
-    base[i.base..<index(after: i).base]
+    base[i.baseRange]
   }
 
   @inlinable
   public func index(after i: Index) -> Index {
     let baseIdx = base.index(
-      i.base, offsetBy: _chunkCount, limitedBy: base.endIndex)
-    return Index(_base: baseIdx ?? base.endIndex)
+      i.baseRange.upperBound, offsetBy: chunkCount,
+      limitedBy: base.endIndex
+    ) ?? base.endIndex
+    return Index(_baseRange: i.baseRange.upperBound..<baseIdx)
   }
 }
 
-extension ChunkedCollection.Index: Comparable {
+extension ChunkedByCount.Index: Comparable {
   @inlinable
-  public static func < (lhs: ChunkedCollection.Index,
-                        rhs: ChunkedCollection.Index) -> Bool {
-    lhs.base < rhs.base
+  public static func < (lhs: ChunkedByCount.Index,
+                        rhs: ChunkedByCount.Index) -> Bool {
+    lhs.baseRange.lowerBound < rhs.baseRange.lowerBound
   }
 }
 
-extension ChunkedCollection:
+extension ChunkedByCount:
   BidirectionalCollection, RandomAccessCollection
 where Base: RandomAccessCollection {
   @inlinable
   public func index(before i: Index) -> Index {
-    if i.base == base.endIndex {
-      let remainder = base.count%_chunkCount
+    var offset = chunkCount
+    if i.baseRange.lowerBound == base.endIndex {
+      let remainder = base.count%chunkCount
       if remainder != 0 {
-        return Index(_base: base.index(i.base, offsetBy: -remainder))
+        offset = remainder
       }
     }
-    return Index(_base: base.index(i.base, offsetBy: -_chunkCount))
+    
+    let baseIdx = base.index(
+      i.baseRange.lowerBound, offsetBy: -offset,
+      limitedBy: base.startIndex
+    ) ?? base.startIndex
+    return Index(_baseRange: baseIdx..<i.baseRange.lowerBound)
   }
 
   @inlinable
   public func distance(from start: Index, to end: Index) -> Int {
-    let distance = base.distance(from: start.base, to: end.base)
+    let distance =
+      base.distance(from: start.baseRange.lowerBound,
+                    to: end.baseRange.lowerBound)
     let (quotient, remainder) =
-      distance.quotientAndRemainder(dividingBy: _chunkCount)
+      distance.quotientAndRemainder(dividingBy: chunkCount)
     // Increment should account for negative distances.
     if remainder < 0 {
       return quotient - 1
     }
     return quotient + (remainder == 0 ? 0 : 1)
   }
-  
+
   @inlinable
-  public func index(_ i: Index, offsetBy n: Int) -> Index {
-    guard n != 0 else { return i }
-    
-    let bound = n > 0 ? base.endIndex : base.startIndex
-    return Index(_base:
-      base.index(i.base, offsetBy: n * _chunkCount, limitedBy: bound) ?? bound
-    )
+  public var count: Int {
+    let (quotient, remainder) =
+      base.count.quotientAndRemainder(dividingBy: chunkCount)
+    return quotient + (remainder == 0 ? 0 : 1)
   }
 }
 
@@ -367,12 +392,17 @@ extension Collection {
   ///
   /// - Complexity: O(1)
   @inlinable
-  public func chunks(ofCount count: Int) -> ChunkedCollection<Self> {
+  public func chunks(ofCount count: Int) -> ChunkedByCount<Self> {
     precondition(count > 0, " Cannot chunk with count <= 0!")
-    return ChunkedCollection(_base: self, _chunkCount: count)
+    return ChunkedByCount(_base: self, _chunkCount: count)
   }
 }
 
 // Conditional conformances.
-extension ChunkedCollection: Equatable where Base: Equatable {}
-extension ChunkedCollection: Hashable where Base: Hashable {}
+extension ChunkedByCount: Equatable where Base: Equatable {}
+extension ChunkedByCount: Hashable where Base: Hashable, Base.Index: Hashable {}
+extension ChunkedByCount.Index: Hashable where Base.Index: Hashable {}
+
+// Lazy conditional conformance.
+extension ChunkedByCount: LazySequenceProtocol where Base: LazySequenceProtocol {}
+extension ChunkedByCount: LazyCollectionProtocol where Base: LazyCollectionProtocol {}

Tests/SwiftAlgorithmsTests/ChunkedTests.swift

@@ -81,7 +81,7 @@ final class ChunkedTests: XCTestCase {
   //===----------------------------------------------------------------------===//
   // Tests for `chunks(ofCount:)`
   //===----------------------------------------------------------------------===//
-  func testChunksOfSize() {
+  func testChunksOfCount() {
     XCTAssertEqualSequences([Int]().chunks(ofCount: 1), [])
     XCTAssertEqualSequences([Int]().chunks(ofCount: 5), [])
 
@@ -96,7 +96,7 @@ final class ChunkedTests: XCTestCase {
                             [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
   }
 
-  func testChunksOfSizeBidirectional() {
+  func testChunksOfCountBidirectional() {
     let collection = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
 
     XCTAssertEqualSequences(collection.chunks(ofCount: 1).reversed(),
@@ -109,7 +109,7 @@ final class ChunkedTests: XCTestCase {
                             [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
   }
 
-  func testChunksOfSizeCount() {
+  func testChunksOfCountCount() {
     XCTAssertEqual([Int]().chunks(ofCount: 1).count, 0)
     XCTAssertEqual([Int]().chunks(ofCount: 5).count, 0)
 
@@ -119,62 +119,24 @@ final class ChunkedTests: XCTestCase {
     XCTAssertEqual(collection.chunks(ofCount: 5).count, 2)
     XCTAssertEqual(collection.chunks(ofCount: 11).count, 1)
   }
-  
-  func testChunksOfSizeDistance() {
+
+  func testEmptyChunksTraversal() {
     let emptyChunks = [Int]().chunks(ofCount: 1)
 
-    XCTAssertEqual(
-      emptyChunks.distance(from: emptyChunks.startIndex,
-                           to: emptyChunks.endIndex), 0)
-    
+    validateIndexTraversals(emptyChunks)
+  }
+  
+  func testChunksOfCountTraversal() {
     let collection = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
-    let chunks = collection.chunks(ofCount: 3)
-    
-    XCTAssertEqual(
-      chunks.distance(from: chunks.startIndex, to: chunks.endIndex), 4)
-    XCTAssertEqual(
-      chunks.distance(from: chunks.endIndex, to: chunks.startIndex), -4)
-    
-    var idx = chunks.startIndex
-    chunks.formIndex(after: &idx)
-    XCTAssertEqual(chunks.distance(from: chunks.startIndex, to: idx), 1)
-    XCTAssertEqual(chunks.distance(from: idx, to: chunks.startIndex), -1)
-
-    chunks.formIndex(after: &idx)
-    XCTAssertEqual(chunks.distance(from: chunks.startIndex, to: idx), 2)
-    XCTAssertEqual(chunks.distance(from: idx, to: chunks.startIndex), -2)
+    let chunks = collection.chunks(ofCount: 2)
+  
+    validateIndexTraversals(chunks)
   }
 
-  func testChunksOfSizeOffset() {
-    let emptyChunks = [Int]().chunks(ofCount: 1)
-    
-    XCTAssertEqual(
-      emptyChunks.index(emptyChunks.startIndex, offsetBy: 1),
-      emptyChunks.endIndex)
-    XCTAssertEqual(
-      emptyChunks.index(emptyChunks.startIndex, offsetBy: -1),
-      emptyChunks.endIndex)
-    
+  func testChunksOfCountWithRemainderTraversal() {
     let collection = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
     let chunks = collection.chunks(ofCount: 3)
 
-    XCTAssertEqual(chunks.index(chunks.startIndex, offsetBy: 1),
-                   chunks.index(after: chunks.startIndex))
-    
-    XCTAssertEqual(chunks.index(chunks.startIndex, offsetBy: 1),
-                   chunks.index(after: chunks.startIndex))
-    
-    var expected = chunks.startIndex
-    chunks.formIndex(after: &expected)
-    chunks.formIndex(after: &expected)
-
-    XCTAssertEqual(chunks.index(chunks.startIndex, offsetBy: 2),
-                   expected)
-    
-    // Offset is bounded
-    XCTAssertEqual(chunks.index(chunks.startIndex, offsetBy: 5),
-                   chunks.endIndex)
-    XCTAssertEqual(chunks.index(chunks.startIndex, offsetBy: -1),
-                   chunks.startIndex)
+    validateIndexTraversals(chunks)
   }
 }