fast/foreignNormalize functions

Author: Lance Parker

stdlib/public/core/StringNormalization.swift

@@ -13,7 +13,6 @@
 import SwiftShims
 
 internal enum _Normalization {
-
   // ICU's NFC unorm2 instance
   //
   // TODO(String performance): Should we cache one on TLS? Is this an expensive
@@ -102,3 +101,258 @@ internal func _tryNormalize(
   }
   return numericCast(count)
 }
+
+internal enum NormalizationResult {
+  case success(SuccessResult)
+  case bufferTooSmall(BufferResizeRequest) // The size needed to normalize the rest of the string
+  
+  struct SuccessResult {
+    var amountFilled: Int
+    var nextReadPosition: String.Index
+  }
+  struct BufferResizeRequest {
+    var newOutputBufferSize: Int
+    var newPreNormalScratchBufferSize: Int
+    var newPostNormalScratchBufferSize: Int
+  }
+  
+  static func bufferTooSmall(count: Int) -> NormalizationResult {
+    let outputBufferSize = count * 9
+    let preNormalBufferSize = count
+    let postNormalBufferSize = count * 3
+    let resizeRequest = BufferResizeRequest(
+      newOutputBufferSize: outputBufferSize, 
+      newPreNormalScratchBufferSize: preNormalBufferSize,
+      newPostNormalScratchBufferSize: postNormalBufferSize
+    )
+    return .bufferTooSmall(resizeRequest)
+  }
+  
+  static func success(
+    amountFilled filled: Int, nextReadPosition index: String.Index
+  ) -> NormalizationResult {
+    let successResult = SuccessResult(amountFilled: filled, nextReadPosition: index)
+    return .success(successResult)
+  }
+}
+
+func unimplemented() -> Never { fatalError("Unimplemented function called") }
+
+internal func fastFill(
+  _ sourceBuffer: UnsafeBufferPointer<UInt8>,
+  _ outputBuffer: UnsafeMutableBufferPointer<UInt8>
+) -> (Int, Int) {
+  // Quick check if a scalar is NFC and a segment starter
+  @inline(__always) func isNFCStarter(_ scalar: Unicode.Scalar) -> Bool {
+    // Fast-path: All scalars up through U+02FF are NFC and have boundaries
+    // before them
+    if scalar.value < 0x300 { return true }
+
+    // Otherwise, consult the properties
+    return scalar._hasNormalizationBoundaryBefore && scalar._isNFCQCYes
+  }
+  
+  var outputBufferThreshold: Int {
+    return outputBuffer.count - 4
+  }
+
+  // TODO: Additional fast-path: All CCC-ascending NFC_QC segments are NFC
+  // TODO: Just freakin do normalization and don't bother with ICU
+  var outputCount = 0
+  let outputEnd = outputBufferThreshold
+  var inputCount = 0
+  let inputEnd = sourceBuffer.count
+  while inputCount < inputEnd && outputCount < outputEnd {
+    // TODO: Slightly faster code-unit scan for latiny (<0xCC)
+
+    // Check scalar-based fast-paths
+    let (scalar, len) = _decodeScalar(sourceBuffer, startingAt: inputCount)
+    _internalInvariant(inputCount &+ len <= inputEnd)
+
+    if _slowPath(
+         !sourceBuffer.hasNormalizationBoundary(before: inputCount &+ len)
+      || !isNFCStarter(scalar)
+    ) {
+      break
+    }
+    inputCount &+= len
+
+    for cu in UTF8.encode(scalar)._unsafelyUnwrappedUnchecked {
+      outputBuffer[outputCount] = cu
+      outputCount &+= 1
+    }
+
+    _internalInvariant(inputCount == outputCount,
+      "non-normalizing UTF-8 fast path should be 1-to-1 in code units")
+  }
+  return (inputCount, outputCount)
+}
+
+internal func transcodeToUTF16(
+  _ sourceBuffer: UnsafeBufferPointer<UInt8>,
+  into outputBuffer: UnsafeMutableBufferPointer<UInt16>
+) -> (Int, Int)? {
+  var readIndex = 0
+  var writeIndex = 0
+  let outputCount = outputBuffer.count
+  let sourceCount = sourceBuffer.count
+  
+  while readIndex < sourceCount {
+    let (scalar, length) = _decodeScalar(sourceBuffer, startingAt: readIndex)
+    
+    if scalar._hasNormalizationBoundaryBefore && readIndex != 0 {
+      break
+    }
+    
+    readIndex += length
+    
+    for cu in scalar.utf16 {
+      if writeIndex < outputCount {
+        outputBuffer[writeIndex] = cu
+        writeIndex += 1
+      } else {
+        return nil
+      }
+    }
+  }
+  
+  return (readIndex, writeIndex)
+}
+
+internal func transcodeToUTF8(
+  _ sourceBuffer: UnsafeBufferPointer<UInt16>,
+  into outputBuffer: UnsafeMutableBufferPointer<UInt8>
+) -> (Int, Int)? {
+  var readIndex = 0
+  var writeIndex = 0
+  let outputCount = outputBuffer.count
+  let sourceCount = sourceBuffer.count
+  
+  while readIndex < sourceCount {
+    let (scalar, length) = _decodeScalar(sourceBuffer, startingAt: readIndex)
+    //we don't need to check for normalization boundaries here because we are only transcoding
+    //a single segment at this point
+    
+    readIndex += length
+    guard scalar.withUTF8CodeUnits({ utf8 in 
+      for cu in utf8 {
+        if writeIndex < outputCount {
+          outputBuffer[writeIndex] = cu
+          writeIndex += 1
+        } else {
+          return false
+        }
+      }
+      return true
+    }) else {
+      return nil
+    }
+  }
+  return (readIndex, writeIndex)
+}
+
+internal func fastNormalize(
+  readIndex: String.Index,
+  guts: _StringGuts,
+  outputBuffer: UnsafeMutableBufferPointer<UInt8>,
+  icuInputBuffer: UnsafeMutableBufferPointer<UInt16>,
+  icuOutputBuffer: UnsafeMutableBufferPointer<UInt16>
+) -> NormalizationResult {
+  _internalInvariant(guts.isFastUTF8)
+  return guts.withFastUTF8 { sourceBuffer in
+    let sourceCount = sourceBuffer.count - readIndex.encodedOffset
+    let start = readIndex.encodedOffset
+    let rebasedSourceBuffer = UnsafeBufferPointer(rebasing: sourceBuffer[start...])
+    do {
+      let (read, filled) = fastFill(rebasedSourceBuffer, outputBuffer)
+        if filled > 0 {
+          let nextIndex = readIndex.encoded(offsetBy: read)
+          _internalInvariant(guts.isOnUnicodeScalarBoundary(nextIndex))
+
+          return .success(amountFilled: filled, nextReadPosition: nextIndex)
+        }
+    }
+    guard let (read, filled) = transcodeToUTF16(rebasedSourceBuffer, into: icuInputBuffer) else {
+      return .bufferTooSmall(count: sourceBuffer.count)
+    }
+    
+    let nextIndex = readIndex.encoded(offsetBy: read)
+    _internalInvariant(guts.isOnUnicodeScalarBoundary(nextIndex))
+    
+    let rebasedICUInputBuffer = UnsafeBufferPointer(rebasing: icuInputBuffer[..<filled])
+    return sharedNormalize(
+      sourceCount, nextIndex, outputBuffer, rebasedICUInputBuffer, icuOutputBuffer
+    )
+  }
+}
+
+internal func foreignNormalize(
+  readIndex: String.Index,
+  guts: _StringGuts,
+  outputBuffer: UnsafeMutableBufferPointer<UInt8>,
+  icuInputBuffer: UnsafeMutableBufferPointer<UInt16>,
+  icuOutputBuffer: UnsafeMutableBufferPointer<UInt16>
+) -> NormalizationResult {
+  let sourceCount = guts.count - readIndex.encodedOffset
+  guard let (read, filled) = foreignFill(readIndex, guts, into: icuInputBuffer) else {
+    return .bufferTooSmall(count: guts.count)
+  }
+  
+  let nextIndex = readIndex.encoded(offsetBy: read)
+  _internalInvariant(guts.isOnUnicodeScalarBoundary(nextIndex))
+  
+  let rebasedICUInputBuffer = UnsafeBufferPointer(rebasing: icuInputBuffer[..<filled])
+  return sharedNormalize(
+    sourceCount, nextIndex, outputBuffer, rebasedICUInputBuffer, icuOutputBuffer
+  )
+}
+
+func foreignFill(
+  _ readIndex: String.Index,
+  _ guts: _StringGuts,
+  into outputBuffer: UnsafeMutableBufferPointer<UInt16>
+) -> (Int, Int)? {
+  var index = readIndex
+  var writeIndex = 0
+  let outputCount = outputBuffer.count
+  let cachedEndIndex = guts.endIndex
+  while index != cachedEndIndex {
+    let (scalar, length) = guts.foreignErrorCorrectedScalar(startingAt: index)
+    if scalar._hasNormalizationBoundaryBefore && index != readIndex {
+      break
+    }
+    
+    index = index.encoded(offsetBy: length)
+    
+    for cu in scalar.utf16 {
+      if writeIndex < outputCount {
+        outputBuffer[writeIndex] = cu
+        writeIndex += 1
+      } else {
+        return nil
+      }
+    }
+  }
+  return (index.encodedOffset - readIndex.encodedOffset, writeIndex)
+}
+
+private func sharedNormalize(
+  _ sourceCount: Int,
+  _ nextIndex: String.Index,
+  _ outputBuffer: UnsafeMutableBufferPointer<UInt8>,
+  _ icuInputBuffer: UnsafeBufferPointer<UInt16>,
+  _ icuOutputBuffer: UnsafeMutableBufferPointer<UInt16>
+) -> NormalizationResult {
+  guard let normalized = _tryNormalize(icuInputBuffer, into: icuOutputBuffer) else {
+    return .bufferTooSmall(count: sourceCount)
+  }
+  
+  guard let (_, transcoded) = transcodeToUTF8(
+    UnsafeBufferPointer<UInt16>(rebasing: icuOutputBuffer[..<normalized]),
+    into: outputBuffer
+  ) else {
+    return .bufferTooSmall(count: sourceCount)
+  }
+
+  return .success(amountFilled: transcoded, nextReadPosition: nextIndex)
+}

