FlattenSequence/distance(from:to:) untrapping Int.min

stdlib/public/core/Flatten.swift

@@ -289,41 +289,63 @@ extension FlattenCollection: Collection {
 
   @inlinable // lazy-performance
   public func distance(from start: Index, to end: Index) -> Int {
+    let distanceIsNegative = start > end
+
     // The following check ensures that distance(from:to:) is invoked on
     // the _base at least once, to trigger a _precondition in forward only
     // collections.
-    if start > end {
+    if distanceIsNegative {
       _ = _base.distance(from: _base.endIndex, to: _base.startIndex)
     }
-    
-    // This handles indices belonging to the same collection.
+
+    // This path handles indices belonging to the same collection.
     if start._outer == end._outer {
       guard let i = start._inner, let j = end._inner else { return 0 }
       return _base[start._outer].distance(from: i, to: j)
     }
-
-    // The following combines the distance of three sections.
-    let range = start <= end ? start ..< end : end ..< start
-    var outer = range.lowerBound._outer
-    var count = 0 as Int // 0...Int.max
-
-    if let inner = range.lowerBound._inner {
-      let collection = _base[outer]
-      count += collection.distance(from: inner, to: collection.endIndex)
-      _base.formIndex(after: &outer)
+
+    // The following path combines the distances of three regions.
+    let lowerBound: Index
+    let upperBound: Index
+
+    let step: Int
+    var distance: Int
+
+    // Note that lowerBound is a valid index because start != end.
+    if distanceIsNegative {
+      step = -1
+      lowerBound = end
+      upperBound = start
+      let lowest = _base[lowerBound._outer]
+      distance = lowest.distance(from: lowest.endIndex, to: lowerBound._inner!)
+    } else {
+      step = 01
+      lowerBound = start
+      upperBound = end
+      let lowest = _base[lowerBound._outer]
+      distance = lowest.distance(from: lowerBound._inner!, to: lowest.endIndex)
     }
-
-    while outer < range.upperBound._outer {
-      count += _base[outer].count
+
+    // We can use each collection's count in the middle region since the
+    // fast path ensures that the other regions cover a nonzero distance,
+    // which means that an extra Int.min distance should trap regardless.
+    var outer = _base.index(after: lowerBound._outer)
+    while outer < upperBound._outer {
+      // 0 ... Int.max can always be negated.
+      distance += _base[outer].count &* step
       _base.formIndex(after: &outer)
     }
-
-    if let inner = range.upperBound._inner {
-      let collection = _base[outer]
-      count += collection.distance(from: collection.startIndex, to: inner)
+
+    // This unwraps if start != endIndex and end != endIndex. We can use the
+    // positive distance for the same reason that we can use the collection's
+    // count in the middle region.
+    if let inner = upperBound._inner {
+      // 0 ... Int.max can always be negated.
+      let highest = _base[upperBound._outer]
+      distance += highest.distance(from: highest.startIndex, to: inner) &* step
     }
-    
-    return start <= end ? count : -count
+
+    return distance
   }
 
   @inline(__always)

test/stdlib/FlattenDistanceFromTo.swift

@@ -0,0 +1,180 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// RUN: %target-run-simple-swift(-parse-as-library)
+// REQUIRES: executable_test
+// END.
+//
+//===----------------------------------------------------------------------===//
+
+import StdlibUnittest
+
+@main
+final class FlattenDistanceFromToTests {
+
+  static func main() {
+    let tests = FlattenDistanceFromToTests()
+    let suite = TestSuite("FlattenDistanceFromToTests")
+    suite.test("EachIndexPair", tests.testEachIndexPair)
+    if #available(SwiftStdlib 6.0, *) {
+      // The random access time complexity was fixed in Swift 6.0.
+      suite.test("MinMaxRandomAccess", tests.testMinMaxRandomAccess)
+    }
+    runAllTests()
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// MARK: - Each Index Pair
+//===----------------------------------------------------------------------===//
+
+extension FlattenDistanceFromToTests {
+
+  /// Performs one `action` per lane size case through `limits`.
+  ///
+  ///     limits: [0,1,2,3]
+  ///     ─────────────────
+  ///     [][ ][   ][     ]
+  ///     [][1][   ][     ]
+  ///     [][ ][2  ][     ]
+  ///     [][1][   ][     ]
+  ///     [][ ][2,2][     ]
+  ///     [][1][2,2][     ]
+  ///     [][ ][   ][3    ]
+  ///     ─────────────────
+  ///     [][1][2,2][3,3,3]
+  ///
+  private func forEachLaneSizeCase(
+    through limits: [Int],
+    perform action: ([[Int]]) -> Void
+  ) {
+    var array = Array(repeating: [Int](), count: limits.count)
+    var index = array.startIndex
+    while index < limits.endIndex {
+      action(array)
+
+      if array[index].count < limits[index] {
+        array[index].append(index)
+      } else {
+        while index < limits.endIndex, array[index].count == limits[index] {
+          array.formIndex(after: &index)
+        }
+
+        if index < limits.endIndex {
+          array[index].append(index)
+
+          while index > array.startIndex {
+            array.formIndex(before: &index)
+            array[index].removeAll(keepingCapacity: true)
+          }
+        }
+      }
+    }
+  }
+
+  /// Performs one `action` per offset-index pair in `collection`.
+  ///
+  ///     collection: [[0],[1,2]].joined()
+  ///     ────────────────────────────────
+  ///     offset: 0, index: 0,0
+  ///     offset: 1, index: 1,0
+  ///     offset: 2, index: 1,1
+  ///     offset: 3, index: 2
+  ///
+  private func forEachEnumeratedIndexIncludingEndIndex<T>(
+    in collection: T,
+    perform action: ((offset: Int, index: T.Index)) -> Void
+  ) where T: Collection {
+    var state = (offset: 0, index: collection.startIndex)
+    while true {
+      action(state)
+
+      if state.index == collection.endIndex {
+        return
+      }
+
+      state.offset += 1
+      collection.formIndex(after: &state.index)
+    }
+  }
+
+  /// Checks the distance between each index pair in various cases.
+  ///
+  /// You need three lanes to exercise the first, the middle, the last region.
+  /// The past-the-end index is a separate lane, so the middle region contains
+  /// one additional lane when the past-the-end index is selected.
+  ///
+  func testEachIndexPair() {
+    var invocations = 0 as Int
+
+    for lanes in 0 ... 3 {
+      let limits = Array(repeating: 3, count: lanes)
+
+      forEachLaneSizeCase(through: limits) { base in
+        let collection: FlattenSequence = base.joined()
+
+        forEachEnumeratedIndexIncludingEndIndex(in: collection) { start in
+          forEachEnumeratedIndexIncludingEndIndex(in: collection) { end in
+            let pair = (from: start.offset, to: end.offset)
+
+            invocations += 1
+
+            expectEqual(
+              collection.distance(from: start.index, to: end.index),
+              end.offset - start.offset,
+              """
+              index distance != offset distance for \(pair) in \(base).joined()
+              """
+            )
+          }
+        }
+      }
+    }
+
+    expectEqual(invocations, 2502, "unexpected workload")
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// MARK: - Min Max Outputs
+//===----------------------------------------------------------------------===//
+
+extension FlattenDistanceFromToTests {
+
+  /// Checks some `Int.min` and `Int.max` distances with random access.
+  ///
+  /// It needs Swift 6.0+ because prior versions find the distance by
+  /// iterating from one index to the other, which takes way too long.
+  ///
+  /// - Note: A distance of `Int.min` requires more than `Int.max` elements.
+  ///
+  @available(SwiftStdlib 6.0, *)
+  func testMinMaxRandomAccess() {
+    for s: FlattenSequence in [
+
+      [-1..<Int.max/1],
+      [00..<Int.max/1, 00..<000000001],
+      [00..<000000001, 01..<Int.max/1, 00..<000000001],
+      [00..<000000001, 00..<Int.max/2, 00..<Int.max/2, 00..<000000001]
+
+    ].map({ $0.joined() }) {
+
+      let a = s.startIndex, b = s.endIndex
+
+      expectEqual(Int.min, s.distance(from: b, to: s.index(a, offsetBy: 00)))
+      expectEqual(Int.max, s.distance(from: s.index(a, offsetBy: 01), to: b))
+
+      expectEqual(Int.min, s.distance(from: s.index(b, offsetBy: 00), to: a))
+      expectEqual(Int.max, s.distance(from: a, to: s.index(b, offsetBy: -1)))
+    }
+  }
+}