Track source locations for alternation pipes

Author: hamishknight

Sources/_MatchingEngine/Regex/AST/AST.swift

@@ -92,11 +92,20 @@ extension AST {
 
   public struct Alternation: Hashable, _ASTNode {
     public let children: [AST]
-    public let location: SourceLocation
+    public let pipes: [SourceLocation]
+
+    public init(_ mems: [AST], pipes: [SourceLocation]) {
+      // An alternation must have at least two branches (though the branches
+      // may be empty AST nodes), and n - 1 pipes.
+      precondition(mems.count >= 2)
+      precondition(pipes.count == mems.count - 1)
 
-    public init(_ mems: [AST], _ location: SourceLocation) {
       self.children = mems
-      self.location = location
+      self.pipes = pipes
+    }
+
+    public var location: SourceLocation {
+      .init(children.first!.location.start ..< children.last!.location.end)
     }
   }
 

Sources/_MatchingEngine/Regex/Parse/Parse.swift

@@ -89,17 +89,20 @@ extension Parser {
 
     if source.isEmpty { return .empty(.init(loc(_start))) }
 
-    var result = Array<AST>(singleElement: try parseConcatenation())
-    while source.tryEat("|") {
-      // TODO: track pipe locations too...
+    var result = [try parseConcatenation()]
+    var pipes: [SourceLocation] = []
+    while true {
+      let pipeStart = source.currentPosition
+      guard source.tryEat("|") else { break }
+      pipes.append(loc(pipeStart))
       result.append(try parseConcatenation())
     }
 
     if result.count == 1 {
       return result[0]
     }
 
-    return .alternation(.init(result, loc(_start)))
+    return .alternation(.init(result, pipes: pipes))
   }
 
   /// Parse a term, potentially separated from others by `|`

Sources/_StringProcessing/ASTBuilder.swift

@@ -17,7 +17,9 @@ AST.
 import _MatchingEngine
 
 func alt(_ asts: [AST]) -> AST {
-  .alternation(.init(asts, .fake))
+  return .alternation(
+    .init(asts, pipes: Array(repeating: .fake, count: asts.count - 1))
+  )
 }
 func alt(_ asts: AST...) -> AST {
   alt(asts)

Tests/RegexTests/ParseTests.swift

@@ -928,6 +928,11 @@ extension RegexTests {
       rangeTest(alt, entireRange)
       rangeTest("(\(alt))", insetRange(by: 1), at: \.children![0].location)
     }
+
+    rangeTest("|", entireRange, at: { $0.as(Alt.self)!.pipes[0] })
+    rangeTest("a|", range(1 ..< 2), at: { $0.as(Alt.self)!.pipes[0] })
+    rangeTest("a|b", range(1 ..< 2), at: { $0.as(Alt.self)!.pipes[0] })
+    rangeTest("|||", range(1 ..< 2), at: { $0.as(Alt.self)!.pipes[1] })
   }
 
   func testParseErrors() {