stdlib/public/core/UnicodeHelpers.swift

@@ -23,6 +23,7 @@ internal func _isTrailingSurrogate(_ cu: UInt16) -> Bool {
   return cu & _surrogateMask == _trailingSurrogateBias
 }
 @inline(__always)
+@usableFromInline
 internal func _isLeadingSurrogate(_ cu: UInt16) -> Bool {
   return cu & _surrogateMask == _leadingSurrogateBias
 }
@@ -84,6 +85,23 @@ internal func _decodeUTF8(
   return Unicode.Scalar(_unchecked: value)
 }
 
+@inlinable
+internal func _decodeScalar(
+  _ utf16: UnsafeBufferPointer<UInt16>, startingAt i: Int
+) -> (Unicode.Scalar, scalarLength: Int) {
+  let high = utf16[i]
+  if i + 1 >= utf16.count {
+    return (Unicode.Scalar(_unchecked: UInt32(high)), 1)
+  }
+  
+  if !_isLeadingSurrogate(high) {
+    return (Unicode.Scalar(_unchecked: UInt32(high)), 1)
+  }
+  
+  let low = utf16[i+1]
+  return (Unicode.Scalar(_unchecked: _decodeSurrogatePair(leading: high, trailing: low)), 2)
+}
+
 @inlinable
 internal func _decodeScalar(
   _ utf8: UnsafeBufferPointer<UInt8>, startingAt i: Int
@@ -148,6 +166,7 @@ internal func _continuationPayload(_ x: UInt8) -> UInt32 {
 }
 
 @inline(__always)
+@usableFromInline
 internal func _decodeSurrogatePair(
   leading high: UInt16, trailing low: UInt16
 ) -> UInt32 {

validation-test/stdlib/String.swift

@@ -2216,8 +2216,28 @@ StringTests.test("NormalizationCheck") {
   let str = "\u{0336}\u{0344}\u{0357}\u{0343}\u{0314}\u{0351}\u{0340}\u{0300}\u{0340}\u{0360}\u{0314}\u{0357}\u{0315}\u{0301}\u{0344}a"
   let nfcCodeUnits = str._nfcCodeUnits
   let expectedCodeUnits: [UInt8] = [0xCC, 0xB6, 0xCC, 0x88, 0xCC, 0x81, 0xCD, 0x97, 0xCC, 0x93, 0xCC, 0x94, 0xCD, 0x91, 0xCC, 0x80, 0xCC, 0x80, 0xCC, 0x80, 0xCC, 0x94, 0xCD, 0x97, 0xCC, 0x81, 0xCC, 0x88, 0xCC, 0x81, 0xCC, 0x95, 0xCD, 0xA0, 0x61]
+
+  expectEqual(expectedCodeUnits, nfcCodeUnits)
+}
+
+StringTests.test("NormalizationCheck/Opaque") {
+  let str = "\u{0336}\u{0344}\u{0357}\u{0343}\u{0314}\u{0351}\u{0340}\u{0300}\u{0340}\u{0360}\u{0314}\u{0357}\u{0315}\u{0301}\u{0344}a"
+  let opaqueString = NSSlowString(string: str) as String
+  let nfcCodeUnits = opaqueString._nfcCodeUnits
+  let expectedCodeUnits: [UInt8] = [0xCC, 0xB6, 0xCC, 0x88, 0xCC, 0x81, 0xCD, 0x97, 0xCC, 0x93, 0xCC, 0x94, 0xCD, 0x91, 0xCC, 0x80, 0xCC, 0x80, 0xCC, 0x80, 0xCC, 0x94, 0xCD, 0x97, 0xCC, 0x81, 0xCC, 0x88, 0xCC, 0x81, 0xCC, 0x95, 0xCD, 0xA0, 0x61]
+
+  expectEqual(expectedCodeUnits, nfcCodeUnits)
+}
+
+StringTests.test("ABC") {
+  let str = "abcdefg"
+  let nfcCodeUnits = str._nfcCodeUnits
+  let expectedCodeUnits = Array(str.utf8)
 
   expectEqual(expectedCodeUnits, nfcCodeUnits)
+  
+  let opaqueString = NSSlowString(string: str) as String
+  expectEqual(expectedCodeUnits, opaqueString._nfcCodeUnits)
 }
 
 runAllTests()