Skip to content

Commit 11c0231

Browse files
committed
FlattenSequence/distance(from:to:) enhancements.
I was looking through Swift's issues tab and found this issue (#57496), which points out the poor performance of `FlattenSequence/count`. The current `count` implementation calls `distance(from:to:)`, which iterates through the entire collection, so I thought I'd improve that method. The new version computes the distance as the sum of three parts, with a fast path when both indices belong to the same underlying collection. Note that it does this by calling either `count` or `distance(from:to:)` on each underlying collection, which makes `[repeatElement(0, count: Int.max)].joined().count` instant, for example.
1 parent 823db1f commit 11c0231

File tree

1 file changed

+24
-17
lines changed

1 file changed

+24
-17
lines changed

stdlib/public/core/Flatten.swift

Lines changed: 24 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -292,28 +292,35 @@ extension FlattenCollection: Collection {
292292
// The following check makes sure that distance(from:to:) is invoked on the
293293
// _base at least once, to trigger a _precondition in forward only
294294
// collections.
295-
if end < start {
295+
if start > end {
296296
_ = _base.distance(from: _base.endIndex, to: _base.startIndex)
297297
}
298-
var _start: Index
299-
let _end: Index
300-
let step: Int
301-
if start > end {
302-
_start = end
303-
_end = start
304-
step = -1
298+
299+
// This handles the case where both indices belong to the same collection.
300+
if start._outer == end._outer {
301+
return start._outer < _base.endIndex ? _base[start._outer].distance(from: start._inner!, to: end._inner!) : Int.zero
305302
}
306-
else {
307-
_start = start
308-
_end = end
309-
step = 1
303+
304+
// The following path combines the prefix, the middle, and the suffix distance.
305+
let range = Range(uncheckedBounds: start <= end ? (start, end) : (end, start))
306+
var outer = range.lowerBound._outer
307+
var count = 0 as Int // 0...Int.max
308+
309+
if let inner = range.lowerBound._inner {
310+
count += _base[outer][inner...].count
311+
_base.formIndex(after: &outer)
312+
}
313+
314+
while outer < range.upperBound._outer {
315+
count += _base[outer].count
316+
_base.formIndex(after: &outer)
310317
}
311-
var count = 0
312-
while _start != _end {
313-
count += step
314-
formIndex(after: &_start)
318+
319+
if let inner = range.upperBound._inner {
320+
count += _base[outer][..<inner].count
315321
}
316-
return count
322+
323+
return start <= end ? count : -count
317324
}
318325

319326
@inline(__always)

0 commit comments

Comments
 (0)