Add UTF-8 byte matching optimization

Author: milseman

Sources/RegexBenchmark/Suite/URLRegex.swift

@@ -12,3 +12,4 @@ extension BenchmarkRunner {
     url.register(&self)
   }
 }
+

Sources/_StringProcessing/ByteCodeGen.swift

@@ -117,6 +117,30 @@ fileprivate extension Compiler.ByteCodeGen {
   }
 
   mutating func emitQuotedLiteral(_ s: String) {
+    // ASCII is normalization-invariant, so is the safe subset for
+    // us to optimize
+    if optimizationsEnabled,
+       !options.usesCanonicalEquivalence || s.utf8.allSatisfy(\._isASCII),
+       !s.isEmpty
+    {
+
+      // TODO: Make an optimizations configuration struct, where
+      // we can enable/disable specific optimizations and change
+      // thresholds
+      let longThreshold = 5
+
+      // Longer content will be matched against UTF-8 in contiguous
+      // memory
+      //
+      // TODO: case-insensitive variant (just add/subtract from
+      // ASCII value)
+      if s.utf8.count >= longThreshold, !options.isCaseInsensitive {
+        let boundaryCheck = options.semanticLevel == .graphemeCluster
+        builder.buildMatchUTF8(Array(s.utf8), boundaryCheck: boundaryCheck)
+        return
+      }
+    }
+
     guard options.semanticLevel == .graphemeCluster else {
       for char in s {
         for scalar in char.unicodeScalars {

Sources/_StringProcessing/Engine/InstPayload.swift

@@ -43,7 +43,7 @@ extension Instruction.Payload {
     // and variables
 
     case string(StringRegister)
-    case sequence(SequenceRegister)
+    case utf8(UTF8Register)
     case position(PositionRegister)
     case optionalString(StringRegister?)
     case int(IntRegister)
@@ -168,10 +168,18 @@ extension Instruction.Payload {
     return (scalar, caseInsensitive: caseInsensitive, boundaryCheck: boundaryCheck)
   }
 
-  init(sequence: SequenceRegister) {
-    self.init(sequence)
+  init(utf8: UTF8Register, boundaryCheck: Bool) {
+    self.init(boundaryCheck ? 1 : 0, utf8)
   }
-  var sequence: SequenceRegister {
+  var matchUTF8Payload: (UTF8Register, boundaryCheck: Bool) {
+    let pair: (UInt64, UTF8Register) = interpretPair()
+    return (pair.1, pair.0 == 1)
+  }
+
+  init(utf8: UTF8Register) {
+    self.init(utf8)
+  }
+  var utf8: UTF8Register {
     interpret()
   }
 

Sources/_StringProcessing/Engine/Instruction.swift

@@ -112,6 +112,17 @@ extension Instruction {
     /// Operands: Scalar value to match against and booleans
     case matchScalar
 
+    /// Match directly (binary semantics) against a series of UTF-8 bytes
+    ///
+    /// NOTE: Compiler should ensure to only emit this instruction when normalization
+    /// is not required. E.g., scalar-semantic mode or when the matched portion is entirely ASCII
+    /// (which is invariant under NFC). Similary, this is case-sensitive.
+    ///
+    /// TODO: should we add case-insensitive?
+    ///
+    ///     matchUTF8(_: UTF8Register, boundaryCheck: Bool)
+    case matchUTF8
+
     /// Match a character or a scalar against a set of valid ascii values stored in a bitset
     ///
     ///     matchBitset(_: AsciiBitsetRegister, isScalar: Bool)

Sources/_StringProcessing/Engine/MEBuilder.swift

@@ -20,7 +20,7 @@ extension MEProgram {
     var enableMetrics = false
 
     var elements = TypedSetVector<Input.Element, _ElementRegister>()
-    var sequences = TypedSetVector<[Input.Element], _SequenceRegister>()
+    var utf8Contents = TypedSetVector<[UInt8], _UTF8Register>()
 
     var asciiBitsets: [DSLTree.CustomCharacterClass.AsciiBitset] = []
     var consumeFunctions: [ConsumeFunction] = []
@@ -198,6 +198,11 @@ extension MEProgram.Builder {
       .match, .init(element: elements.store(e), isCaseInsensitive: isCaseInsensitive)))
   }
 
+  mutating func buildMatchUTF8(_ utf8: Array<UInt8>, boundaryCheck: Bool) {
+    instructions.append(.init(.matchUTF8, .init(
+      utf8: utf8Contents.store(utf8), boundaryCheck: boundaryCheck)))
+  }
+
   mutating func buildMatchScalar(_ s: Unicode.Scalar, boundaryCheck: Bool) {
     instructions.append(.init(.matchScalar, .init(scalar: s, caseInsensitive: false, boundaryCheck: boundaryCheck)))
   }
@@ -416,7 +421,7 @@ extension MEProgram.Builder {
 
     var regInfo = MEProgram.RegisterInfo()
     regInfo.elements = elements.count
-    regInfo.sequences = sequences.count
+    regInfo.utf8Contents = utf8Contents.count
     regInfo.ints = nextIntRegister.rawValue
     regInfo.values = nextValueRegister.rawValue
     regInfo.positions = nextPositionRegister.rawValue
@@ -430,7 +435,7 @@ extension MEProgram.Builder {
     return MEProgram(
       instructions: InstructionList(instructions),
       staticElements: elements.stored,
-      staticSequences: sequences.stored,
+      staticUTF8Contents: utf8Contents.stored,
       staticBitsets: asciiBitsets,
       staticConsumeFunctions: consumeFunctions,
       staticTransformFunctions: transformFunctions,

Sources/_StringProcessing/Engine/MEProgram.swift

@@ -23,7 +23,7 @@ struct MEProgram {
   var instructions: InstructionList<Instruction>
 
   var staticElements: [Input.Element]
-  var staticSequences: [[Input.Element]]
+  var staticUTF8Contents: [[UInt8]]
   var staticBitsets: [DSLTree.CustomCharacterClass.AsciiBitset]
   var staticConsumeFunctions: [ConsumeFunction]
   var staticTransformFunctions: [TransformFunction]

Sources/_StringProcessing/Engine/Processor.swift

@@ -320,6 +320,24 @@ extension Processor {
     return true
   }
 
+  // TODO: bytes should be a Span or RawSpan
+  mutating func matchUTF8(
+    _ bytes: Array<UInt8>,
+    boundaryCheck: Bool
+  ) -> Bool {
+    guard let next = input.matchUTF8(
+      bytes,
+      at: currentPosition,
+      limitedBy: end,
+      boundaryCheck: boundaryCheck
+    ) else {
+      signalFailure()
+      return false
+    }
+    currentPosition = next
+    return true
+  }
+
   // If we have a bitset we know that the CharacterClass only matches against
   // ascii characters, so check if the current input element is ascii then
   // check if it is set in the bitset
@@ -542,6 +560,15 @@ extension Processor {
         controller.step()
       }
 
+    case .matchUTF8:
+      let (utf8Reg, boundaryCheck) = payload.matchUTF8Payload
+      let utf8Content = registers[utf8Reg]
+      if matchUTF8(
+        utf8Content, boundaryCheck: boundaryCheck
+      ) {
+        controller.step()
+      }
+
     case .matchBitset:
       let (isScalar, reg) = payload.bitsetPayload
       let bitset = registers[reg]
@@ -752,6 +779,27 @@ extension String {
     return idx
   }
 
+  func matchUTF8(
+    _ bytes: Array<UInt8>,
+    at pos: Index,
+    limitedBy end: Index,
+    boundaryCheck: Bool
+  ) -> Index? {
+    var cur = pos
+    for b in bytes {
+      guard cur < end, self.utf8[cur] == b else { return nil }
+      self.utf8.formIndex(after: &cur)
+    }
+
+    guard cur <= end else { return nil }
+
+    if boundaryCheck && !isOnGraphemeClusterBoundary(cur) {
+      return nil
+    }
+
+    return cur
+  }
+
   func matchASCIIBitset(
     _ bitset: DSLTree.CustomCharacterClass.AsciiBitset,
     at pos: Index,

Sources/_StringProcessing/Engine/Registers.swift

@@ -24,11 +24,9 @@ extension Processor {
     // Verbatim elements to compare against
     var elements: [Element]
 
-    // Verbatim sequences to compare against
-    //
-    // TODO: Degenericize Processor and store Strings
-    var sequences: [[Element]] = []
-    
+    // Verbatim bytes to compare against
+    var utf8Contents: [[UInt8]]
+
     var bitsets: [DSLTree.CustomCharacterClass.AsciiBitset]
 
     var consumeFunctions: [MEProgram.ConsumeFunction]
@@ -55,9 +53,6 @@ extension Processor {
 extension Processor.Registers {
   typealias Input = String
 
-  subscript(_ i: SequenceRegister) -> [Input.Element] {
-    sequences[i.rawValue]
-  }
   subscript(_ i: IntRegister) -> Int {
     get { ints[i.rawValue] }
     set {
@@ -82,6 +77,9 @@ extension Processor.Registers {
   subscript(_ i: ElementRegister) -> Input.Element {
     elements[i.rawValue]
   }
+  subscript(_ i: UTF8Register) -> [UInt8] {
+    utf8Contents[i.rawValue]
+  }
   subscript(
     _ i: AsciiBitsetRegister
   ) -> DSLTree.CustomCharacterClass.AsciiBitset {
@@ -110,8 +108,8 @@ extension Processor.Registers {
     self.elements = program.staticElements
     assert(elements.count == info.elements)
 
-    self.sequences = program.staticSequences
-    assert(sequences.count == info.sequences)
+    self.utf8Contents = program.staticUTF8Contents
+    assert(utf8Contents.count == info.utf8Contents)
 
     self.bitsets = program.staticBitsets
     assert(bitsets.count == info.bitsets)
@@ -156,7 +154,7 @@ extension MutableCollection {
 extension MEProgram {
   struct RegisterInfo {
     var elements = 0
-    var sequences = 0
+    var utf8Contents = 0
     var bools = 0
     var strings = 0
     var bitsets = 0

Sources/_StringProcessing/MatchingOptions.swift

@@ -125,7 +125,15 @@ extension MatchingOptions {
       ? .graphemeCluster
       : .unicodeScalar
   }
-  
+
+  /// Whether matching needs to honor canonical equivalence.
+  ///
+  /// Currently, this is synonymous with grapheme-cluster semantics, but could
+  /// become its own option in the future
+  var usesCanonicalEquivalence: Bool {
+    semanticLevel == .graphemeCluster
+  }
+
   var usesNSRECompatibleDot: Bool {
     stack.last!.contains(.nsreCompatibleDot)
   }

Sources/_StringProcessing/Utility/TypedInt.swift

@@ -122,8 +122,9 @@ enum _SavePointAddress {}
 typealias ElementRegister = TypedInt<_ElementRegister>
 enum _ElementRegister {}
 
-typealias SequenceRegister = TypedInt<_SequenceRegister>
-enum _SequenceRegister {}
+/// The register number for a sequence of UTF-8 bytes
+typealias UTF8Register = TypedInt<_UTF8Register>
+enum _UTF8Register {}
 
 /// The register number for a stored boolean value
 ///

Tests/RegexTests/CompileTests.swift

@@ -41,6 +41,7 @@ enum DecodedInstr {
   case matchAnyNonNewline
   case matchBitset
   case matchBuiltin
+  case matchUTF8
   case consumeBy
   case assertBy
   case matchBy
@@ -141,6 +142,8 @@ extension DecodedInstr {
       return .captureValue
     case .matchBuiltin:
       return .matchBuiltin
+    case .matchUTF8:
+      return .matchUTF8
     }
   }
 }
@@ -443,10 +446,30 @@ extension RegexTests {
       contains: [.matchScalarUnchecked],
       doesNotContain: [.match, .consumeBy, .matchScalar])
     expectProgram(
-      for: "aaa\u{301}",
+      for: "a\u{301}",
       semanticLevel: .unicodeScalar,
       contains: [.matchScalarUnchecked],
       doesNotContain: [.match, .consumeBy, .matchScalar])
+    expectProgram(
+      for: "abcdefg",
+      semanticLevel: .unicodeScalar,
+      contains: [.matchUTF8],
+      doesNotContain: [.match, .consumeBy, .matchScalar])
+    expectProgram(
+      for: "abcdefg",
+      semanticLevel: .graphemeCluster,
+      contains: [.matchUTF8],
+      doesNotContain: [.match, .consumeBy, .matchScalar])
+    expectProgram(
+      for: "aaa\u{301}",
+      semanticLevel: .unicodeScalar,
+      contains: [.matchUTF8],
+      doesNotContain: [.match, .consumeBy, .matchScalar])
+    expectProgram(
+      for: "aaa\u{301}",
+      semanticLevel: .graphemeCluster,
+      contains: [.match],
+      doesNotContain: [.matchUTF8, .consumeBy])
   }
 
   func testCaseInsensitivityCompilation() {