Use a bitset for ascii-only character classes

Sources/_StringProcessing/ByteCodeGen.swift

@@ -643,8 +643,12 @@ fileprivate extension Compiler.ByteCodeGen {
   mutating func emitCustomCharacterClass(
     _ ccc: DSLTree.CustomCharacterClass
   ) throws {
-    let consumer = try ccc.generateConsumer(options)
-    builder.buildConsume(by: consumer)
+    if let asciiBitset = ccc.asAsciiBitset(options) {
+      builder.buildMatchAsciiBitset(asciiBitset)
+    } else {
+      let consumer = try ccc.generateConsumer(options)
+      builder.buildConsume(by: consumer)
+    }
   }
 
   @discardableResult

Sources/_StringProcessing/ConsumerInterface.swift

@@ -51,7 +51,26 @@ extension DSLTree.Node {
   }
 }
 
+extension DSLTree._AST.Atom {
+  var asciiValue: UInt8? {
+    return ast.asciiValue
+  }
+}
+
 extension DSLTree.Atom {
+  var asciiValue: UInt8? {
+    switch self {
+    case let .char(c) where c != "\r\n":
+      return c.asciiValue
+    case let .scalar(s) where s.isASCII:
+      return UInt8(ascii: s)
+    case let .unconverted(atom):
+      return atom.asciiValue
+    default:
+      return nil
+    }
+  }
+
   // TODO: If ByteCodeGen switches first, then this is unnecessary for
   // top-level nodes, but it's also invoked for `.atom` members of a custom CC
   func generateConsumer(
@@ -177,7 +196,18 @@ extension AST.Atom {
     default: return nil
     }
   }
-
+
+  var asciiValue: UInt8? {
+    switch kind {
+    case let .char(c) where c != "\r\n":
+      return c.asciiValue
+    case let .scalar(s) where s.value.isASCII:
+      return UInt8(ascii: s.value)
+    default:
+      return nil
+    }
+  }
+
   func generateConsumer(
     _ opts: MatchingOptions
   ) throws -> MEProgram.ConsumeFunction? {
@@ -235,6 +265,34 @@ extension AST.Atom {
 }
 
 extension DSLTree.CustomCharacterClass.Member {
+  func asAsciiBitset(
+    _ opts: MatchingOptions,
+    _ isInverted: Bool
+  ) -> DSLTree.CustomCharacterClass.AsciiBitset? {
+    switch self {
+    case let .atom(a):
+      if let val = a.asciiValue {
+        return DSLTree.CustomCharacterClass.AsciiBitset(
+          val,
+          isInverted,
+          opts.isCaseInsensitive
+        )
+      }
+    case let .range(low, high):
+      if let lowVal = low.asciiValue, let highVal = high.asciiValue {
+        return DSLTree.CustomCharacterClass.AsciiBitset(
+          low: lowVal,
+          high: highVal,
+          isInverted: isInverted,
+          isCaseInsensitive: opts.isCaseInsensitive
+        )
+      }
+    default:
+      return nil
+    }
+    return nil
+  }
+
   func generateConsumer(
     _ opts: MatchingOptions
   ) throws -> MEProgram.ConsumeFunction {
@@ -342,6 +400,19 @@ extension DSLTree.CustomCharacterClass.Member {
 }
 
 extension DSLTree.CustomCharacterClass {
+  func asAsciiBitset(_ opts: MatchingOptions) -> AsciiBitset? {
+    return members.reduce(
+      .init(isInverted: isInverted),
+      {result, member in
+        if let next = member.asAsciiBitset(opts, isInverted) {
+          return result?.union(next)
+        } else {
+          return nil
+        }
+      }
+    )
+  }
+
   func generateConsumer(
     _ opts: MatchingOptions
   ) throws -> MEProgram.ConsumeFunction {

Sources/_StringProcessing/Engine/InstPayload.swift

@@ -50,6 +50,7 @@ extension Instruction.Payload {
     case bool(BoolRegister)
     case element(ElementRegister)
     case consumer(ConsumeFunctionRegister)
+    case bitset(AsciiBitsetRegister)
     case assertion(AssertionFunctionRegister)
     case addr(InstructionAddress)
     case capture(CaptureRegister)
@@ -196,6 +197,13 @@ extension Instruction.Payload {
     interpret()
   }
 
+  init(bitset: AsciiBitsetRegister) {
+    self.init(bitset)
+  }
+  var bitset: AsciiBitsetRegister {
+    interpret()
+  }
+
   init(consumer: ConsumeFunctionRegister) {
     self.init(consumer)
   }

Sources/_StringProcessing/Engine/Instruction.swift

@@ -84,6 +84,10 @@ extension Instruction {
     /// Operand: Sequence register to compare against.
     case matchSequence
 
+    /// Match against a set of valid ascii values stored in a bitset
+    /// Operand: Ascii bitset register containing the bitset
+    case matchBitset
+
     /// TODO: builtin assertions and anchors
     case builtinAssertion
 

Sources/_StringProcessing/Engine/MEBuilder.swift

@@ -18,6 +18,7 @@ extension MEProgram {
     var elements = TypedSetVector<Input.Element, _ElementRegister>()
     var sequences = TypedSetVector<[Input.Element], _SequenceRegister>()
 
+    var asciiBitsets: [DSLTree.CustomCharacterClass.AsciiBitset] = []
     var consumeFunctions: [ConsumeFunction] = []
     var assertionFunctions: [AssertionFunction] = []
     var transformFunctions: [TransformFunction] = []
@@ -147,6 +148,15 @@ extension MEProgram.Builder {
       .init(sequence: sequences.store(.init(s)))))
   }
 
+  mutating func buildMatchAsciiBitset(
+    _ b: DSLTree.CustomCharacterClass.AsciiBitset
+  ) {
+    instructions.append(
+      Instruction.init(
+        Instruction.OpCode.matchBitset,
+        Instruction.Payload.init(bitset: makeAsciiBitset(b))))
+  }
+
   mutating func buildConsume(
     by p: @escaping MEProgram.ConsumeFunction
   ) {
@@ -273,6 +283,7 @@ extension MEProgram.Builder {
     regInfo.sequences = sequences.count
     regInfo.ints = nextIntRegister.rawValue
     regInfo.values = nextValueRegister.rawValue
+    regInfo.bitsets = asciiBitsets.count
     regInfo.consumeFunctions = consumeFunctions.count
     regInfo.assertionFunctions = assertionFunctions.count
     regInfo.transformFunctions = transformFunctions.count
@@ -283,6 +294,7 @@ extension MEProgram.Builder {
       instructions: InstructionList(instructions),
       staticElements: elements.stored,
       staticSequences: sequences.stored,
+      staticBitsets: asciiBitsets,
       staticConsumeFunctions: consumeFunctions,
       staticAssertionFunctions: assertionFunctions,
       staticTransformFunctions: transformFunctions,
@@ -414,6 +426,13 @@ extension MEProgram.Builder {
   // TODO: A register-mapping helper struct, which could release
   // registers without monotonicity required
 
+  mutating func makeAsciiBitset(
+    _ b: DSLTree.CustomCharacterClass.AsciiBitset
+  ) -> AsciiBitsetRegister {
+    defer { asciiBitsets.append(b) }
+    return AsciiBitsetRegister(asciiBitsets.count)
+  }
+
   mutating func makeConsumeFunction(
     _ f: @escaping MEProgram.ConsumeFunction
   ) -> ConsumeFunctionRegister {

Sources/_StringProcessing/Engine/MEProgram.swift

@@ -26,6 +26,7 @@ struct MEProgram {
 
   var staticElements: [Input.Element]
   var staticSequences: [[Input.Element]]
+  var staticBitsets: [DSLTree.CustomCharacterClass.AsciiBitset]
   var staticConsumeFunctions: [ConsumeFunction]
   var staticAssertionFunctions: [AssertionFunction]
   var staticTransformFunctions: [TransformFunction]

Sources/_StringProcessing/Engine/Processor.swift

@@ -226,6 +226,20 @@ extension Processor {
     }
     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
+  mutating func matchBitset(
+    _ bitset: DSLTree.CustomCharacterClass.AsciiBitset
+  ) -> Bool {
+    guard let cur = load(), bitset.matches(char: cur) else {
+      signalFailure()
+      return false
+    }
+    _uncheckedForcedConsumeOne()
+    return true
+  }
 
   mutating func signalFailure() {
     guard let (pc, pos, stackEnd, capEnds, intRegisters) =
@@ -364,6 +378,13 @@ extension Processor {
         controller.step()
       }
 
+    case .matchBitset:
+      let reg = payload.bitset
+      let bitset = registers[reg]
+      if matchBitset(bitset) {
+        controller.step()
+      }
+
     case .consumeBy:
       let reg = payload.consumer
       guard currentPosition < searchBounds.upperBound,

Sources/_StringProcessing/Engine/Registers.swift

@@ -28,6 +28,8 @@ extension Processor {
     //
     // TODO: Degenericize Processor and store Strings
     var sequences: [[Element]] = []
+
+    var bitsets: [DSLTree.CustomCharacterClass.AsciiBitset]
 
     var consumeFunctions: [MEProgram.ConsumeFunction]
 
@@ -67,6 +69,11 @@ extension Processor.Registers {
   subscript(_ i: ElementRegister) -> Input.Element {
     elements[i.rawValue]
   }
+  subscript(
+    _ i: AsciiBitsetRegister
+  ) -> DSLTree.CustomCharacterClass.AsciiBitset {
+    bitsets[i.rawValue]
+  }
   subscript(_ i: ConsumeFunctionRegister) -> MEProgram.ConsumeFunction {
     consumeFunctions[i.rawValue]
   }
@@ -94,6 +101,9 @@ extension Processor.Registers {
     self.sequences = program.staticSequences
     assert(sequences.count == info.sequences)
 
+    self.bitsets = program.staticBitsets
+    assert(bitsets.count == info.bitsets)
+
     self.consumeFunctions = program.staticConsumeFunctions
     assert(consumeFunctions.count == info.consumeFunctions)
 
@@ -133,6 +143,7 @@ extension MEProgram {
     var sequences = 0
     var bools = 0
     var strings = 0
+    var bitsets = 0
     var consumeFunctions = 0
     var assertionFunctions = 0
     var transformFunctions = 0

Sources/_StringProcessing/Regex/DSLTree.swift

@@ -164,6 +164,87 @@ extension DSLTree {
       indirect case subtraction(CustomCharacterClass, CustomCharacterClass)
       indirect case symmetricDifference(CustomCharacterClass, CustomCharacterClass)
     }
+
+    public struct AsciiBitset {
+      let isInverted: Bool
+      var a: UInt64 = 0
+      var b: UInt64 = 0
+
+      init(isInverted: Bool) {
+        self.isInverted = isInverted
+      }
+
+      init(_ val: UInt8, _ isInverted: Bool, _ isCaseInsensitive: Bool) {
+        self.isInverted = isInverted
+        add(val, isCaseInsensitive)
+      }
+
+      init(low: UInt8, high: UInt8, isInverted: Bool, isCaseInsensitive: Bool) {
+        self.isInverted = isInverted
+        for val in low...high {
+          add(val, isCaseInsensitive)
+        }
+      }
+
+      internal init(
+        a: UInt64,
+        b: UInt64,
+        isInverted: Bool
+      ) {
+        self.isInverted = isInverted
+        self.a = a
+        self.b = b
+      }
+
+      internal mutating func add(_ val: UInt8, _ isCaseInsensitive: Bool) {
+        setBit(val)
+        if isCaseInsensitive {
+          if val >= 64 && val <= 90 {
+            setBit(val + 32)
+          }
+          if val >= 97 && val <= 122 {
+            setBit(val - 32)
+          }
+        }
+      }
+
+      internal mutating func setBit(_ val: UInt8) {
+        if val < 64 {
+          a = a | 1 << val
+        } else {
+          b = b | 1 << (val - 64)
+        }
+      }
+
+      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
+          }
+        } else {
+          ret = false
+        }
+
+        if isInverted {
+          return !ret
+        }
+
+        return ret
+      }
+
+      /// Joins another bitset from a Member of the same CustomCharacterClass
+      internal func union(_ other: AsciiBitset) -> AsciiBitset {
+        precondition(self.isInverted == other.isInverted)
+        return AsciiBitset(
+          a: self.a | other.a,
+          b: self.b | other.b,
+          isInverted: self.isInverted
+        )
+      }
+    }
   }
 
   @_spi(RegexBuilder)

Sources/_StringProcessing/Utility/TypedInt.swift

@@ -134,6 +134,10 @@ enum _BoolRegister {}
 typealias StringRegister = TypedInt<_StringRegister>
 enum _StringRegister {}
 
+/// Used for matching sets of ascii values via bitsets
+typealias AsciiBitsetRegister = TypedInt<_AsciiBitsetRegister>
+enum _AsciiBitsetRegister {}
+
 /// Used for consume functions, e.g. character classes
 typealias ConsumeFunctionRegister = TypedInt<_ConsumeFunctionRegister>
 enum _ConsumeFunctionRegister {}