Skip to content

Commit a27d9c4

Browse files
glessardglessard
authored
Merge pull request #71912 from oscbyspro/better-joined-distance-from-to-int-min
FlattenSequence/distance(from:to:) untrapping Int.min
2 parents 79d2c8b + af0c9fa commit a27d9c4

File tree

2 files changed

+224
-22
lines changed

2 files changed

+224
-22
lines changed

‎stdlib/public/core/Flatten.swift

Lines changed: 44 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -297,41 +297,63 @@ extension FlattenCollection: Collection {
297297

298298
@inlinable // lazy-performance
299299
public func distance(from start: Index, to end: Index) -> Int {
300+
let distanceIsNegative = start > end
301+
300302
// The following check ensures that distance(from:to:) is invoked on
301303
// the _base at least once, to trigger a _precondition in forward only
302304
// collections.
303-
if start > end {
305+
if distanceIsNegative {
304306
_ = _base.distance(from: _base.endIndex, to: _base.startIndex)
305307
}
306-
307-
// This handles indices belonging to the same collection.
308+
309+
// This path handles indices belonging to the same collection.
308310
if start._outer == end._outer {
309311
guard let i = start._inner, let j = end._inner else { return 0 }
310312
return _base[start._outer].distance(from: i, to: j)
311313
}
312-
313-
// The following combines the distance of three sections.
314-
let range = start <= end ? start ..< end : end ..< start
315-
var outer = range.lowerBound._outer
316-
var count = 0 as Int // 0...Int.max
317-
318-
if let inner = range.lowerBound._inner {
319-
let collection = _base[outer]
320-
count += collection.distance(from: inner, to: collection.endIndex)
321-
_base.formIndex(after: &outer)
314+
315+
// The following path combines the distances of three regions.
316+
let lowerBound: Index
317+
let upperBound: Index
318+
319+
let step: Int
320+
var distance: Int
321+
322+
// Note that lowerBound is a valid index because start != end.
323+
if distanceIsNegative {
324+
step = -1
325+
lowerBound = end
326+
upperBound = start
327+
let lowest = _base[lowerBound._outer]
328+
distance = lowest.distance(from: lowest.endIndex, to: lowerBound._inner!)
329+
} else {
330+
step = 01
331+
lowerBound = start
332+
upperBound = end
333+
let lowest = _base[lowerBound._outer]
334+
distance = lowest.distance(from: lowerBound._inner!, to: lowest.endIndex)
322335
}
323-
324-
while outer < range.upperBound._outer {
325-
count += _base[outer].count
336+
337+
// We can use each collection's count in the middle region since the
338+
// fast path ensures that the other regions cover a nonzero distance,
339+
// which means that an extra Int.min distance should trap regardless.
340+
var outer = _base.index(after: lowerBound._outer)
341+
while outer < upperBound._outer {
342+
// 0 ... Int.max can always be negated.
343+
distance += _base[outer].count &* step
326344
_base.formIndex(after: &outer)
327345
}
328-
329-
if let inner = range.upperBound._inner {
330-
let collection = _base[outer]
331-
count += collection.distance(from: collection.startIndex, to: inner)
346+
347+
// This unwraps if start != endIndex and end != endIndex. We can use the
348+
// positive distance for the same reason that we can use the collection's
349+
// count in the middle region.
350+
if let inner = upperBound._inner {
351+
// 0 ... Int.max can always be negated.
352+
let highest = _base[upperBound._outer]
353+
distance += highest.distance(from: highest.startIndex, to: inner) &* step
332354
}
333-
334-
return start <= end ? count : -count
355+
356+
return distance
335357
}
336358

337359
@inline(__always)

‎test/stdlib/FlattenDistanceFromTo.swift

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
//===----------------------------------------------------------------------===//
2+
//
3+
// This source file is part of the Swift.org open source project
4+
//
5+
// Copyright (c) 2024 Apple Inc. and the Swift project authors
6+
// Licensed under Apache License v2.0 with Runtime Library Exception
7+
//
8+
// See https://swift.org/LICENSE.txt for license information
9+
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
10+
//
11+
//===----------------------------------------------------------------------===//
12+
//
13+
// RUN: %target-run-simple-swift(-parse-as-library)
14+
// REQUIRES: executable_test
15+
// END.
16+
//
17+
//===----------------------------------------------------------------------===//
18+
19+
import StdlibUnittest
20+
21+
@main
22+
final class FlattenDistanceFromToTests {
23+
24+
static func main() {
25+
let tests = FlattenDistanceFromToTests()
26+
let suite = TestSuite("FlattenDistanceFromToTests")
27+
suite.test("EachIndexPair", tests.testEachIndexPair)
28+
if #available(SwiftStdlib 6.0, *) {
29+
// The random access time complexity was fixed in Swift 6.0.
30+
suite.test("MinMaxRandomAccess", tests.testMinMaxRandomAccess)
31+
}
32+
runAllTests()
33+
}
34+
}
35+
36+
//===----------------------------------------------------------------------===//
37+
// MARK: - Each Index Pair
38+
//===----------------------------------------------------------------------===//
39+
40+
extension FlattenDistanceFromToTests {
41+
42+
/// Performs one `action` per lane size case through `limits`.
43+
///
44+
/// limits: [0,1,2,3]
45+
/// ─────────────────
46+
/// [][ ][ ][ ]
47+
/// [][1][ ][ ]
48+
/// [][ ][2 ][ ]
49+
/// [][1][ ][ ]
50+
/// [][ ][2,2][ ]
51+
/// [][1][2,2][ ]
52+
/// [][ ][ ][3 ]
53+
/// ─────────────────
54+
/// [][1][2,2][3,3,3]
55+
///
56+
private func forEachLaneSizeCase(
57+
through limits: [Int],
58+
perform action: ([[Int]]) -> Void
59+
) {
60+
var array = Array(repeating: [Int](), count: limits.count)
61+
var index = array.startIndex
62+
while index < limits.endIndex {
63+
action(array)
64+
65+
if array[index].count < limits[index] {
66+
array[index].append(index)
67+
} else {
68+
while index < limits.endIndex, array[index].count == limits[index] {
69+
array.formIndex(after: &index)
70+
}
71+
72+
if index < limits.endIndex {
73+
array[index].append(index)
74+
75+
while index > array.startIndex {
76+
array.formIndex(before: &index)
77+
array[index].removeAll(keepingCapacity: true)
78+
}
79+
}
80+
}
81+
}
82+
}
83+
84+
/// Performs one `action` per offset-index pair in `collection`.
85+
///
86+
/// collection: [[0],[1,2]].joined()
87+
/// ────────────────────────────────
88+
/// offset: 0, index: 0,0
89+
/// offset: 1, index: 1,0
90+
/// offset: 2, index: 1,1
91+
/// offset: 3, index: 2
92+
///
93+
private func forEachEnumeratedIndexIncludingEndIndex<T>(
94+
in collection: T,
95+
perform action: ((offset: Int, index: T.Index)) -> Void
96+
) where T: Collection {
97+
var state = (offset: 0, index: collection.startIndex)
98+
while true {
99+
action(state)
100+
101+
if state.index == collection.endIndex {
102+
return
103+
}
104+
105+
state.offset += 1
106+
collection.formIndex(after: &state.index)
107+
}
108+
}
109+
110+
/// Checks the distance between each index pair in various cases.
111+
///
112+
/// You need three lanes to exercise the first, the middle, the last region.
113+
/// The past-the-end index is a separate lane, so the middle region contains
114+
/// one additional lane when the past-the-end index is selected.
115+
///
116+
func testEachIndexPair() {
117+
var invocations = 0 as Int
118+
119+
for lanes in 0 ... 3 {
120+
let limits = Array(repeating: 3, count: lanes)
121+
122+
forEachLaneSizeCase(through: limits) { base in
123+
let collection: FlattenSequence = base.joined()
124+
125+
forEachEnumeratedIndexIncludingEndIndex(in: collection) { start in
126+
forEachEnumeratedIndexIncludingEndIndex(in: collection) { end in
127+
let pair = (from: start.offset, to: end.offset)
128+
129+
invocations += 1
130+
131+
expectEqual(
132+
collection.distance(from: start.index, to: end.index),
133+
end.offset - start.offset,
134+
"""
135+
index distance != offset distance for \(pair) in \(base).joined()
136+
"""
137+
)
138+
}
139+
}
140+
}
141+
}
142+
143+
expectEqual(invocations, 2502, "unexpected workload")
144+
}
145+
}
146+
147+
//===----------------------------------------------------------------------===//
148+
// MARK: - Min Max Outputs
149+
//===----------------------------------------------------------------------===//
150+
151+
extension FlattenDistanceFromToTests {
152+
153+
/// Checks some `Int.min` and `Int.max` distances with random access.
154+
///
155+
/// It needs Swift 6.0+ because prior versions find the distance by
156+
/// iterating from one index to the other, which takes way too long.
157+
///
158+
/// - Note: A distance of `Int.min` requires more than `Int.max` elements.
159+
///
160+
@available(SwiftStdlib 6.0, *)
161+
func testMinMaxRandomAccess() {
162+
for s: FlattenSequence in [
163+
164+
[-1..<Int.max/1],
165+
[00..<Int.max/1, 00..<000000001],
166+
[00..<000000001, 01..<Int.max/1, 00..<000000001],
167+
[00..<000000001, 00..<Int.max/2, 00..<Int.max/2, 00..<000000001]
168+
169+
].map({ $0.joined() }) {
170+
171+
let a = s.startIndex, b = s.endIndex
172+
173+
expectEqual(Int.min, s.distance(from: b, to: s.index(a, offsetBy: 00)))
174+
expectEqual(Int.max, s.distance(from: s.index(a, offsetBy: 01), to: b))
175+
176+
expectEqual(Int.min, s.distance(from: s.index(b, offsetBy: 00), to: a))
177+
expectEqual(Int.max, s.distance(from: a, to: s.index(b, offsetBy: -1)))
178+
}
179+
}
180+
}

0 commit comments

Comments
 (0)