Fix two crashers with quantification on transformed capture.

This patch fixes two issues:
- When an optional quantification matches 1 occurrence, the capture type constructed by the engine unexpectedly drops the `Optional` wrapper type. `Capture.value` is fixed by including the `Optional` wrapper type.
- When an optional or array quantification matches 0 occurrence, the capture type constructed by the engine unexpectedly uses `Substring` as the atom type instead of the transformed type from `CaptureTransform`. This is fixed by propagating capture transform result type information from the DSL API down to `captureNil` and `captureArray` instructions in the bytecode. The engine will use the type information to construct the correct types for `nil` and `[]`.

Author: rxwei

Sources/_MatchingEngine/Regex/AST/AST.swift

@@ -94,10 +94,13 @@ extension AST {
 
   /// Whether this node has nested somewhere inside it a capture
   public var hasCapture: Bool {
-    if case let .group(g) = self, g.kind.value.isCapturing {
+    switch self {
+    case .group(let g) where g.kind.value.isCapturing,
+         .groupTransform(let g, _) where g.kind.value.isCapturing:
       return true
+    default:
+      break
     }
-
     return self.children?.any(\.hasCapture) ?? false
   }
 }
@@ -207,14 +210,18 @@ extension AST {
 
 // FIXME: Get this out of here
 public struct CaptureTransform: Equatable, Hashable, CustomStringConvertible {
+  public let resultType: Any.Type
   public let closure: (Substring) -> Any
 
-  public init(_ closure: @escaping (Substring) -> Any) {
+  public init(resultType: Any.Type, _ closure: @escaping (Substring) -> Any) {
+    self.resultType = resultType
     self.closure = closure
   }
 
   public func callAsFunction(_ input: Substring) -> Any {
-    closure(input)
+    let result = closure(input)
+    assert(type(of: result) == resultType)
+    return result
   }
 
   public static func == (lhs: CaptureTransform, rhs: CaptureTransform) -> Bool {
@@ -229,7 +236,6 @@ public struct CaptureTransform: Equatable, Hashable, CustomStringConvertible {
   }
 
   public var description: String {
-    "<transform>"
+    "<transform result_type=\(resultType)>"
   }
 }
-

Sources/_MatchingEngine/Regex/Parse/CaptureStructure.swift

@@ -11,7 +11,7 @@
 
 // A tree representing the type of some captures.
 public enum CaptureStructure: Equatable {
-  case atom(name: String? = nil)
+  case atom(name: String? = nil, type: AnyType? = nil)
   indirect case array(CaptureStructure)
   indirect case optional(CaptureStructure)
   indirect case tuple([CaptureStructure])
@@ -47,6 +47,17 @@ extension AST {
       default:
         return innerCaptures
       }
+    case .groupTransform(let group, let transform):
+      let innerCaptures = group.child.captureStructure
+      switch group.kind.value {
+      case .capture:
+        return .atom(type: AnyType(transform.resultType)) + innerCaptures
+      case .namedCapture(let name):
+        return .atom(name: name.value, type: AnyType(transform.resultType))
+          + innerCaptures
+      default:
+        return innerCaptures
+      }
     case .conditional(let c):
       // A conditional's capture structure is effectively that of an alternation
       // between the true and false branches. However the condition may also
@@ -67,8 +78,6 @@ extension AST {
         quantification.amount.value == .zeroOrOne
           ? CaptureStructure.optional
           : CaptureStructure.array)
-    case .groupTransform:
-      fatalError("Unreachable. Case will be removed later.")
     case .quote, .trivia, .atom, .customCharacterClass, .empty:
       return .empty
     }
@@ -135,8 +144,10 @@ extension CaptureStructure {
 
   public func type(withAtomType atomType: Any.Type) -> Any.Type {
     switch self {
-    case .atom:
+    case .atom(_, type: nil):
       return atomType
+    case .atom(_, type: let type?):
+      return type.base
     case .array(let child):
       return TypeConstruction.arrayType(of: child.type(withAtomType: atomType))
     case .optional(let child):
@@ -213,16 +224,18 @@ extension CaptureStructure {
     func encode(_ node: CaptureStructure, isTopLevel: Bool = false) {
       switch node {
       // 〚`T` (atom)〛 ==> .atom
-      case .atom(name: nil):
+      case .atom(name: nil, type: nil):
         append(.atom)
       // 〚`name: T` (atom)〛 ==> .atom, `name`, '\0'
-      case .atom(name: let name?):
+      case .atom(name: let name?, type: nil):
         append(.namedAtom)
         let nameCString = name.utf8CString
         let nameSlot = UnsafeMutableRawBufferPointer(
           rebasing: buffer[offset ..< offset+nameCString.count])
         nameCString.withUnsafeBytes(nameSlot.copyMemory(from:))
         offset += nameCString.count
+      case .atom(_, _?):
+        fatalError("Cannot encode a capture structure with explicit types")
       // 〚`[T]`〛 ==> 〚`T`〛, .formArray
       case .array(let child):
         encode(child)

Sources/_MatchingEngine/Utility/Misc.swift

@@ -149,3 +149,19 @@ extension BinaryInteger {
   }
 }
 
+/// A wrapper of an existential metatype, equatable and hashable by reference.
+public struct AnyType: Equatable, Hashable {
+  public var base: Any.Type
+
+  public init(_ type: Any.Type) {
+    base = type
+  }
+
+  public static func == (lhs: AnyType, rhs: AnyType) -> Bool {
+    lhs.base == rhs.base
+  }
+
+  public func hash(into hasher: inout Hasher) {
+    hasher.combine(ObjectIdentifier(base))
+  }
+}

Sources/_StringProcessing/Capture.swift

@@ -51,7 +51,10 @@ extension Capture {
       }
       return _openExistential(childType.base, do: helper)
     case .some(let subcapture):
-      return subcapture.value
+      func helper<T>(_ value: T) -> Any {
+        Optional(value) as Any
+      }
+      return _openExistential(subcapture.value, do: helper)
     case .none(let childType):
       func helper<T>(_: T.Type) -> Any {
         nil as T? as Any

Sources/_StringProcessing/RegexDSL/DSL.swift

@@ -179,7 +179,7 @@ public struct CapturingGroup<Match: MatchProtocol>: RegexProtocol {
     self.regex = .init(ast:
       .groupTransform(
         .init(.init(faking: .capture), component.regex.ast, .fake),
-        transform: CaptureTransform {
+        transform: CaptureTransform(resultType: NewCapture.self) {
           transform($0) as Any
         }))
   }

Tests/RegexTests/RegexDSLTests.swift

@@ -80,6 +80,55 @@ class RegexDSLTests: XCTestCase {
         == Tuple3("b", "cccc", ["d", "d", "d"]))
   }
 
+  func testQuantificationWithTransformedCapture() throws {
+    // This test is to make sure transformed capture type information is
+    // correctly propagated from the DSL into the bytecode and that the engine
+    // is reconstructing the right types upon quantification (both empty and
+    // non-empty).
+    enum Word: Int32 {
+      case apple
+      case orange
+
+      init?(_ string: Substring) {
+        switch string {
+        case "apple": self = .apple
+        case "orange": self = .orange
+        default: return nil
+        }
+      }
+    }
+    let regex = Regex {
+      "a".+
+      OneOrMore(.whitespace)
+      Optionally {
+        OneOrMore(.digit).capture { Int($0)! }
+      }
+      Repeat {
+        OneOrMore(.whitespace)
+        OneOrMore(.word).capture { Word($0)! }
+      }
+    }
+    // Assert the inferred capture type.
+    let _: Tuple3<Substring, Int?, [Word]>.Type
+      = type(of: regex).Match.self
+    do {
+      let input = "aaa 123 apple orange apple"
+      let match = input.match(regex)?.match.tuple
+      let (whole, number, words) = try XCTUnwrap(match)
+      XCTAssertTrue(whole == input)
+      XCTAssertEqual(number, 123)
+      XCTAssertEqual(words, [.apple, .orange, .apple])
+    }
+    do {
+      let input = "aaa   "
+      let match = input.match(regex)?.match.tuple
+      let (whole, number, words) = try XCTUnwrap(match)
+      XCTAssertTrue(whole == input)
+      XCTAssertEqual(number, nil)
+      XCTAssertTrue(words.isEmpty)
+    }
+  }
+
   // Note: Types of nested captures should be flat, but are currently nested
   // due to the lack of variadic generics. Without it, we cannot effectively
   // express type constraints to concatenate splatted tuples.
@@ -174,39 +223,46 @@ class RegexDSLTests: XCTestCase {
     let line = """
       A6F0..A6F1    ; Extend # Mn   [2] BAMUM COMBINING MARK KOQNDON..BAMUM COMBINING MARK TUKWENTIS
       """
-    let regex = Regex {
-      OneOrMore(CharacterClass.hexDigit).capture()
+    
+    let regexWithCapture = Regex {
+      OneOrMore(CharacterClass.hexDigit).capture(Unicode.Scalar.init(hex:))
       Optionally {
         ".."
-        OneOrMore(CharacterClass.hexDigit).capture()
+        OneOrMore(CharacterClass.hexDigit).capture(Unicode.Scalar.init(hex:))
       }
       OneOrMore(CharacterClass.whitespace)
       ";"
       OneOrMore(CharacterClass.whitespace)
       OneOrMore(CharacterClass.word).capture()
       Repeat(CharacterClass.any)
+    } // Regex<(Substring, Unicode.Scalar?, Unicode.Scalar??, Substring)>
+    do {
+      // Assert the inferred capture type.
+      typealias ExpectedMatch = Tuple4<
+        Substring, Unicode.Scalar?, Unicode.Scalar??, Substring
+      >
+      let _: ExpectedMatch.Type = type(of: regexWithCapture).Match.self
+      let maybeMatchResult = line.match(regexWithCapture)
+      let matchResult = try XCTUnwrap(maybeMatchResult)
+      let (wholeMatch, lower, upper, propertyString) = matchResult.match.tuple
+      XCTAssertEqual(wholeMatch, Substring(line))
+      XCTAssertEqual(lower, Unicode.Scalar(0xA6F0))
+      XCTAssertEqual(upper, Unicode.Scalar(0xA6F1))
+      XCTAssertEqual(propertyString, "Extend")
     }
-    // Assert the inferred capture type.
-    typealias ExpectedMatch = Tuple4<
-      Substring, Substring, Substring?, Substring
-    >
-    let _: ExpectedMatch.Type = type(of: regex).Match.self
-    func run<R: RegexProtocol>(
-      _ regex: R
-    ) throws where R.Match == ExpectedMatch {
-      let maybeMatchResult = line.match(regex)
+
+    do {
+      let regexLiteral = try MockRegexLiteral(
+        #"([0-9A-F]+)(?:\.\.([0-9A-F]+))?\s+;\s+(\w+).*"#,
+        matching: Tuple4<Substring, Substring, Substring?, Substring>.self)
+      let maybeMatchResult = line.match(regexLiteral)
       let matchResult = try XCTUnwrap(maybeMatchResult)
       let (wholeMatch, lower, upper, propertyString) = matchResult.match.tuple
       XCTAssertEqual(wholeMatch, Substring(line))
       XCTAssertEqual(lower, "A6F0")
       XCTAssertEqual(upper, "A6F1")
       XCTAssertEqual(propertyString, "Extend")
     }
-    let regexLiteral = try MockRegexLiteral(
-      #"([0-9A-F]+)(?:\.\.([0-9A-F]+))?\s+;\s+(\w+).*"#,
-      matching: Tuple4<Substring, Substring, Substring?, Substring>.self)
-    try run(regex)
-    try run(regexLiteral)
   }
 
   func testDynamicCaptures() throws {
@@ -233,3 +289,13 @@ class RegexDSLTests: XCTestCase {
     }
   }
 }
+
+extension Unicode.Scalar {
+  // Convert a hexadecimal string to a scalar
+  public init?<S: StringProtocol>(hex: S) {
+    guard let val = UInt32(hex, radix: 16), let scalar = Self(val) else {
+      return nil
+    }
+    self = scalar
+  }
+}