Separate out ASCII and UTF8 in NSString accessors, and use the ASCII one where we really need ASCII

Author: Catfish-Man

stdlib/public/core/SmallString.swift

@@ -372,6 +372,28 @@ extension _SmallString {
     }
     self._invariantCheck()
   }
+  
+  @_effects(readonly) // @opaque
+  @usableFromInline // testable
+  internal init?(taggedASCIICocoa cocoa: AnyObject) {
+    self.init()
+    var success = true
+    self.withMutableCapacity {
+      /*
+       For regular NSTaggedPointerStrings we will always succeed here, but
+       tagged NSLocalizedStrings may not fit in a SmallString
+       */
+      if let len = _bridgeTaggedASCII(cocoa, intoUTF8: $0) {
+        return len
+      }
+      success = false
+      return 0
+    }
+    if !success {
+      return nil
+    }
+    self._invariantCheck()
+  }
 }
 #endif
 

stdlib/public/core/StringBridge.swift

@@ -167,8 +167,9 @@ internal func _cocoaStringSubscript(
 }
 
 @_effects(releasenone)
-private func _NSStringCopyUTF8(
+private func _NSStringCopyBytes(
   _ o: _StringSelectorHolder,
+  encoding: UInt,
   into bufPtr: UnsafeMutableRawBufferPointer
 ) -> Int? {
   let ptr = bufPtr.baseAddress._unsafelyUnwrappedUnchecked
@@ -179,7 +180,7 @@ private func _NSStringCopyUTF8(
     ptr,
     maxLength: bufPtr.count,
     usedLength: &usedLen,
-    encoding: _cocoaUTF8Encoding,
+    encoding: encoding,
     options: 0,
     range: _SwiftNSRange(location: 0, length: len),
     remaining: &remainingRange
@@ -195,7 +196,23 @@ internal func _cocoaStringCopyUTF8(
   _ target: _CocoaString,
   into bufPtr: UnsafeMutableRawBufferPointer
 ) -> Int? {
-  return _NSStringCopyUTF8(_objc(target), into: bufPtr)
+  return _NSStringCopyUTF8(
+    _objc(target),
+    encoding: _cocoaUTF8Encoding,
+    into: bufPtr
+  )
+}
+
+@_effects(releasenone)
+internal func _cocoaStringCopyASCII(
+  _ target: _CocoaString,
+  into bufPtr: UnsafeMutableRawBufferPointer
+) -> Int? {
+  return _NSStringCopyASCII(
+    _objc(target),
+    encoding: _cocoaASCIIEncoding,
+    into: bufPtr
+  )
 }
 
 @_effects(readonly)
@@ -346,17 +363,30 @@ internal func _bridgeTagged(
   _internalInvariant(_isObjCTaggedPointer(cocoa))
   return _cocoaStringCopyUTF8(cocoa, into: bufPtr)
 }
+
+@_effects(releasenone) // @opaque
+internal func _bridgeTaggedASCII(
+  _ cocoa: _CocoaString,
+  intoUTF8 bufPtr: UnsafeMutableRawBufferPointer
+) -> Int? {
+  _internalInvariant(_isObjCTaggedPointer(cocoa))
+  return _cocoaStringCopyASCII(cocoa, into: bufPtr)
+}
 #endif
 
 @_effects(readonly)
 private func _NSStringASCIIPointer(_ str: _StringSelectorHolder) -> UnsafePointer<UInt8>? {
- // TODO(String bridging): Is there a better interface here? Ideally we'd be
-  // able to ask for UTF8 rather than just ASCII
   //TODO(String bridging): Unconditionally asking for nul-terminated contents is
   // overly conservative and hurts perf with some NSStrings
   return str._fastCStringContents(1)?._asUInt8
 }
 
+@_effects(readonly)
+private func _NSStringUTF8Pointer(_ str: _StringSelectorHolder) -> UnsafePointer<UInt8>? {
+  //We don't have a way to ask for UTF8 here currently
+  return _NSStringASCIIPointer(str)
+}
+
 @_effects(readonly) // @opaque
 private func _withCocoaASCIIPointer<R>(
   _ str: _CocoaString,
@@ -371,7 +401,7 @@ private func _withCocoaASCIIPointer<R>(
     if requireStableAddress {
       return nil // tagged pointer strings don't support _fastCStringContents
     }
-    if let smol = _SmallString(taggedCocoa: str) {
+    if let smol = _SmallString(taggedASCIICocoa: str) {
       return _StringGuts(smol).withFastUTF8 {
         work($0.baseAddress._unsafelyUnwrappedUnchecked)
       }
@@ -385,6 +415,34 @@ private func _withCocoaASCIIPointer<R>(
   return nil
 }
 
+@_effects(readonly) // @opaque
+private func _withCocoaUTF8Pointer<R>(
+  _ str: _CocoaString,
+  requireStableAddress: Bool,
+  work: (UnsafePointer<UInt8>) -> R?
+) -> R? {
+  #if !(arch(i386) || arch(arm) || arch(arm64_32))
+  if _isObjCTaggedPointer(str) {
+    if let ptr = getConstantTaggedCocoaContents(str)?.asciiContentsPointer {
+      return work(ptr)
+    }
+    if requireStableAddress {
+      return nil // tagged pointer strings don't support _fastCStringContents
+    }
+    if let smol = _SmallString(taggedCocoa: str) {
+      return _StringGuts(smol).withFastUTF8 {
+        work($0.baseAddress._unsafelyUnwrappedUnchecked)
+      }
+    }
+  }
+  #endif
+  defer { _fixLifetime(str) }
+  if let ptr = _NSStringUTF8Pointer(_objc(str)) {
+    return work(ptr)
+  }
+  return nil
+}
+
 @_effects(readonly) // @opaque
 internal func withCocoaASCIIPointer<R>(
   _ str: _CocoaString,
@@ -393,12 +451,26 @@ internal func withCocoaASCIIPointer<R>(
   return _withCocoaASCIIPointer(str, requireStableAddress: false, work: work)
 }
 
+@_effects(readonly) // @opaque
+internal func withCocoaUTF8Pointer<R>(
+  _ str: _CocoaString,
+  work: (UnsafePointer<UInt8>) -> R?
+) -> R? {
+  return _withCocoaUTF8Pointer(str, requireStableAddress: false, work: work)
+}
+
 @_effects(readonly)
 internal func stableCocoaASCIIPointer(_ str: _CocoaString)
   -> UnsafePointer<UInt8>? {
   return _withCocoaASCIIPointer(str, requireStableAddress: true, work: { $0 })
 }
 
+@_effects(readonly)
+internal func stableCocoaUTF8Pointer(_ str: _CocoaString)
+  -> UnsafePointer<UInt8>? {
+  return _withCocoaUTF8Pointer(str, requireStableAddress: true, work: { $0 })
+}
+
 private enum CocoaStringPointer {
   case ascii(UnsafePointer<UInt8>)
   case utf8(UnsafePointer<UInt8>)

stdlib/public/core/StringObject.swift

@@ -915,7 +915,7 @@ extension _StringObject {
     _internalInvariant(largeFastIsShared)
 #if _runtime(_ObjC)
     if largeIsCocoa {
-      return stableCocoaASCIIPointer(cocoaObject)._unsafelyUnwrappedUnchecked
+      return stableCocoaUTF8Pointer(cocoaObject)._unsafelyUnwrappedUnchecked
     }
 #endif
 

stdlib/public/core/StringStorageBridge.swift

@@ -14,8 +14,8 @@ import SwiftShims
 
 #if _runtime(_ObjC)
 
-internal let _cocoaASCIIEncoding:UInt = 1 /* NSASCIIStringEncoding */
-internal let _cocoaUTF8Encoding:UInt = 4 /* NSUTF8StringEncoding */
+internal let _cocoaASCIIEncoding:UInt { 1 } /* NSASCIIStringEncoding */
+internal let _cocoaUTF8Encoding:UInt { 4 } /* NSUTF8StringEncoding */
 
 extension String {
   @available(SwiftStdlib 5.6, *)