Limit bitset payloads in quantifier instruction

Sources/_StringProcessing/ByteCodeGen+DSLList.swift

@@ -674,7 +674,8 @@ fileprivate extension Compiler.ByteCodeGen {
       guard let bitset = ccc.asAsciiBitset(options) else {
         return false
       }
-      builder.buildQuantify(bitset: bitset, kind, minTrips, maxExtraTrips, isScalarSemantics: isScalarSemantics)
+      guard builder.buildQuantify(bitset: bitset, kind, minTrips, maxExtraTrips, isScalarSemantics: isScalarSemantics)
+      else { return false }
 
     case .atom(let atom):
       switch atom {
@@ -684,7 +685,8 @@ fileprivate extension Compiler.ByteCodeGen {
           guard let bitset = DSLTree.CustomCharacterClass(members: [.atom(atom)]).asAsciiBitset(options) else {
             return false
           }
-          builder.buildQuantify(bitset: bitset, kind, minTrips, maxExtraTrips, isScalarSemantics: isScalarSemantics)
+          guard builder.buildQuantify(bitset: bitset, kind, minTrips, maxExtraTrips, isScalarSemantics: isScalarSemantics)
+          else { return false }
         } else {
           // Uncased character OR case-sensitive matching; match as a single scalar ascii value character
           guard let val = c._singleScalarAsciiValue else {

Sources/_StringProcessing/Engine/InstPayload.swift

@@ -96,22 +96,23 @@ extension Instruction.Payload {
   // for 1990s Unicode it's good enough for us.
   //
   // TODO: but really, let's come up with something
-  private var firstSplitMask: UInt64 { 0x0000_FFFF }
-  private var secondSplitMask: UInt64 { 0xFFFF_0000 }
+  private static var firstSplitMask: UInt64 { 0x00FF_FFFF }
+  private static var secondSplitMask: UInt64 { 0xFFFF_FF00_0000 }
+  private static var secondSplitShift: UInt64 { 24 }
 
   private var split: (first: UInt64, second: UInt64) {
-    assert(rawValue == ((firstSplitMask|secondSplitMask) & rawValue))
+    assert(rawValue == ((Self.firstSplitMask|Self.secondSplitMask) & rawValue))
 
     // TODO: Which order is better?
-    let first = rawValue & firstSplitMask
-    let second = (rawValue & secondSplitMask) &>> 16
+    let first = rawValue & Self.firstSplitMask
+    let second = rawValue &>> Self.secondSplitShift
     return (first, second)
   }
 
   private init(_ a: UInt64, _ b: UInt64) {
-    self.init(a | (b &<< 16))
-    assert(a == a & firstSplitMask)
-    assert(b == b & firstSplitMask)
+    self.init(a | (b &<< Self.secondSplitShift))
+    assert(a == a & Self.firstSplitMask)
+    assert(b &<< Self.secondSplitShift == (b &<< Self.secondSplitShift) & Self.secondSplitMask)
   }
   private init<👻>(_ a: UInt64, _ b: TypedInt<👻>) {
     self.init(a, b.bits)
@@ -410,11 +411,11 @@ struct QuantifyPayload: RawRepresentable {
   static var maxStorableTrips: UInt64 { (1 << 8) - 1 }
   static var isScalarSemanticsBit: UInt64 { 1 &<< 38 }
 
-  var quantKindMask: UInt64  { 3 }
-  var maxExtraTripsMask: UInt64 { 0x1FF }
-  var minTripsMask: UInt64   { 0xFF }
-  var typeMask: UInt64       { 7 }
-  var payloadMask: UInt64    { 0xFF_FF }
+  static var quantKindMask: UInt64  { 3 }
+  static var maxExtraTripsMask: UInt64 { 0x1FF }
+  static var minTripsMask: UInt64   { 0xFF }
+  static var typeMask: UInt64       { 7 }
+  static var quantifyPayloadMask: UInt64    { 0xFF_FF }
 
   // Calculate the maximum number of trips, else UInt64.max if unbounded
   var maxTrips: UInt64 {
@@ -466,7 +467,7 @@ struct QuantifyPayload: RawRepresentable {
     _ maxExtraTrips: Int?,
     isScalarSemantics: Bool
   ) {
-    assert(bitset.bits <= _payloadMask)
+    assert(bitset.bits <= Self.quantifyPayloadMask)
     self.rawValue = bitset.bits
       + QuantifyPayload.packInfoValues(kind, minTrips, maxExtraTrips, .asciiBitset, isScalarSemantics: isScalarSemantics)
   }
@@ -513,7 +514,7 @@ struct QuantifyPayload: RawRepresentable {
   }
 
   var quantKind: AST.Quantification.Kind {
-    switch (self.rawValue >> QuantifyPayload.quantKindShift) & quantKindMask {
+    switch (self.rawValue >> Self.quantKindShift) & Self.quantKindMask {
     case 0: return .eager
     case 1: return .reluctant
     case 2: return .possessive
@@ -523,11 +524,11 @@ struct QuantifyPayload: RawRepresentable {
   }
 
   var minTrips: UInt64 {
-    (self.rawValue >> QuantifyPayload.minTripsShift) & minTripsMask
+    (self.rawValue >> Self.minTripsShift) & Self.minTripsMask
   }
 
   var maxExtraTrips: UInt64? {
-    let val = (self.rawValue >> QuantifyPayload.maxExtraTripsShift) & maxExtraTripsMask
+    let val = (self.rawValue >> Self.maxExtraTripsShift) & Self.maxExtraTripsMask
     if val == 1 {
       return nil
     } else {
@@ -540,11 +541,11 @@ struct QuantifyPayload: RawRepresentable {
   }
 
   var bitset: AsciiBitsetRegister {
-    TypedInt(self.rawValue & payloadMask)
+    TypedInt(self.rawValue & Self.quantifyPayloadMask)
   }
 
   var asciiChar: UInt8 {
-    UInt8(asserting: self.rawValue & payloadMask)
+    UInt8(asserting: self.rawValue & Self.quantifyPayloadMask)
   }
 
   var anyMatchesNewline: Bool {

Sources/_StringProcessing/Engine/MEBuilder.swift

@@ -215,15 +215,21 @@ extension MEProgram.Builder {
   mutating func buildMatchAsciiBitset(
     _ b: DSLTree.CustomCharacterClass.AsciiBitset
   ) {
+    guard let bitset = makeAsciiBitset(b, indexLimit: Int(_payloadMask)) else {
+      fatalError("Bitset exceeded maximum count")
+    }
     instructions.append(.init(
-      .matchBitset, .init(bitset: makeAsciiBitset(b), isScalar: false)))
+      .matchBitset, .init(bitset: bitset, isScalar: false)))
   }
 
   mutating func buildScalarMatchAsciiBitset(
     _ b: DSLTree.CustomCharacterClass.AsciiBitset
   ) {
+    guard let bitset = makeAsciiBitset(b, indexLimit: Int(_payloadMask)) else {
+      fatalError("Bitset exceeded maximum count")
+    }
     instructions.append(.init(
-      .matchBitset, .init(bitset: makeAsciiBitset(b), isScalar: true)))
+      .matchBitset, .init(bitset: bitset, isScalar: true)))
   }
 
   mutating func buildMatchBuiltin(model: _CharacterClassModel) {
@@ -262,10 +268,14 @@ extension MEProgram.Builder {
     _ minTrips: Int,
     _ maxExtraTrips: Int?,
     isScalarSemantics: Bool
-  ) {
+  ) -> Bool {
+    guard let bitsetRegister = makeAsciiBitset(bitset, indexLimit: Int(QuantifyPayload.quantifyPayloadMask)) else {
+      return false
+    }
     instructions.append(.init(
       .quantify,
-      .init(quantify: .init(bitset: makeAsciiBitset(bitset), kind, minTrips, maxExtraTrips, isScalarSemantics: isScalarSemantics))))
+      .init(quantify: .init(bitset: bitsetRegister, kind, minTrips, maxExtraTrips, isScalarSemantics: isScalarSemantics))))
+    return true
   }
 
   mutating func buildQuantify(
@@ -578,8 +588,10 @@ extension MEProgram.Builder {
   // registers without monotonicity required
 
   mutating func makeAsciiBitset(
-    _ b: DSLTree.CustomCharacterClass.AsciiBitset
-  ) -> AsciiBitsetRegister {
+    _ b: DSLTree.CustomCharacterClass.AsciiBitset,
+    indexLimit: Int
+  ) -> AsciiBitsetRegister? {
+    guard asciiBitsets.count < indexLimit else { return nil }
     defer { asciiBitsets.append(b) }
     return AsciiBitsetRegister(asciiBitsets.count)
   }

Tests/RegexTests/CompileTests.swift

@@ -405,6 +405,21 @@ extension RegexTests {
       contains: [.consumeBy],
       doesNotContain: [.matchBuiltin, .matchBitset, .matchBitsetScalar])
 
+    // A bitset can be a payload in a quantify instruction...
+    expectProgram(
+      for: "[ab][ab][ab][ab][cd]+",
+      contains: [.matchBitset, .quantify],
+      doesNotContain: [.match, .consumeBy])
+    expectProgram(
+      for: String(repeating: "[ab]", count: 65534) + "[cd]+",
+      contains: [.matchBitset, .quantify],
+      doesNotContain: [.match, .consumeBy])
+    // ...But not if its register index is larger than a quantify payload
+    expectProgram(
+      for: String(repeating: "[ab]", count: 65535) + "[cd]+",
+      contains: [.matchBitset],
+      doesNotContain: [.match, .consumeBy, .quantify])
+
     // Must have new stdlib for character class ranges.
     guard ensureNewStdlib() else { return }
 

Tests/RegexTests/MatchTests.swift

@@ -1136,7 +1136,6 @@ extension RegexTests {
       ("j", "j")
     )
 
-
     // These can't compile in grapheme semantic mode, but make sure they work in
     // scalar semantic mode.
     firstMatchTests(
@@ -3015,4 +3014,22 @@ extension RegexTests {
     XCTAssertNotNil("testS".wholeMatch(of: regex))
     XCTAssertNotNil("tesTs".wholeMatch(of: regex))
   }
+
+  func testQuantifierPayload() throws {
+    // rdar://185730811
+    func test(_ n: Int, line: UInt = #line) throws {
+      let pattern = "^" + String(repeating: "[ac]", count: n) + "[bd]+"
+      let regex = try Regex(pattern)
+      let goodInput = String(repeating: "a", count: n) + "b"
+      let badInput = String(repeating: "a", count: n + 1)
+
+      XCTAssertTrue(goodInput.contains(regex), "Missed match in 'a....b'", line: line)
+      XCTAssertFalse(badInput.contains(regex), "Incorrect match in 'a....'", line: line)
+    }
+
+    try test(4)
+    try test(65534)
+    try test(65535)
+    try test(65536)
+  }
 }