Add regex-specific Matches and Ranges collections (swiftlang#460)

This prepares for adopting an opaque result type for matches(of:)
and ranges(of:). The old, CollectionConsumer-based model moves 
index-by-index, and isn't aware of the regex's semantic level, 
which results in inaccurate results for regexes that match at a 
mid-character index.

Author: natecook1000

Sources/_StringProcessing/Algorithms/Algorithms/Ranges.swift

@@ -226,17 +226,53 @@ extension BidirectionalCollection where Element: Comparable {
 //  }
 }
 
+@available(SwiftStdlib 5.7, *)
+struct RegexRangesCollection<Output> {
+  let base: RegexMatchesCollection<Output>
+  
+  init(string: Substring, regex: Regex<Output>) {
+    self.base = RegexMatchesCollection(base: string, regex: regex)
+  }
+}
+
+@available(SwiftStdlib 5.7, *)
+extension RegexRangesCollection: Sequence {
+  struct Iterator: IteratorProtocol {
+    var matchesBase: RegexMatchesCollection<Output>.Iterator
+    
+    mutating func next() -> Range<String.Index>? {
+      matchesBase.next().map(\.range)
+    }
+  }
+  
+  func makeIterator() -> Iterator {
+    Iterator(matchesBase: base.makeIterator())
+  }
+}
+
+@available(SwiftStdlib 5.7, *)
+extension RegexRangesCollection: Collection {
+  typealias Index = RegexMatchesCollection<Output>.Index
+
+  var startIndex: Index { base.startIndex }
+  var endIndex: Index { base.endIndex }
+  func index(after i: Index) -> Index { base.index(after: i) }
+  subscript(position: Index) -> Range<String.Index> { base[position].range }
+}
+
 // MARK: Regex algorithms
 
-extension BidirectionalCollection where SubSequence == Substring {
+extension Collection where SubSequence == Substring {
   @available(SwiftStdlib 5.7, *)
   @_disfavoredOverload
   func _ranges<R: RegexComponent>(
     of regex: R
-  ) -> RangesCollection<RegexConsumer<R, Self>> {
-    _ranges(of: RegexConsumer(regex))
+  ) -> RegexRangesCollection<R.RegexOutput> {
+    RegexRangesCollection(string: self[...], regex: regex.regex)
   }
+}
 
+extension BidirectionalCollection where SubSequence == Substring {
   @available(SwiftStdlib 5.7, *)
   func _rangesFromBack<R: RegexComponent>(
     of regex: R

Sources/_StringProcessing/Algorithms/Algorithms/Replace.swift

@@ -12,21 +12,21 @@
 // MARK: `CollectionSearcher` algorithms
 
 extension RangeReplaceableCollection {
-  func _replacing<Searcher: CollectionSearcher, Replacement: Collection>(
-    _ searcher: Searcher,
+  func _replacing<Ranges: Collection, Replacement: Collection>(
+    _ ranges: Ranges,
     with replacement: Replacement,
-    subrange: Range<Index>,
     maxReplacements: Int = .max
-  ) -> Self where Searcher.Searched == SubSequence,
+  ) -> Self where Ranges.Element == Range<Index>,
                   Replacement.Element == Element
   {
     precondition(maxReplacements >= 0)
 
-    var index = subrange.lowerBound
     var result = Self()
-    result.append(contentsOf: self[..<index])
+    var index = startIndex
 
-    for range in self[subrange]._ranges(of: searcher).prefix(maxReplacements) {
+    // `maxRanges` is a workaround for https://github.com/apple/swift/issues/59522
+    let maxRanges = ranges.prefix(maxReplacements)
+    for range in maxRanges {
       result.append(contentsOf: self[index..<range.lowerBound])
       result.append(contentsOf: replacement)
       index = range.upperBound
@@ -36,29 +36,15 @@ extension RangeReplaceableCollection {
     return result
   }
 
-  func _replacing<Searcher: CollectionSearcher, Replacement: Collection>(
-    _ searcher: Searcher,
-    with replacement: Replacement,
-    maxReplacements: Int = .max
-  ) -> Self where Searcher.Searched == SubSequence,
-                  Replacement.Element == Element
-  {
-    _replacing(
-      searcher,
-      with: replacement,
-      subrange: startIndex..<endIndex,
-      maxReplacements: maxReplacements)
-  }
-  
   mutating func _replace<
-    Searcher: CollectionSearcher, Replacement: Collection
+    Ranges: Collection, Replacement: Collection
   >(
-    _ searcher: Searcher,
+    _ ranges: Ranges,
     with replacement: Replacement,
     maxReplacements: Int = .max
-  ) where Searcher.Searched == SubSequence, Replacement.Element == Element {
+  ) where Ranges.Element == Range<Index>, Replacement.Element == Element {
     self = _replacing(
-      searcher,
+      ranges,
       with: replacement,
       maxReplacements: maxReplacements)
   }
@@ -85,9 +71,8 @@ extension RangeReplaceableCollection where Element: Equatable {
     maxReplacements: Int = .max
   ) -> Self where C.Element == Element, Replacement.Element == Element {
     _replacing(
-      ZSearcher(pattern: Array(other), by: ==),
+      self[subrange]._ranges(of: other),
       with: replacement,
-      subrange: subrange,
       maxReplacements: maxReplacements)
   }
 
@@ -143,9 +128,8 @@ extension RangeReplaceableCollection
     maxReplacements: Int = .max
   ) -> Self where C.Element == Element, Replacement.Element == Element {
     _replacing(
-      PatternOrEmpty(searcher: TwoWaySearcher(pattern: Array(other))),
+      self[subrange]._ranges(of: other),
       with: replacement,
-      subrange: subrange,
       maxReplacements: maxReplacements)
   }
 
@@ -195,9 +179,8 @@ extension RangeReplaceableCollection where SubSequence == Substring {
     maxReplacements: Int = .max
   ) -> Self where Replacement.Element == Element {
     _replacing(
-      RegexConsumer(regex),
+      self[subrange]._ranges(of: regex),
       with: replacement,
-      subrange: subrange,
       maxReplacements: maxReplacements)
   }
 

Sources/_StringProcessing/Algorithms/Matching/Matches.swift

@@ -183,13 +183,155 @@ extension BidirectionalCollection {
 
 // MARK: Regex algorithms
 
+@available(SwiftStdlib 5.7, *)
+struct RegexMatchesCollection<Output> {
+  let input: Substring
+  let regex: Regex<Output>
+  let startIndex: Index
+  
+  init(base: Substring, regex: Regex<Output>) {
+    self.input = base
+    self.regex = regex
+    self.startIndex = base.firstMatch(of: regex).map(Index.match) ?? .end
+  }
+}
+
+@available(SwiftStdlib 5.7, *)
+extension RegexMatchesCollection: Sequence {
+  /// Returns the index to start searching for the next match after `match`.
+  fileprivate func searchIndex(after match: Regex<Output>.Match) -> String.Index? {
+    if !match.range.isEmpty {
+      return match.range.upperBound
+    }
+    
+    // If the last match was an empty match, advance by one position and
+    // run again, unless at the end of `input`.
+    if match.range.lowerBound == input.endIndex {
+      return nil
+    }
+    
+    switch regex.initialOptions.semanticLevel {
+    case .graphemeCluster:
+      return input.index(after: match.range.upperBound)
+    case .unicodeScalar:
+      return input.unicodeScalars.index(after: match.range.upperBound)
+    }
+  }
+
+  struct Iterator: IteratorProtocol {
+    let base: RegexMatchesCollection
+    
+    // Because `RegexMatchesCollection` eagerly computes the first match for
+    // its `startIndex`, the iterator can use that match for its initial
+    // iteration. For subsequent calls to `next()`, this value is `false`, and
+    // `nextStart` is used to search for the next match.
+    var initialIteration = true
+    var nextStart: String.Index?
+    
+    init(_ matches: RegexMatchesCollection) {
+      self.base = matches
+      self.nextStart = base.startIndex.match.flatMap(base.searchIndex(after:))
+    }
+    
+    mutating func next() -> Regex<Output>.Match? {
+      // Initial case with pre-computed first match
+      if initialIteration {
+        initialIteration = false
+        return base.startIndex.match
+      }
+      
+      // `nextStart` is `nil` when iteration has completed
+      guard let start = nextStart else {
+        return nil
+      }
+      
+      // Otherwise, find the next match (if any) and compute `nextStart`
+      let match = try? base.regex.firstMatch(in: base.input[start...])
+      nextStart = match.flatMap(base.searchIndex(after:))
+      return match
+    }
+  }
+  
+  func makeIterator() -> Iterator {
+    Iterator(self)
+  }
+}
+
+@available(SwiftStdlib 5.7, *)
+extension RegexMatchesCollection: Collection {
+  enum Index: Comparable {
+    case match(Regex<Output>.Match)
+    case end
+    
+    var match: Regex<Output>.Match? {
+      switch self {
+      case .match(let match): return match
+      case .end: return nil
+      }
+    }
+    
+    static func == (lhs: Self, rhs: Self) -> Bool {
+      switch (lhs, rhs) {
+      case (.match(let lhs), .match(let rhs)):
+        return lhs.range == rhs.range
+      case (.end, .end):
+        return true
+      case (.end, .match), (.match, .end):
+        return false
+      }
+    }
+    
+    static func < (lhs: Self, rhs: Self) -> Bool {
+      switch (lhs, rhs) {
+      case (.match(let lhs), .match(let rhs)):
+        // This implementation uses a tuple comparison so that an empty
+        // range `i..<i` will be ordered before a non-empty range at that
+        // same starting point `i..<j`. As of 2022-05-30, `Regex` does not
+        // return matches of this kind, but that is one behavior under
+        // discussion for regexes like /a*|b/ when matched against "b".
+        return (lhs.range.lowerBound, lhs.range.upperBound)
+          < (rhs.range.lowerBound, rhs.range.upperBound)
+      case (.match, .end):
+        return true
+      case (.end, .match), (.end, .end):
+        return false
+      }
+    }
+  }
+  
+  var endIndex: Index {
+    Index.end
+  }
+  
+  func index(after i: Index) -> Index {
+    guard let currentMatch = i.match else {
+      fatalError("Can't advance past the 'endIndex' of a match collection.")
+    }
+    
+    guard
+      let start = searchIndex(after: currentMatch),
+      let nextMatch = try? regex.firstMatch(in: input[start...])
+    else {
+      return .end
+    }
+    return Index.match(nextMatch)
+  }
+  
+  subscript(position: Index) -> Regex<Output>.Match {
+    guard let match = position.match else {
+      fatalError("Can't subscript the 'endIndex' of a match collection.")
+    }
+    return match
+  }
+}
+
 extension BidirectionalCollection where SubSequence == Substring {
   @available(SwiftStdlib 5.7, *)
   @_disfavoredOverload
   func _matches<R: RegexComponent>(
     of regex: R
-  ) -> MatchesCollection<RegexConsumer<R, Self>> {
-    _matches(of: RegexConsumer(regex))
+  ) -> RegexMatchesCollection<R.RegexOutput> {
+    RegexMatchesCollection(base: self[...], regex: regex.regex)
   }
 
   @available(SwiftStdlib 5.7, *)
@@ -207,30 +349,6 @@ extension BidirectionalCollection where SubSequence == Substring {
   public func matches<Output>(
     of r: some RegexComponent<Output>
   ) -> [Regex<Output>.Match] {
-    let slice = self[...]
-    var start = self.startIndex
-    let end = self.endIndex
-    let regex = r.regex
-
-    var result = [Regex<Output>.Match]()
-    while start <= end {
-      guard let match = try? regex._firstMatch(
-        slice.base, in: start..<end
-      ) else {
-        break
-      }
-      result.append(match)
-      if match.range.isEmpty {
-        if match.range.upperBound == end {
-          break
-        }
-        // FIXME: semantic level
-        start = slice.index(after: match.range.upperBound)
-      } else {
-        start = match.range.upperBound
-      }
-    }
-    return result
+    Array(_matches(of: r))
   }
-
 }

Tests/RegexTests/AlgorithmsInternalsTests.swift

@@ -44,4 +44,39 @@ extension AlgorithmTests {
     XCTAssertEqual("x", "axb"._trimming(r))
     XCTAssertEqual("x", "axbb"._trimming(r))
   }
+  
+  func testMatchesCollection() {
+    let r = try! Regex("a|b+|c*", as: Substring.self)
+    
+    let str = "zaabbbbbbcde"
+    let matches = str._matches(of: r)
+    let expected: [Substring] = [
+      "", // before 'z'
+      "a",
+      "a",
+      "bbbbbb",
+      "c",
+      "", // after 'c'
+      "", // after 'd'
+      "", // after 'e'
+    ]
+
+    // Make sure we're getting the right collection type
+    let _: RegexMatchesCollection<Substring> = matches
+
+    XCTAssertEqual(matches.map(\.output), expected)
+    
+    let i = matches.index(matches.startIndex, offsetBy: 3)
+    XCTAssertEqual(matches[i].output, expected[3])
+    let j = matches.index(i, offsetBy: 5)
+    XCTAssertEqual(j, matches.endIndex)
+    
+    var index = matches.startIndex
+    while index < matches.endIndex {
+      XCTAssertEqual(
+        matches[index].output,
+        expected[matches.distance(from: matches.startIndex, to: index)])
+      matches.formIndex(after: &index)
+    }
+  }
 }