MatchingEngine support for (flat) captures with full history (#105)

 Capture and backreference support in backend and compiler

Author: milseman

Sources/_MatchingEngine/Engine/Backtracking.swift

@@ -0,0 +1,54 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2021-2022 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
+//
+//===----------------------------------------------------------------------===//
+
+extension Processor {
+
+  // TODO: What all do we want to save? Configurable?
+  // TODO: Do we need to save any registers?
+  // TODO: Is this the right place to do function stack unwinding?
+  struct SavePoint {
+    var pc: InstructionAddress
+    var pos: Position?
+
+    // The end of the call stack, so we can slice it off
+    // when failing inside a call.
+    //
+    // NOTE: Alternatively, also place return addresses on the
+    // save point stack
+    var stackEnd: CallStackAddress
+
+    // FIXME: Save minimal info (e.g. stack position and
+    // perhaps current start)
+    var captureEnds: [_StoredCapture]
+
+    var destructure: (
+      pc: InstructionAddress,
+      pos: Position?,
+      stackEnd: CallStackAddress,
+      captureEnds: [_StoredCapture]
+    ) {
+      (pc, pos, stackEnd, captureEnds)
+    }
+  }
+
+  func makeSavePoint(
+    _ pc: InstructionAddress,
+    addressOnly: Bool = false
+  ) -> SavePoint {
+    SavePoint(
+      pc: pc,
+      pos: addressOnly ? nil : currentPosition,
+      stackEnd: .init(callStack.count),
+      captureEnds: storedCaptures)
+  }
+}
+
+

Sources/_MatchingEngine/Engine/Builder.swift

@@ -27,6 +27,7 @@ extension Program where Input.Element: Hashable {
     var nextBoolRegister = BoolRegister(0)
     var nextIntRegister = IntRegister(0)
     var nextPositionRegister = PositionRegister(0)
+    var nextCaptureRegister = CaptureRegister(0)
 
     // Special addresses or instructions
     var failAddressToken: AddressToken? = nil
@@ -81,8 +82,7 @@ extension Program.Builder {
   public mutating func buildMoveImmediate(
     _ value: Int, into: IntRegister
   ) {
-    let uint = UInt64(truncatingIfNeeded: value)
-    assert(uint == value)
+    let uint = UInt64(asserting: value)
     buildMoveImmediate(uint, into: into)
   }
 
@@ -207,6 +207,27 @@ extension Program.Builder {
     instructions.append(.init(.print, .init(string: s)))
   }
 
+  public mutating func buildBeginCapture(
+    _ cap: CaptureRegister
+  ) {
+    instructions.append(
+      .init(.beginCapture, .init(capture: cap)))
+  }
+
+  public mutating func buildEndCapture(
+    _ cap: CaptureRegister
+  ) {
+    instructions.append(
+      .init(.endCapture, .init(capture: cap)))
+  }
+
+  public mutating func buildBackreference(
+    _ cap: CaptureRegister
+  ) {
+    instructions.append(
+      .init(.backreference, .init(capture: cap)))
+  }
+
   // TODO: Mutating because of fail address fixup, drop when
   // that's removed
   public mutating func assemble() -> Program {
@@ -263,6 +284,7 @@ extension Program.Builder {
     regInfo.positions = nextPositionRegister.rawValue
     regInfo.consumeFunctions = consumeFunctions.count
     regInfo.assertionFunctions = assertionFunctions.count
+    regInfo.captures = nextCaptureRegister.rawValue
 
     return Program(
       instructions: InstructionList(instructions),
@@ -341,6 +363,11 @@ extension Program.Builder {
 
 // Register helpers
 extension Program.Builder {
+  public mutating func makeCapture() -> CaptureRegister {
+    defer { nextCaptureRegister.rawValue += 1 }
+    return nextCaptureRegister
+  }
+
   public mutating func makeBoolRegister() -> BoolRegister {
     defer { nextBoolRegister.rawValue += 1 }
     return nextBoolRegister

Sources/_MatchingEngine/Engine/Capture.swift

@@ -0,0 +1,116 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2021-2022 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
+//
+//===----------------------------------------------------------------------===//
+
+/*
+
+ TODO: Specialized data structure for all captures:
+
+ - We want to be able to refer to COW prefixes for which
+   simple appends do not invalidate
+ - We want a compact save-point representation
+
+ TODO: Conjectures:
+
+ - We should be able to remove the entire capture history,
+   lazily recomputing it on-request from the initial stored
+   save point
+ - We should be able to keep these flat and simple, lazily
+   constructing structured types on-request
+
+ */
+
+
+extension Processor {
+  struct _StoredCapture {
+    // Set whenever we push the very first capture, allows us
+    // to theoretically re-compute anything we want to later.
+    fileprivate var startState: SavePoint? = nil
+
+    // Save the entire history as we go, so that backtracking
+    // can just lop-off aborted runs.
+    //
+    // Backtracking entries can specify a per-capture stack
+    // index so that we can abort anything that came after.
+    //
+    // By remembering the entire history, we waste space, but
+    // we get flexibility for now.
+    //
+    fileprivate var stack: Array<Range<Position>> = []
+
+    // An in-progress capture start
+    fileprivate var currentCaptureBegin: Position? = nil
+
+    fileprivate func _invariantCheck() {
+      if startState == nil {
+        assert(stack.isEmpty)
+        assert(currentCaptureBegin == nil)
+      } else {
+        assert(!stack.isEmpty || currentCaptureBegin != nil)
+      }
+    }
+
+    // MARK: - IPI
+
+    var isEmpty: Bool { stack.isEmpty }
+
+    var history: Array<Range<Position>> {
+      stack
+    }
+
+    var latest: Range<Position>? { stack.last }
+
+    /// Start a new capture. If the previously started one was un-ended,
+    /// will clear it and restart. If this is the first start, will save `initial`.
+    mutating func startCapture(
+      _ idx: Position, initial: SavePoint
+    ) {
+      _invariantCheck()
+      defer { _invariantCheck() }
+
+      if self.startState == nil {
+        self.startState = initial
+      }
+      currentCaptureBegin = idx
+    }
+
+    mutating func endCapture(_ idx: Position) {
+      _invariantCheck()
+      assert(currentCaptureBegin != nil)
+      defer { _invariantCheck() }
+
+      stack.append(currentCaptureBegin! ..< idx)
+    }
+
+    mutating func fail(truncatingAt stackIdx: Int) {
+      _invariantCheck()
+      assert(stackIdx <= stack.endIndex)
+      defer { _invariantCheck() }
+
+      stack.removeSubrange(stackIdx...)
+      if stack.isEmpty {
+        startState = nil
+      }
+    }
+  }
+}
+
+public struct CaptureList {
+  var caps: Array<Array<Range<String.Index>>>
+
+  func extract(from s: String) -> Array<Array<Substring>> {
+    caps.map { $0.map { s[$0] }  }
+  }
+
+  func latest(from s: String) -> Array<Substring?> {
+    // TODO: If empty, probably need empty range or something...
+    extract(from: s).map { $0.last }
+  }
+}

Sources/_MatchingEngine/Engine/Consume.swift

@@ -22,16 +22,20 @@ extension Engine {
       matchMode: matchMode,
       isTracingEnabled: enableTracing)
   }
+}
 
-  public func consume(_ input: Input) -> Input.Index? {
+extension Engine where Input == String {
+  public func consume(
+    _ input: Input
+  ) -> (Input.Index, CaptureList)? {
     consume(input, in: input.startIndex ..< input.endIndex)
   }
 
   public func consume(
     _ input: Input,
     in range: Range<Input.Index>,
     matchMode: MatchMode = .prefix
-  ) -> Input.Index? {
+  ) -> (Input.Index, CaptureList)? {
     if enableTracing {
       print("Consume: \(input)")
     }
@@ -56,7 +60,10 @@ extension Engine {
         print("Result: nil")
       }
     }
-    return result
+    guard let result = result else { return nil }
+
+    let capList = cpu.storedCaptures.map { $0.history }
+    return (result, CaptureList(caps: capList))
   }
 }
 

Sources/_MatchingEngine/Engine/InstPayload.swift

@@ -52,6 +52,7 @@ extension Instruction.Payload {
     case consumer(ConsumeFunctionRegister)
     case assertion(AssertionFunctionRegister)
     case addr(InstructionAddress)
+    case capture(CaptureRegister)
 
     case packedImmInt(Int, IntRegister)
     case packedAddrBool(InstructionAddress, BoolRegister)
@@ -121,6 +122,7 @@ extension Instruction.Payload {
     // TODO: We'd like to use shadow bits to assert on kind
     return TypedInt(rawValue)
   }
+
   private func interpretPair<👻>(
     secondAs: TypedInt<👻>.Type = TypedInt<👻>.self
   ) -> (UInt64, TypedInt<👻>) {
@@ -212,6 +214,13 @@ extension Instruction.Payload {
     interpret()
   }
 
+  init(capture: CaptureRegister) {
+    self.init(capture)
+  }
+  var capture: CaptureRegister {
+    interpret()
+  }
+
 
   // MARK: Packed operand payloads
 

Sources/_MatchingEngine/Engine/Instruction.swift

@@ -228,6 +228,24 @@ extension Instruction {
     ///
     case splitSaving
 
+    /// Begin the given capture
+    ///
+    ///     beginCapture(_:CapReg)
+    ///
+    case beginCapture
+
+    /// End the given capture
+    ///
+    ///     endCapture(_:CapReg)
+    ///
+    case endCapture
+
+    /// Match a previously captured value
+    ///
+    ///     backreference(_:CapReg)
+    ///
+    case backreference
+
     // MARK: Matching: State transitions
 
     // TODO: State transitions need more work. We want