Skip to content

Commit 8fe1c9d

Browse files
[Chunked] Code review adjustments
1 parent 05df988 commit 8fe1c9d

File tree

1 file changed

+29
-25
lines changed

1 file changed

+29
-25
lines changed

Sources/Algorithms/Chunked.swift

Lines changed: 29 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,7 @@ public struct ChunkedByCount<Base: Collection> {
271271
internal let chunkCount: Int
272272

273273
@usableFromInline
274-
internal var computedStartIndex: Index
274+
internal var computedStartEnd: Base.Index
275275

276276
/// Creates a view instance that presents the elements of `base`
277277
/// in `SubSequence` chunks of the given count.
@@ -284,13 +284,10 @@ public struct ChunkedByCount<Base: Collection> {
284284

285285
// Compute the start index upfront in order to make
286286
// start index a O(1) lookup.
287-
let baseEnd = _base.index(
287+
self.computedStartEnd = _base.index(
288288
_base.startIndex, offsetBy: _chunkCount,
289289
limitedBy: _base.endIndex
290290
) ?? _base.endIndex
291-
292-
self.computedStartIndex =
293-
Index(_baseRange: _base.startIndex..<baseEnd)
294291
}
295292
}
296293

@@ -307,7 +304,9 @@ extension ChunkedByCount: Collection {
307304

308305
/// - Complexity: O(1)
309306
@inlinable
310-
public var startIndex: Index { computedStartIndex }
307+
public var startIndex: Index {
308+
Index(_baseRange: base.startIndex..<computedStartEnd)
309+
}
311310
@inlinable
312311
public var endIndex: Index {
313312
Index(_baseRange: base.endIndex..<base.endIndex)
@@ -331,6 +330,12 @@ extension ChunkedByCount: Collection {
331330
}
332331

333332
extension ChunkedByCount.Index: Comparable {
333+
@inlinable
334+
public static func == (lhs: ChunkedByCount.Index,
335+
rhs: ChunkedByCount.Index) -> Bool {
336+
lhs.baseRange.lowerBound == rhs.baseRange.lowerBound
337+
}
338+
334339
@inlinable
335340
public static func < (lhs: ChunkedByCount.Index,
336341
rhs: ChunkedByCount.Index) -> Bool {
@@ -391,47 +396,46 @@ where Base: RandomAccessCollection {
391396

392397
@usableFromInline
393398
internal func offsetForward(_ i: Index, offsetBy distance: Int) -> Index {
399+
assert(distance > 0)
394400
return makeOffsetIndex(
395-
from: i, baseBound: base.endIndex, distance: distance
401+
from: i, baseBound: base.endIndex, baseDistance: distance * chunkCount
396402
)
397403
}
398404

399405
@usableFromInline
400406
internal func offsetBackward(_ i: Index, offsetBy distance: Int) -> Index {
401-
var idx = i
402-
var distance = distance
403-
// If we know that the last chunk is the only one that can possible
404-
// have a variadic count. So in order to simplify and avoid another
405-
// calculation of offsets(that is already done at `index(before:)`)
406-
// we just move one position already so the index can be calculated
407-
// since all remaining chunks have the same size.
407+
assert(distance < 0)
408+
// Distance "minus" one(at this point distance is negative) because we
409+
// need to adjust for the last position that may have variadic(remainder)
410+
// number of elements.
411+
var baseDistance = (distance + 1) * chunkCount
408412
if i.baseRange.lowerBound == base.endIndex {
409-
formIndex(before: &idx)
410-
distance += 1
411-
// If the offset was simply one, we are done.
412-
guard distance != 0 else {
413-
return idx
413+
let remainder = base.count%chunkCount
414+
// We have to take it into account when calculating offsets.
415+
if remainder != 0 {
416+
baseDistance -= remainder
417+
return makeOffsetIndex(
418+
from: i, baseBound: base.startIndex, baseDistance: baseDistance
419+
)
414420
}
415421
}
416-
417422
return makeOffsetIndex(
418-
from: idx, baseBound: base.startIndex, distance: distance
423+
from: i, baseBound: base.startIndex, baseDistance: baseDistance - chunkCount
419424
)
420425
}
421426

422427
// Helper to compute index(offsetBy:) index.
423428
@inline(__always)
424429
private func makeOffsetIndex(
425-
from i: Index, baseBound: Base.Index, distance: Int
430+
from i: Index, baseBound: Base.Index, baseDistance: Int
426431
) -> Index {
427432
let baseStartIdx = base.index(
428-
i.baseRange.lowerBound, offsetBy: distance * chunkCount,
433+
i.baseRange.lowerBound, offsetBy: baseDistance,
429434
limitedBy: baseBound
430435
) ?? baseBound
431436

432437
let baseEndIdx = base.index(
433-
i.baseRange.lowerBound, offsetBy: (distance + 1) * chunkCount,
434-
limitedBy: base.endIndex
438+
baseStartIdx, offsetBy: chunkCount, limitedBy: base.endIndex
435439
) ?? base.endIndex
436440

437441
return Index(_baseRange: baseStartIdx..<baseEndIdx)

0 commit comments

Comments
 (0)