Fix matching for non-ASCII-equivalent ASCII chars

Sources/_StringProcessing/ByteCodeGen+DSLList.swift

@@ -694,7 +694,10 @@ fileprivate extension Compiler.ByteCodeGen {
           builder.buildQuantify(bitset: bitset, kind, minTrips, maxExtraTrips, isScalarSemantics: isScalarSemantics)
         } else {
           // Uncased character OR case-sensitive matching; match as a single scalar ascii value character
-          guard let val = c._singleScalarAsciiValue else {
+          guard let val = c._singleScalarAsciiValue,
+                !options.usesCanonicalEquivalence
+                  || !val._hasNonASCIICanonicalEquivalent
+          else {
             return false
           }
           builder.buildQuantify(asciiChar: val, kind, minTrips, maxExtraTrips, isScalarSemantics: isScalarSemantics)

Sources/_StringProcessing/ByteCodeGen.swift

@@ -97,31 +97,69 @@ extension Compiler.ByteCodeGen {
   }
 
   mutating func emitQuotedLiteral(_ s: String) {
-    // ASCII is normalization-invariant, so is the safe subset for
-    // us to optimize
+    guard !s.isEmpty else { return }
+
+    // Matching bytes directly is only safe when matching case sensitively
+    // and one of the following is true:
+    // - we're using unicodeScalar semantics, or
+    // - the literal is ASCII with no non-ASCII scalar equivalents
     if optimizationsEnabled,
-       !options.usesCanonicalEquivalence || s.utf8.allSatisfy(\._isASCII),
-       !s.isEmpty
+       !options.isCaseInsensitive,
+       s._canMatchUTF8Directly(
+         usesCanonicalEquivalence: options.usesCanonicalEquivalence)
     {
+      let boundaryCheck = options.semanticLevel == .graphemeCluster
+      builder.buildMatchUTF8(Array(s.utf8), boundaryCheck: boundaryCheck)
+      return
+    }
 
-      // TODO: Make an optimizations configuration struct, where
-      // we can enable/disable specific optimizations and change
-      // thresholds
+    // Special case: the literal is ASCII, but contains at least one character
+    // that a non-ASCII scalar is equivalent to. Here we can split around the
+    // special characters: for example, `"...TIKEUT..."` is emitted as:
+    //
+    // - `matchUTF8("...TI")`
+    // - `matchScalarEither("K", "\u{212A}")`
+    // - `matchUTF8("EUT...")`
+    if optimizationsEnabled,
+       !options.isCaseInsensitive,
+       s.utf8.allSatisfy(\._isASCII)
+    {
+      assert(options.usesCanonicalEquivalence)
+      assert(s.utf8.contains(where: \._hasNonASCIICanonicalEquivalent))
       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
+
+      var run: [UInt8] = []
+      func flushRun() {
+        if run.count >= longThreshold {
+          builder.buildMatchUTF8(run, boundaryCheck: true)
+        } else {
+          // Only the last byte needs a boundary check, since the run is all ASCII.
+          for byte in run.dropLast() {
+            builder.buildMatchScalar(
+              Unicode.Scalar(byte), boundaryCheck: false)
+          }
+          if let byte = run.last {
+            builder.buildMatchScalar(
+              Unicode.Scalar(byte), boundaryCheck: true)
+          }
+        }
+        run = []
+      }
+
+      for byte in s.utf8 {
+        if byte._hasNonASCIICanonicalEquivalent {
+          flushRun()
+          emitCharacter(Character(Unicode.Scalar(byte)))
+        } else {
+          run.append(byte)
+        }
       }
+      flushRun()
+      return
     }
 
-    guard options.semanticLevel == .graphemeCluster else {
+    // Naive path for unicodeScalar semantics
+    guard options.usesCanonicalEquivalence else {
       for char in s {
         for scalar in char.unicodeScalars {
           emitMatchScalar(scalar)
@@ -131,15 +169,21 @@ extension Compiler.ByteCodeGen {
     }
 
     // Fast path for eliding boundary checks for an all ascii quoted literal
-    if optimizationsEnabled && s.allSatisfy(\.isASCII) && !s.isEmpty {
-      let lastIdx = s.unicodeScalars.indices.last!
-      for idx in s.unicodeScalars.indices {
-        let boundaryCheck = idx == lastIdx
-        let scalar = s.unicodeScalars[idx]
+    if optimizationsEnabled,
+       s._isDirectlyComparableASCII
+    {
+      for scalar in s.unicodeScalars.dropLast() {
+        if options.isCaseInsensitive && scalar.properties.isCased {
+          builder.buildMatchScalarCaseInsensitive(scalar, boundaryCheck: false)
+        } else {
+          builder.buildMatchScalar(scalar, boundaryCheck: false)
+        }
+      }
+      if let scalar = s.unicodeScalars.last {
         if options.isCaseInsensitive && scalar.properties.isCased {
-          builder.buildMatchScalarCaseInsensitive(scalar, boundaryCheck: boundaryCheck)
+          builder.buildMatchScalarCaseInsensitive(scalar, boundaryCheck: true)
         } else {
-          builder.buildMatchScalar(scalar, boundaryCheck: boundaryCheck)
+          builder.buildMatchScalar(scalar, boundaryCheck: true)
         }
       }
       return
@@ -198,7 +242,12 @@ extension Compiler.ByteCodeGen {
   mutating func emitMatchScalar(_ s: UnicodeScalar) {
     assert(options.semanticLevel == .unicodeScalar)
     if options.isCaseInsensitive && s.properties.isCased {
-      builder.buildMatchScalarCaseInsensitive(s, boundaryCheck: false)
+      if s.isASCII {
+        let otherValue = s.value ^ 32
+        builder.buildMatchScalarEither(s, orScalar: UnicodeScalar(otherValue)!, boundaryCheck: false)
+      } else {
+        builder.buildMatchScalarCaseInsensitive(s, boundaryCheck: false)
+      }
     } else {
       builder.buildMatchScalar(s, boundaryCheck: false)
     }
@@ -227,14 +276,28 @@ extension Compiler.ByteCodeGen {
       return
     }
 
-    if optimizationsEnabled && c.isASCII {
+    if optimizationsEnabled, c.isASCII,
+       !c._containsASCIIWithNonASCIIEquivalent {
       let lastIdx = c.unicodeScalars.indices.last!
       for idx in c.unicodeScalars.indices {
         builder.buildMatchScalar(c.unicodeScalars[idx], boundaryCheck: idx == lastIdx)
       }
       return
     }
-
+
+    // An ASCII character with a non-ASCII scalar canonically equivalent to it
+    // can still be matched as a scalar, by accepting either one, rather than
+    // falling back to materializing and comparing `Character`s.
+    if optimizationsEnabled,
+       let ascii = c._singleScalarAsciiValue,
+       let equivalent = ascii._nonASCIICanonicalEquivalent {
+      builder.buildMatchScalarEither(
+        Unicode.Scalar(ascii),
+        orScalar: equivalent,
+        boundaryCheck: true)
+      return
+    }
+
     builder.buildMatch(c, isCaseInsensitive: false)
   }
 

Sources/_StringProcessing/Engine/InstPayload.swift

@@ -168,6 +168,31 @@ extension Instruction.Payload {
     return (scalar, caseInsensitive: caseInsensitive, boundaryCheck: boundaryCheck)
   }
 
+  /// Two scalars plus a boundary-check flag.
+  ///
+  /// The first scalar occupies the low 21 bits, as in `scalarPayload`; the
+  /// second goes in bits 32-52, which that layout leaves free. The
+  /// boundary-check flag stays at bit 54 so it lines up with `scalarPayload`.
+  init(
+    scalar: Unicode.Scalar,
+    orScalar other: Unicode.Scalar,
+    boundaryCheck: Bool
+  ) {
+    let raw = UInt64(scalar.value)
+      | (UInt64(other.value) << 32)
+      | (boundaryCheck ? 1 << 54 : 0)
+    self.init(raw)
+  }
+  var scalarEitherPayload: (
+    Unicode.Scalar, Unicode.Scalar, boundaryCheck: Bool
+  ) {
+    let boundaryCheck = (self.rawValue >> 54) & 1 == 1
+    let first = Unicode.Scalar(_value: UInt32(self.rawValue & 0x1F_FFFF))
+    let second = Unicode.Scalar(
+      _value: UInt32((self.rawValue >> 32) & 0x1F_FFFF))
+    return (first, second, boundaryCheck: boundaryCheck)
+  }
+
   init(utf8: UTF8Register, boundaryCheck: Bool) {
     self.init(boundaryCheck ? 1 : 0, utf8)
   }

Sources/_StringProcessing/Engine/Instruction.swift

@@ -112,6 +112,13 @@ extension Instruction {
     /// Operands: Scalar value to match against and booleans
     case matchScalar
 
+    /// Match against one of two scalars, and possibly perform a boundary check.
+    ///
+    ///     matchScalarEither(_: Unicode.Scalar, _: Unicode.Scalar, boundaryCheck: Bool)
+    ///
+    /// Operands: the two scalar values to match against and a boolean
+    case matchScalarEither
+
     /// Match directly (binary semantics) against a series of UTF-8 bytes
     ///
     /// NOTE: Compiler should ensure to only emit this instruction when normalization

Sources/_StringProcessing/Engine/MEBuilder.swift

@@ -207,6 +207,17 @@ extension MEProgram.Builder {
     instructions.append(.init(.matchScalar, .init(scalar: s, caseInsensitive: false, boundaryCheck: boundaryCheck)))
   }
 
+  /// Match `s`, or the non-ASCII scalar canonically equivalent to it.
+  mutating func buildMatchScalarEither(
+    _ s: Unicode.Scalar,
+    orScalar other: Unicode.Scalar,
+    boundaryCheck: Bool
+  ) {
+    instructions.append(.init(
+      .matchScalarEither,
+      .init(scalar: s, orScalar: other, boundaryCheck: boundaryCheck)))
+  }
+
   mutating func buildMatchScalarCaseInsensitive(_ s: Unicode.Scalar, boundaryCheck: Bool) {
     instructions.append(.init(.matchScalar, .init(scalar: s, caseInsensitive: true, boundaryCheck: boundaryCheck)))
   }

Sources/_StringProcessing/Engine/Processor.swift

@@ -294,6 +294,24 @@ extension Processor {
     return true
   }
 
+  mutating func matchScalarEither(
+    _ first: Unicode.Scalar,
+    _ second: Unicode.Scalar,
+    boundaryCheck: Bool
+  ) -> Bool {
+    guard let next = input.matchScalarEither(
+      first, second,
+      at: currentPosition,
+      limitedBy: end,
+      boundaryCheck: boundaryCheck
+    ) else {
+      signalFailure()
+      return false
+    }
+    currentPosition = next
+    return true
+  }
+
   mutating func matchScalar(
     _ s: Unicode.Scalar,
     boundaryCheck: Bool,
@@ -587,6 +605,14 @@ extension Processor {
         controller.step()
       }
 
+    case .matchScalarEither:
+      let (first, second, boundaryCheck) = payload.scalarEitherPayload
+      if matchScalarEither(
+        first, second, boundaryCheck: boundaryCheck
+      ) {
+        controller.step()
+      }
+
     case .matchUTF8:
       let (utf8Reg, boundaryCheck) = payload.matchUTF8Payload
       let utf8Content = registers[utf8Reg]
@@ -775,6 +801,28 @@ extension String {
     return cur
   }
 
+  /// Matches either `first` or `second` as a single scalar.
+  func matchScalarEither(
+    _ first: Unicode.Scalar,
+    _ second: Unicode.Scalar,
+    at pos: Index,
+    limitedBy end: String.Index,
+    boundaryCheck: Bool
+  ) -> Index? {
+    guard pos < end else { return nil }
+    let curScalar = unicodeScalars[pos]
+    guard curScalar == first || curScalar == second else { return nil }
+
+    let idx = unicodeScalars.index(after: pos)
+    assert(idx <= end, "Input is a substring with a sub-scalar endIndex.")
+
+    if boundaryCheck && !isOnGraphemeClusterBoundary(idx) {
+      return nil
+    }
+
+    return idx
+  }
+
   func matchScalar(
     _ scalar: Unicode.Scalar,
     at pos: Index,