Add scalar mode support for matching bitsets + fix bug

Author: rctcwyvrn

Sources/_StringProcessing/ByteCodeGen.swift

@@ -665,11 +665,12 @@ fileprivate extension Compiler.ByteCodeGen {
     _ ccc: DSLTree.CustomCharacterClass
   ) throws {
     if let asciiBitset = ccc.asAsciiBitset(options),
-        options.semanticLevel == .graphemeCluster,
         optimizationsEnabled {
-      // future work: add a bit to .matchBitset to consume either a character
-      // or a scalar so we can have this optimization in scalar mode
-      builder.buildMatchAsciiBitset(asciiBitset)
+      if options.semanticLevel == .unicodeScalar {
+        builder.buildScalarMatchAsciiBitset(asciiBitset)
+      } else {
+        builder.buildMatchAsciiBitset(asciiBitset)
+      }
     } else {
       let consumer = try ccc.generateConsumer(options)
       builder.buildConsume(by: consumer)

Sources/_StringProcessing/ConsumerInterface.swift

@@ -11,6 +11,16 @@
 
 @_implementationOnly import _RegexParser
 
+extension Character {
+  var singleScalarAsciiValue: UInt8? {
+    if let val = asciiValue, self != "\r\n" {
+      return val
+    }
+    return nil
+  }
+}
+
+
 extension DSLTree.Node {
   /// Attempt to generate a consumer from this AST node
   ///
@@ -60,8 +70,8 @@ extension DSLTree._AST.Atom {
 extension DSLTree.Atom {
   var singleScalarASCIIValue: UInt8? {
     switch self {
-    case let .char(c) where c != "\r\n":
-      return c.asciiValue
+    case let .char(c):
+      return c.singleScalarAsciiValue
     case let .scalar(s) where s.isASCII:
       return UInt8(ascii: s)
     case let .unconverted(atom):
@@ -214,8 +224,8 @@ extension AST.Atom {
 
   var singleScalarASCIIValue: UInt8? {
     switch kind {
-    case let .char(c) where c != "\r\n":
-      return c.asciiValue
+    case let .char(c):
+      return c.singleScalarAsciiValue
     case let .scalar(s) where s.value.isASCII:
       return UInt8(ascii: s.value)
     default:

Sources/_StringProcessing/Engine/Instruction.swift

@@ -90,6 +90,8 @@ extension Instruction {
     /// Match against a set of valid ascii values stored in a bitset
     /// Operand: Ascii bitset register containing the bitset
     case matchBitset
+    /// Match a single scalar instead of a character
+    case matchBitsetScalar
 
     /// TODO: builtin assertions and anchors
     case builtinAssertion

Sources/_StringProcessing/Engine/MEBuilder.swift

@@ -163,6 +163,13 @@ extension MEProgram.Builder {
       .matchBitset, .init(bitset: makeAsciiBitset(b))))
   }
 
+  mutating func buildScalarMatchAsciiBitset(
+    _ b: DSLTree.CustomCharacterClass.AsciiBitset
+  ) {
+    instructions.append(.init(
+      .matchBitsetScalar, .init(bitset: makeAsciiBitset(b))))
+  }
+
   mutating func buildConsume(
     by p: @escaping MEProgram.ConsumeFunction
   ) {

Sources/_StringProcessing/Engine/Processor.swift

@@ -255,8 +255,20 @@ extension Processor {
   // ascii characters, so check if the current input element is ascii then
   // check if it is set in the bitset
   mutating func matchBitset(
-    _ bitset: DSLTree.CustomCharacterClass.AsciiBitset
+    _ bitset: DSLTree.CustomCharacterClass.AsciiBitset,
+    scalar: Bool
   ) -> Bool {
+    if scalar {
+      guard let curScalar = loadScalar(),
+              bitset.matches(scalar: curScalar),
+            let idx = nextScalarIndex(offsetBy: 1, boundaryCheck: false) else {
+        signalFailure()
+        return false
+      }
+      currentPosition = idx
+      return true
+    }
+    
     guard let cur = load(), bitset.matches(char: cur) else {
       signalFailure()
       return false
@@ -416,7 +428,13 @@ extension Processor {
     case .matchBitset:
       let reg = payload.bitset
       let bitset = registers[reg]
-      if matchBitset(bitset) {
+      if matchBitset(bitset, scalar: false) {
+        controller.step()
+      }
+    case .matchBitsetScalar:
+      let reg = payload.bitset
+      let bitset = registers[reg]
+      if matchBitset(bitset, scalar: true) {
         controller.step()
       }
 

Sources/_StringProcessing/Regex/DSLTree.swift

@@ -215,23 +215,41 @@ extension DSLTree {
         }
       }
 
+      private func matches(_ val: UInt8) -> Bool {
+        if val < 64 {
+          return (a >> val) & 1 == 1
+        } else {
+          return (b >> (val - 64)) & 1 == 1
+        }
+      }
+
       internal func matches(char: Character) -> Bool {
-        let ret: Bool
-        if let val = char.asciiValue {
-          if val < 64 {
-            ret = (a >> val) & 1 == 1
-          } else {
-            ret =  (b >> (val - 64)) & 1 == 1
-          }
+        let matched: Bool
+        if let val = char.singleScalarAsciiValue {
+          matched = matches(val)
         } else {
-          ret = false
+          matched = false
         }
 
         if isInverted {
-          return !ret
+          return !matched
+        }
+        return matched
+      }
+
+      internal func matches(scalar: Unicode.Scalar) -> Bool {
+        let matched: Bool
+        if scalar.isASCII {
+          let val = UInt8(ascii: scalar)
+          matched = matches(val)
+        } else {
+          matched = false
         }
 
-        return ret
+        if isInverted {
+          return !matched
+        }
+        return matched
       }
 
       /// Joins another bitset from a Member of the same CustomCharacterClass

Tests/RegexTests/CompileTests.swift

@@ -206,6 +206,6 @@ extension RegexTests {
     expectProgram(for: "[abc]", doesNotContain: [.consumeBy])
 
     expectProgram(for: "[abc]", semanticLevel: .unicodeScalar, doesNotContain: [.matchBitset])
-    expectProgram(for: "[abc]", semanticLevel: .unicodeScalar, contains: [.consumeBy])
+    expectProgram(for: "[abc]", semanticLevel: .unicodeScalar, contains: [.matchBitsetScalar])
   }
 }

Tests/RegexTests/MatchTests.swift

@@ -633,7 +633,8 @@ extension RegexTests {
       ("\r", true))
     matchTest("[\n\r]",
       ("\n", true),
-      ("\r", true))
+      ("\r", true),
+      ("\r\n", false))
 
     firstMatchTest("[-]", input: "123-abcxyz", match: "-")