Native implementation of -lengthOfBytesUsingEncoding, plus handling ASCII-subset MacRoman in a few places (#81791)

Fixes rdar://154341146

Author: Catfish-Man

benchmark/single-source/ObjectiveCBridging.swift

@@ -100,6 +100,24 @@ public let benchmarks = [
   BenchmarkInfo(name: "NSDictionary.bridged.enumerate",
                   runFunction: run_BridgedNSDictionaryEnumerate, tags: t,
                   setUpFunction: setup_bridgedDictionaries),
+  BenchmarkInfo(name: "NSString.bridged.byteCount.ascii.ascii",
+                  runFunction: run_BridgedNSStringLengthASCII_ASCII, tags: ts,
+                  setUpFunction: setup_bridgedStrings),
+  BenchmarkInfo(name: "NSString.bridged.byteCount.ascii.utf8",
+                  runFunction: run_BridgedNSStringLengthASCII_UTF8, tags: ts,
+                  setUpFunction: setup_bridgedStrings),
+  BenchmarkInfo(name: "NSString.bridged.byteCount.ascii.utf16",
+                  runFunction: run_BridgedNSStringLengthASCII_UTF16, tags: ts,
+                  setUpFunction: setup_bridgedStrings),
+  BenchmarkInfo(name: "NSString.bridged.byteCount.ascii.macroman",
+                  runFunction: run_BridgedNSStringLengthASCII_MacRoman, tags: ts,
+                  setUpFunction: setup_bridgedStrings),
+  BenchmarkInfo(name: "NSString.bridged.byteCount.utf8.utf8",
+                  runFunction: run_BridgedNSStringLengthUTF8_UTF8, tags: ts,
+                  setUpFunction: setup_bridgedStrings),
+  BenchmarkInfo(name: "NSString.bridged.byteCount.utf8.utf16",
+                  runFunction: run_BridgedNSStringLengthUTF8_UTF16, tags: ts,
+                  setUpFunction: setup_bridgedStrings),
 ]
 
 #if _runtime(_ObjC)
@@ -801,6 +819,8 @@ var bridgedDictionaryOfNumbersToNumbers:NSDictionary! = nil
 var bridgedArrayMutableCopy:NSMutableArray! = nil
 var nsArray:NSArray! = nil
 var nsArrayMutableCopy:NSMutableArray! = nil
+var bridgedASCIIString:NSString! = nil
+var bridgedUTF8String:NSString! = nil
 #endif
 
 public func setup_bridgedArrays() {
@@ -822,6 +842,14 @@ public func setup_bridgedDictionaries() {
   bridgedDictionaryOfNumbersToNumbers = numDict as NSDictionary
 }
 
+public func setup_bridgedStrings() {
+  #if _runtime(_ObjC)
+  let str = Array(repeating: "The quick brown fox jumps over the lazy dog.", count: 100).joined()
+  bridgedASCIIString = str as NSString
+  let str2 = Array(repeating: "The quick brown fox jumps over the lazy dög.", count: 100).joined()
+  bridgedUTF8String = str2 as NSString
+  #endif
+}
 
 @inline(never)
 public func run_BridgedNSArrayObjectAtIndex(_ n: Int) {
@@ -914,3 +942,40 @@ public func run_RealNSArrayMutableCopyObjectAtIndex(_ n: Int) {
   #endif
 }
 
+@inline(__always)
+fileprivate func run_BridgedNSStringLength(_ asciiBase: Bool, _ enc: UInt, _ n: Int) {
+  let str = asciiBase ? bridgedASCIIString : bridgedUTF8String
+  for _ in 0 ..< n * 100 {
+    blackHole(str!.lengthOfBytes(using: enc))
+  }
+}
+
+@inline(never)
+public func run_BridgedNSStringLengthASCII_ASCII(_ n: Int) {
+  run_BridgedNSStringLength(true, 1 /* NSASCIIStringEncoding */, n)
+}
+
+@inline(never)
+public func run_BridgedNSStringLengthASCII_UTF8(_ n: Int) {
+  run_BridgedNSStringLength(true, 4 /* NSUTF8StringEncoding */, n)
+}
+
+@inline(never)
+public func run_BridgedNSStringLengthASCII_UTF16(_ n: Int) {
+  run_BridgedNSStringLength(true, 10 /* NSUnicodeStringEncoding */, n)
+}
+
+@inline(never)
+public func run_BridgedNSStringLengthASCII_MacRoman(_ n: Int) {
+  run_BridgedNSStringLength(true, 30 /* NSMacOSRomanStringEncoding */, n)
+}
+
+@inline(never)
+public func run_BridgedNSStringLengthUTF8_UTF8(_ n: Int) {
+  run_BridgedNSStringLength(false, 4 /* NSUTF8StringEncoding */, n)
+}
+
+@inline(never)
+public func run_BridgedNSStringLengthUTF8_UTF16(_ n: Int) {
+  run_BridgedNSStringLength(false, 10 /* NSUnicodeStringEncoding */, n)
+}

stdlib/public/SwiftShims/swift/shims/CoreFoundationShims.h

@@ -35,6 +35,8 @@ typedef unsigned long _swift_shims_CFHashCode;
 typedef signed long _swift_shims_CFIndex;
 #endif
 
+typedef unsigned long _swift_shims_NSUInteger;
+
 // Consider creating SwiftMacTypes.h for these
 typedef unsigned char _swift_shims_Boolean;
 typedef __swift_uint8_t _swift_shims_UInt8;
@@ -80,6 +82,11 @@ const void * _Nullable
 _swift_stdlib_CreateIndirectTaggedPointerString(const __swift_uint8_t * _Nonnull bytes,
                                                 _swift_shims_CFIndex len);
 
+SWIFT_RUNTIME_STDLIB_API
+const _swift_shims_NSUInteger
+_swift_stdlib_NSStringLengthOfBytesInEncodingTrampoline(id _Nonnull obj,
+                                                        unsigned long encoding);
+
 #endif // __OBJC2__
 
 #ifdef __cplusplus

stdlib/public/core/StringBridge.swift

@@ -287,6 +287,13 @@ internal func _cocoaCStringUsingEncodingTrampoline(
   return _swift_stdlib_NSStringCStringUsingEncodingTrampoline(string, encoding)
 }
 
+@_effects(readonly)
+internal func _cocoaLengthOfBytesInEncodingTrampoline(
+  _ string: _CocoaString, _ encoding: UInt
+) -> UInt {
+  return _swift_stdlib_NSStringLengthOfBytesInEncodingTrampoline(string, encoding)
+}
+
 @_effects(releasenone)
 internal func _cocoaGetCStringTrampoline(
   _ string: _CocoaString,

stdlib/public/core/StringStorageBridge.swift

@@ -16,6 +16,8 @@ import SwiftShims
 
 internal var _cocoaASCIIEncoding:UInt { 1 } /* NSASCIIStringEncoding */
 internal var _cocoaUTF8Encoding:UInt { 4 } /* NSUTF8StringEncoding */
+internal var _cocoaUTF16Encoding:UInt { 10 } /* NSUTF16StringEncoding and NSUnicodeStringEncoding*/
+internal var _cocoaMacRomanEncoding:UInt { 30 } /* NSMacOSRomanStringEncoding */
 
 extension String {
   @available(SwiftStdlib 5.6, *)
@@ -69,13 +71,14 @@ extension _AbstractStringStorage {
   ) -> Int8 {
     switch (encoding, isASCII) {
     case (_cocoaASCIIEncoding, true),
+         (_cocoaMacRomanEncoding, true),
          (_cocoaUTF8Encoding, _):
       guard maxLength >= count + 1 else { return 0 }
       unsafe outputPtr.initialize(from: start, count: count)
       unsafe outputPtr[count] = 0
       return 1
     default:
-      return  unsafe _cocoaGetCStringTrampoline(self, outputPtr, maxLength, encoding)
+      return unsafe _cocoaGetCStringTrampoline(self, outputPtr, maxLength, encoding)
     }
   }
 
@@ -84,12 +87,35 @@ extension _AbstractStringStorage {
   internal func _cString(encoding: UInt) -> UnsafePointer<UInt8>? {
     switch (encoding, isASCII) {
     case (_cocoaASCIIEncoding, true),
+         (_cocoaMacRomanEncoding, true),
          (_cocoaUTF8Encoding, _):
       return unsafe start
     default:
       return _cocoaCStringUsingEncodingTrampoline(self, encoding)
     }
   }
+  
+  @_effects(readonly)
+  internal func _lengthOfBytes(using encoding: UInt) -> UInt {
+    switch encoding {
+    case _cocoaASCIIEncoding:
+      if unsafe isASCII || _allASCII(UnsafeBufferPointer(start: start, count: count)) {
+        return UInt(count)
+      }
+      return 0
+    case _cocoaUTF8Encoding:
+      return UInt(count)
+    case _cocoaUTF16Encoding:
+      return UInt(UTF16Length)
+    case _cocoaMacRomanEncoding:
+      if unsafe isASCII || _allASCII(UnsafeBufferPointer(start: start, count: count)) {
+        return UInt(count)
+      }
+      fallthrough
+    default:
+      return _cocoaLengthOfBytesInEncodingTrampoline(self, encoding)
+    }
+  }
 
   @_effects(readonly)
   internal func _nativeIsEqual<T:_AbstractStringStorage>(
@@ -232,6 +258,12 @@ extension __StringStorage {
       return _cocoaUTF8Encoding
     }
   }
+  
+  @objc(lengthOfBytesUsingEncoding:)
+  @_effects(readonly)
+  final internal func lengthOfBytes(using encoding: UInt) -> UInt {
+    _lengthOfBytes(using: encoding)
+  }
 
   @objc(isEqualToString:)
   @_effects(readonly)
@@ -297,6 +329,12 @@ extension __SharedStringStorage {
       return _cocoaUTF8Encoding
     }
   }
+  
+  @objc(lengthOfBytesUsingEncoding:)
+  @_effects(readonly)
+  final internal func lengthOfBytes(using encoding: UInt) -> UInt {
+    _lengthOfBytes(using: encoding)
+  }
 
   @objc(_fastCStringContents:)
   @_effects(readonly)

stdlib/public/stubs/FoundationHelpers.mm

@@ -106,6 +106,19 @@ typedef __swift_uint8_t (*getCStringImplPtr)(id,
 
 }
 
+SWIFT_RUNTIME_STDLIB_API
+const _swift_shims_NSUInteger
+_swift_stdlib_NSStringLengthOfBytesInEncodingTrampoline(id _Nonnull obj,
+                                                        unsigned long encoding) {
+  typedef _swift_shims_NSUInteger (*getLengthImplPtr)(id,
+                                                      SEL,
+                                                      unsigned long);
+  SEL sel = @selector(lengthOfBytesUsingEncoding:);
+  getLengthImplPtr imp = (getLengthImplPtr)class_getMethodImplementation([obj superclass], sel);
+  
+  return imp(obj, sel, encoding);
+}
+
 __swift_uint8_t
 _swift_stdlib_dyld_is_objc_constant_string(const void *addr) {
   return (SWIFT_RUNTIME_WEAK_CHECK(_dyld_is_objc_constant)

test/abi/macOS/arm64/stdlib.swift

@@ -1086,3 +1086,6 @@ Added: _$ss6MirrorV12DisplayStyleO16foreignReferenceyA2DmFWC
 // var InlineArray._protectedAddress
 Added: _$ss11InlineArrayVsRi__rlE16_protectedBufferSRyq_GvpMV
 Added: _$ss11InlineArrayVsRi__rlE17_protectedAddressSPyq_GvpMV
+
+// lengthOfBytes(using:)
+Added: __swift_stdlib_NSStringLengthOfBytesInEncodingTrampoline

test/abi/macOS/x86_64/stdlib.swift

@@ -1086,3 +1086,6 @@ Added: _$ss6MirrorV12DisplayStyleO16foreignReferenceyA2DmFWC
 // var InlineArray._protectedAddress
 Added: _$ss11InlineArrayVsRi__rlE16_protectedBufferSRyq_GvpMV
 Added: _$ss11InlineArrayVsRi__rlE17_protectedAddressSPyq_GvpMV
+
+// lengthOfBytes(using:)
+Added: __swift_stdlib_NSStringLengthOfBytesInEncodingTrampoline