Switch to lazily performing capture transforms

Sources/_StringProcessing/ByteCodeGen+DSLList.swift

@@ -779,25 +779,33 @@ fileprivate extension Compiler.ByteCodeGen {
       if let value {
         builder.buildMove(value, into: cap)
       }
-      // If there's a capture transform, apply it now.
+      // If there's a capture transform, apply it only after matching.
+      // A failable transform needs to run during matching, so that a
+      // `nil` result can backtrack –– its result is cached with the capture
+      // and used when the match is complete.
       if let transform = transform {
-        let fn = builder.makeTransformFunction { input, cap in
+        let fn: MEProgram.TransformFunction = { input, range, customValue in
           // If it's a substring capture with no custom value, apply the
           // transform directly to the substring to avoid existential traffic.
           //
           // FIXME: separate out this code path. This is fragile,
           // slow, and these are clearly different constructs
-          if let range = cap.range, cap.value == nil {
+          if let range, customValue == nil {
             return try transform(input[range])
           }
 
           let value = constructExistentialOutputComponent(
             from: input,
-            component: cap.deconstructed,
+            component: range.map { ($0, customValue) },
             optionalCount: 0)
           return try transform(value)
         }
-        builder.buildTransformCapture(cap, fn)
+        let reg = builder.makeTransformFunction(fn)
+        if transform.isFailable {
+          builder.buildTransformCapture(cap, reg)
+        } else {
+          builder.captureTransforms[cap.rawValue] = reg
+        }
       }
 
     case let .nonCapturingGroup(kind):

Sources/_StringProcessing/Engine/MEBuilder.swift

@@ -25,6 +25,7 @@ extension MEProgram {
     var asciiBitsets: [DSLTree.CustomCharacterClass.AsciiBitset] = []
     var consumeFunctions: [ConsumeFunction] = []
     var transformFunctions: [TransformFunction] = []
+    var captureTransforms: [Int: TransformRegister] = [:]
     var matcherFunctions: [MatcherFunction] = []
 
     // Map tokens to actual addresses
@@ -439,6 +440,7 @@ extension MEProgram.Builder {
       enableMetrics: enableMetrics,
       captureList: captureList,
       referencedCaptureOffsets: referencedCaptureOffsets,
+      captureTransforms: captureTransforms,
       initialOptions: initialOptions,
       canOnlyMatchAtStart: canOnlyMatchAtStart,
       registers: regs)

Sources/_StringProcessing/Engine/MECapture.swift

@@ -15,22 +15,35 @@ extension Processor {
   struct _StoredCapture {
     var range: Range<Position>? = nil
 
-    var value: Any? = nil
+    /// The index in `Registers.customValues` of this capture's value, if
+    /// already created.
+    ///
+    /// The value at `valueIndex` is produced by a custom matcher, is the
+    /// transformed result of a custom matcher, or is the result of a
+    /// `TryCapture` failable transform. Non-failable transforms aren't run
+    /// until a match is complete.
+    var valueIndex: Int? = nil
 
     // An in-progress capture start
     var currentCaptureBegin: Position? = nil
 
     fileprivate func _invariantCheck() {
       if range == nil {
-        assert(value == nil)
+        assert(valueIndex == nil)
       }
     }
 
     // MARK: - IPI
 
-    var deconstructed: (range: Range<Position>, value: Any?)? {
+    func deconstructed(
+      customValues: [Any]
+    ) -> (range: Range<Position>, value: Any?)? {
       guard let r = range else { return nil }
-      return (r, value)
+      return (r, value(in: customValues))
+    }
+
+    func value(in customValues: [Any]) -> Any? {
+      valueIndex.map { customValues[$0] }
     }
 
     /// Start a new capture. If the previously started one was un-ended,
@@ -53,15 +66,15 @@ extension Processor {
       }
 
       range = low..<idx
-      value = nil // TODO: cleaner IPI around this...
+      valueIndex = nil
       currentCaptureBegin = nil
     }
 
-    mutating func registerValue(_ value: Any) {
+    mutating func registerValue(at index: Int) {
       _invariantCheck()
       defer { _invariantCheck() }
 
-      self.value = value
+      valueIndex = index
     }
   }
 }

Sources/_StringProcessing/Engine/MEProgram.swift

@@ -15,8 +15,10 @@ struct MEProgram {
   typealias Input = String
 
   typealias ConsumeFunction = (Input, Range<Input.Index>) -> Input.Index?
+  /// Transforms a capture's range, and the custom matcher value for the
+  /// capture if there is one.
   typealias TransformFunction =
-    (Input, Processor._StoredCapture) throws -> Any?
+    (Input, Range<Input.Index>?, Any?) throws -> Any?
   typealias MatcherFunction =
     (Input, Input.Index, Range<Input.Index>) throws -> (Input.Index, Any)?
 
@@ -28,6 +30,10 @@ struct MEProgram {
 
   let captureList: CaptureList
   let referencedCaptureOffsets: [ReferenceID: Int]
+
+  /// The registers of non-failable capture transforms, by capture index.
+  /// Failable transforms run while matching, and aren't included.
+  let captureTransforms: [Int: TransformRegister]
 
   var initialOptions: MatchingOptions
   var canOnlyMatchAtStart: Bool

Sources/_StringProcessing/Engine/Processor.swift

@@ -435,8 +435,10 @@ extension Processor {
   }
 
   mutating func setCaptureValue(_ capNum: Int, _ value: Any) {
+    let index = registers.customValues.count
+    registers.customValues.append(value)
     updateRegister(at: CaptureRegister(capNum)) {
-      $0.registerValue(value)
+      $0.registerValue(at: index)
     }
   }
 
@@ -696,15 +698,18 @@ extension Processor {
     case .transformCapture:
       let (cap, trans) = payload.pairedCaptureTransform
       let transform = registers[trans]
-      let capNum = Int(asserting: cap.rawValue)
 
       do {
         // FIXME: Pass input or the slice?
-        guard let value = try transform(input, registers.storedCaptures[cap]) else {
+        let stored = registers.storedCaptures[cap]
+        guard let value = try transform(
+          input, stored.range, stored.value(in: registers.customValues)
+        ) else {
           signalFailure()
           return
         }
-        setCaptureValue(capNum, value)
+        // The result replaces the capture's value, which it was computed from.
+        setCaptureValue(Int(asserting: cap.rawValue), value)
         controller.step()
       } catch {
         abort(error)

Sources/_StringProcessing/Engine/Registers.swift

@@ -49,6 +49,9 @@ extension Processor {
 
     var storedCaptures: UndoableArray<CaptureRegister, Processor._StoredCapture>
 
+    /// Values produced by custom matchers and `TryCapture` transforms.
+    var customValues: [Any] = []
+
     var isDirty = false
 
     init(
@@ -166,6 +169,7 @@ extension Processor.Registers {
     values.reset(to: SentinelValue())
     positions.reset(to: Processor.Registers.sentinelIndex)
     storedCaptures.reset(to: Processor._StoredCapture())
+    customValues.removeAll(keepingCapacity: true)
     isDirty = false
   }
 

Sources/_StringProcessing/Engine/Structuralize.swift

@@ -17,8 +17,10 @@ extension Executor {
     _ program: MEProgram,
     matchRange: Range<String.Index>,
     storedCaptures: [Processor._StoredCapture],
+    customValues: [Any],
+    input: String,
     wholeMatchValue: Any?
-  ) -> [AnyRegexOutput.ElementRepresentation] {
+  ) throws -> [AnyRegexOutput.ElementRepresentation] {
     let capList = program.captureList
     let capOffsets = program.referencedCaptureOffsets
 
@@ -37,9 +39,18 @@ extension Executor {
     for (i, (cap, meStored)) in zip(
       capList.captures.dropFirst(), storedCaptures
     ).enumerated() {
+      var content = meStored.deconstructed(customValues: customValues)
+      // A failable transform already ran while matching, and its result is the
+      // capture's value. Any other transform needs to run now that matching has
+      // succeeded.
+      if let reg = program.captureTransforms[i], content != nil {
+        let transform = program.registers[reg]
+        let value = try transform(input, meStored.range, content?.value)
+        content?.value = value
+      }
       let element = AnyRegexOutput.ElementRepresentation(
         optionalDepth: cap.optionalDepth,
-        content: meStored.deconstructed,
+        content: content,
         name: cap.name,
         referenceID: capOffsets.first { $1 == i }?.key,
         visibleInTypedOutput: cap.visibleInTypedOutput

Sources/_StringProcessing/Executor.swift

@@ -198,10 +198,12 @@ extension Executor {
     } else {
       wholeMatchValue = nil
     }
-    let aroElements = Executor.createExistentialElements(
+    let aroElements = try Executor.createExistentialElements(
       program,
       matchRange: startPosition..<endIdx,
       storedCaptures: cpu.registers.storedCaptures.values,
+      customValues: cpu.registers.customValues,
+      input: cpu.input,
       wholeMatchValue: wholeMatchValue)
 
     let anyRegexOutput = AnyRegexOutput(

Sources/_StringProcessing/Regex/DSLTree.swift

@@ -490,6 +490,14 @@ struct CaptureTransform: Hashable, CustomStringConvertible {
       closure: closure)
   }
 
+  /// Whether a `nil` result from this transform means the capture failed.
+  var isFailable: Bool {
+    switch closure {
+    case .failable, .substringFailable: return true
+    case .nonfailable, .substringNonfailable: return false
+    }
+  }
+
   func callAsFunction(_ input: Any) throws -> Any? {
     switch closure {
     case .nonfailable(let transform):

Tests/RegexBuilderTests/CustomTests.swift

@@ -224,6 +224,43 @@ extension Concat: BidirectionalCollection {
 class CustomRegexComponentTests: XCTestCase {
   // TODO: Refactor below into more exhaustive, declarative
   // tests.
+  func testCustomValuesSurviveBacktracking() throws {
+    // Captured custom-matcher values must follow the capture across
+    // backtracking and be what a capture transform receives.
+    let r = Regex {
+      OneOrMore {
+        Capture(Numbler()) { $0 * 10 }
+      }
+      "!"
+    }
+    XCTAssertEqual(try r.wholeMatch(in: "123!")?.1, 30)
+
+    let choice = Regex {
+      ChoiceOf {
+        Regex { Capture(Numbler()); "a" }
+        Regex { Capture(Numbler()); "b" }
+      }
+    }
+    let b = try XCTUnwrap(choice.wholeMatch(in: "7b"))
+    XCTAssertNil(b.1)
+    XCTAssertEqual(b.2, 7)
+    let a = try XCTUnwrap(choice.wholeMatch(in: "8a"))
+    XCTAssertEqual(a.1, 8)
+    XCTAssertNil(a.2)
+
+    // A failable transform of a custom matcher's value, and a non-failable one.
+    let tried = Regex {
+      OneOrMore {
+        TryCapture(Numbler()) { (n: Int) -> Int? in n > 5 ? n * 2 : nil }
+      }
+      Capture(Numbler()) { (n: Int) in n + 100 }
+    }
+    let tm = try XCTUnwrap(tried.wholeMatch(in: "793"))
+    XCTAssertEqual(tm.1, 18)
+    XCTAssertEqual(tm.2, 103)
+    XCTAssertNil(try tried.wholeMatch(in: "3"))
+  }
+
   func testCustomRegexComponents() throws {
     customTest(
       Regex {

Tests/RegexBuilderTests/RegexDSLTests.swift

@@ -1990,6 +1990,60 @@ extension RegexDSLTests {
     XCTAssertNotNil(clip2.contains(pattern))
   }
 
+  func testLazyCaptureTransforms() throws {
+    struct TransformError: Error {}
+
+    // A throwing transform on a captured value still aborts the match.
+    let throwing = Regex {
+      Capture(OneOrMore(.digit)) { (s: Substring) -> Int in
+        if s == "13" { throw TransformError() }
+        return Int(s)!
+      }
+    }
+    XCTAssertEqual(try throwing.firstMatch(in: "a42")?.1, 42)
+    XCTAssertThrowsError(try throwing.firstMatch(in: "a13"))
+
+    // A failing `TryCapture` transform backtracks, and the transformed value
+    // of the final match is the one that's reported.
+    var calls = 0
+    let backtracking = Regex {
+      TryCapture(OneOrMore(.digit)) { (s: Substring) -> Int? in
+        calls += 1
+        return s.count == 2 ? Int(s) : nil
+      }
+      "x"
+    }
+    XCTAssertNil(try backtracking.wholeMatch(in: "123x"))
+    calls = 0
+    XCTAssertEqual(try backtracking.firstMatch(in: "1 12x")?.1, 12)
+    XCTAssertGreaterThan(calls, 0)
+
+    // A failable transform runs exactly once, and its result is the one
+    // that's reported, even if it isn't deterministic.
+    var counter = 0
+    let once = Regex {
+      TryCapture("a") { (_: Substring) -> Int? in
+        counter += 1
+        return counter
+      }
+    }
+    XCTAssertEqual(try XCTUnwrap(once.wholeMatch(in: "a")).1, 1)
+    XCTAssertEqual(counter, 1)
+
+    // Transforms of captures that don't participate aren't called, and
+    // captures within repetitions report the last iteration.
+    let repeated = Regex {
+      OneOrMore {
+        Capture(OneOrMore(.digit)) { Int($0)! }
+        ","
+      }
+      Optionally { Capture("z") { $0.count } }
+    }
+    let m = try XCTUnwrap(repeated.wholeMatch(in: "1,22,333,"))
+    XCTAssertEqual(m.1, 333)
+    XCTAssertNil(m.2)
+  }
+
   func testIssue83022() throws {
     // Original report from https://github.com/swiftlang/swift/issues/83022
     // rdar://155710126