Allow subcripting a match result on a Reference.

To simplify the call site, I chose to precondition instead of returning optionals.

Examples:
```
let a = Reference()
let result = regex.match {
  tryCapture(as: a) {
    oneOrMore(.digit)
  } transform: { Int($0) }
}
result[a] // => Any
result[a, as: Int.self] // => Int
```

Author: rxwei

Sources/_StringProcessing/Capture.swift

@@ -31,6 +31,19 @@ extension Capture {
   }
 }
 
+extension Capture {
+  subscript(_ index: Int) -> Capture {
+    if case .tuple(let elements) = self {
+      assert(
+        (elements.startIndex..<elements.endIndex).contains(index),
+        "Capture index out of bounds"
+      )
+      return elements[index]
+    }
+    return self
+  }
+}
+
 extension Capture {
   static func tupleOrAtom(_ elements: [Capture]) -> Self {
     elements.count == 1 ? elements[0] : .tuple(elements)

Sources/_StringProcessing/Engine/Consume.swift

@@ -63,7 +63,8 @@ extension Engine where Input == String {
     guard let result = result else { return nil }
 
     let capList = cpu.storedCaptures
-    return (result, CaptureList(caps: capList))
+    return (result, CaptureList(
+      values: capList, referencedCaptureOffsets: program.referencedCaptureOffsets))
   }
 }
 

Sources/_StringProcessing/Engine/MEBuilder.swift

@@ -44,7 +44,7 @@ extension MEProgram where Input.Element: Hashable {
 
     // Symbolic reference resolution
     var unresolvedReferences: [Reference.ID: [InstructionAddress]] = [:]
-    var captureOffsets: [Reference.ID: Int] = [:]
+    var referencedCaptureOffsets: [Reference.ID: Int] = [:]
     var captureCount: Int {
       // We currently deduce the capture count from the capture register number.
       nextCaptureRegister.rawValue
@@ -352,7 +352,8 @@ extension MEProgram.Builder {
       staticTransformFunctions: transformFunctions,
       staticMatcherFunctions: matcherFunctions,
       registerInfo: regInfo,
-      captureStructure: captureStructure)
+      captureStructure: captureStructure,
+      referencedCaptureOffsets: referencedCaptureOffsets)
   }
 
   public mutating func reset() { self = Self() }
@@ -424,7 +425,7 @@ extension MEProgram.Builder {
 fileprivate extension MEProgram.Builder {
   mutating func resolveReferences() throws {
     for (id, uses) in unresolvedReferences {
-      guard let offset = captureOffsets[id] else {
+      guard let offset = referencedCaptureOffsets[id] else {
         throw RegexCompilationError.uncapturedReference
       }
       for use in uses {
@@ -441,7 +442,7 @@ extension MEProgram.Builder {
     defer { nextCaptureRegister.rawValue += 1 }
     // Register the capture for later lookup via symbolic references.
     if let id = id {
-      let preexistingValue = captureOffsets.updateValue(
+      let preexistingValue = referencedCaptureOffsets.updateValue(
         captureCount, forKey: id)
       assert(preexistingValue == nil)
     }

Sources/_StringProcessing/Engine/MECapture.swift

@@ -143,14 +143,15 @@ extension Processor._StoredCapture: CustomStringConvertible {
 }
 
 public struct CaptureList {
-  var caps: Array<Processor<String>._StoredCapture>
+  var values: Array<Processor<String>._StoredCapture>
+  var referencedCaptureOffsets: [Reference.ID: Int]
 
 //  func extract(from s: String) -> Array<Array<Substring>> {
 //    caps.map { $0.map { s[$0] }  }
 //  }
 //
   func latestUntyped(from s: String) -> Array<Substring?> {
-    caps.map {
+    values.map {
       guard let last = $0.latest else {
         return nil
       }

Sources/_StringProcessing/Engine/MEProgram.swift

@@ -35,6 +35,7 @@ public struct MEProgram<Input: Collection> where Input.Element: Equatable {
   var enableTracing: Bool = false
 
   let captureStructure: CaptureStructure
+  let referencedCaptureOffsets: [Reference.ID: Int]
 }
 
 extension MEProgram: CustomStringConvertible {

Sources/_StringProcessing/Engine/Processor.swift

@@ -14,6 +14,8 @@ public enum MatchMode {
   case partialFromFront
 }
 
+typealias Program = MEProgram<String>
+
 /// A concrete CU. Somehow will run the concrete logic and
 /// feed stuff back to generic code
 struct Controller {

Sources/_StringProcessing/Engine/StringProcessor.swift

@@ -29,21 +29,21 @@ private struct Fabricator {
 
   mutating func next(
   ) throws -> Processor<String>._StoredCapture {
-    guard curIdx < list.caps.endIndex else {
+    guard curIdx < list.values.endIndex else {
       // TODO: Is `throws` a bit much here?
       // Maybe just precondition or hard trap
       throw Unreachable("Capture count mismatch")
     }
-    defer { list.caps.formIndex(after: &curIdx) }
-    return list.caps[curIdx]
+    defer { list.values.formIndex(after: &curIdx) }
+    return list.values[curIdx]
   }
 
   var currentIsEmpty: Bool {
-    guard curIdx < list.caps.endIndex else {
+    guard curIdx < list.values.endIndex else {
       fatalError("Capture count mismatch")
     }
 
-    return list.caps[curIdx].isEmpty
+    return list.values[curIdx].isEmpty
   }
 
   mutating func formValue(

Sources/_StringProcessing/Engine/Structuralize.swift

@@ -11,6 +11,7 @@
 
 import _MatchingEngine
 
+  // FIXME: Public for prototype
 public struct Executor {
   // TODO: consider let, for now lets us toggle tracing
   var engine: Engine<String>
@@ -19,11 +20,35 @@ public struct Executor {
     self.engine = Engine(program, enableTracing: enablesTracing)
   }
 
+  // FIXME: Public for prototype
+  public struct Result {
+    public var range: Range<String.Index>
+    var captures: Capture
+    var referencedCaptureOffsets: [Reference.ID: Int]
+
+    var destructure: (
+      matched: Range<String.Index>,
+      captures: Capture,
+      referencedCaptureOffsets: [Reference.ID: Int]
+    ) {
+      (range, captures, referencedCaptureOffsets)
+    }
+
+    init(
+      _ matched: Range<String.Index>, _ captures: Capture,
+      _ referencedCaptureOffsets: [Reference.ID: Int]
+    ) {
+      self.range = matched
+      self.captures = captures
+      self.referencedCaptureOffsets = referencedCaptureOffsets
+    }
+  }
+
   public func execute(
     input: String,
     in range: Range<String.Index>,
     mode: MatchMode = .wholeString
-  ) -> MatchResult? {
+  ) -> Result? {
     guard let (endIdx, capList) = engine.consume(
       input, in: range, matchMode: mode
     ) else {
@@ -32,15 +57,16 @@ public struct Executor {
     let capStruct = engine.program.captureStructure
     do {
       let caps = try capStruct.structuralize(capList, input)
-      return MatchResult(range.lowerBound..<endIdx, caps)
+      return Result(
+        range.lowerBound..<endIdx, caps, capList.referencedCaptureOffsets)
     } catch {
       fatalError(String(describing: error))
     }
   }
   public func execute(
     input: Substring,
     mode: MatchMode = .wholeString
-  ) -> MatchResult? {
+  ) -> Result? {
     self.execute(
       input: input.base,
       in: input.startIndex..<input.endIndex,

Sources/_StringProcessing/Executor.swift

@@ -11,8 +11,24 @@
 
 @dynamicMemberLookup
 public struct RegexMatch<Match> {
+  let input: String
   public let range: Range<String.Index>
-  public let match: Match
+  let rawCaptures: Capture
+  let referencedCaptureOffsets: [Reference.ID: Int]
+
+  public var match: Match {
+    if Match.self == (Substring, DynamicCaptures).self {
+      return (input[range], DynamicCaptures(rawCaptures)) as! Match
+    } else if Match.self == Substring.self {
+      // FIXME: Plumb whole match (`.0`) through the matching engine.
+      return input[range] as! Match
+    } else {
+      let typeErasedMatch = rawCaptures.matchValue(
+        withWholeMatch: input[range]
+      )
+      return typeErasedMatch as! Match
+    }
+  }
 
   public subscript<T>(dynamicMember keyPath: KeyPath<Match, T>) -> T {
     match[keyPath: keyPath]
@@ -25,6 +41,26 @@ public struct RegexMatch<Match> {
   ) -> Match {
     match
   }
+
+  public subscript(_ reference: Reference) -> Any {
+    guard let offset = referencedCaptureOffsets[reference.id] else {
+      preconditionFailure(
+        "Reference did not capture any match in the regex")
+    }
+    return rawCaptures[offset].value
+  }
+
+  public subscript<T>(_ reference: Reference, as _: T.Type) -> T {
+    let typeErasedMatch = self[reference]
+    guard let result = typeErasedMatch as? T else {
+      preconditionFailure(
+        """
+        The referenced capture has type '\(type(of: typeErasedMatch))', not
+        '\(T.self)'.
+        """)
+    }
+    return result
+  }
 }
 
 extension RegexProtocol {
@@ -43,23 +79,16 @@ extension RegexProtocol {
     mode: MatchMode = .wholeString
   ) -> RegexMatch<Match>? {
     let executor = Executor(program: regex.program.loweredProgram)
-    guard let (range, captures) = executor.execute(
+    guard let (range, captures, captureOffsets) = executor.execute(
       input: input, in: inputRange, mode: mode
     )?.destructure else {
       return nil
     }
-    let convertedMatch: Match
-    if Match.self == (Substring, DynamicCaptures).self {
-      convertedMatch = (input[range], DynamicCaptures(captures)) as! Match
-    } else if Match.self == Substring.self {
-      convertedMatch = input[range] as! Match
-    } else {
-      let typeErasedMatch = captures.matchValue(
-        withWholeMatch: input[range]
-      )
-      convertedMatch = typeErasedMatch as! Match
-    }
-    return RegexMatch(range: range, match: convertedMatch)
+    return RegexMatch(
+      input: input,
+      range: range,
+      rawCaptures: captures,
+      referencedCaptureOffsets: captureOffsets)
   }
 }
 

Sources/_StringProcessing/RegexDSL/Match.swift

@@ -582,6 +582,36 @@ class RegexDSLTests: XCTestCase {
       }
     }
 
+    // Match result referencing a `Reference`.
+    do {
+      let a = Reference()
+      let b = Reference()
+      let regex = Regex {
+        capture("abc", as: a)
+        zeroOrMore {
+          tryCapture(as: b) {
+            "#"
+            oneOrMore(.digit)
+          } transform: {
+            Int($0.dropFirst())
+          }
+        }
+        a
+        zeroOrMore {
+          capture(b)
+        }
+        optionally {
+          capture(a)
+        }
+      }
+      let input = "abc#41#42abc#42#42"
+      let result = try XCTUnwrap(input.match(regex))
+      XCTAssertEqual(result[a] as? Substring, "abc")
+      XCTAssertEqual(result[b] as? Int, 42)
+      XCTAssertEqual(result[a, as: Substring.self], "abc")
+      XCTAssertEqual(result[b, as: Int.self], 42)
+    }
+
     // Post-hoc captured references
     // #"(?:\w\1|:(\w):)+"#
     try _testDSLCaptures(