[DNM] Experimentally change how Dictionary casting handles Strings, and speed up String bridging a bit to compensate

benchmark/single-source/ObjectiveCBridging.swift

@@ -311,7 +311,6 @@ func createNSDictionary() -> NSDictionary {
 func testObjectiveCBridgeFromNSDictionaryAnyObject() {
   let nsDictionary = createNSDictionary()
   let nsString = NSString(cString: "NSString that does not fit in tagged pointer", encoding: String.Encoding.utf8.rawValue)!
-
   var nativeInt : Int?
   for _ in 0 ..< 10_000 {
     if let nativeDictionary : [NSString : NSNumber] = conditionalCast(nsDictionary) {

lib/IRGen/GenEnum.cpp

@@ -175,6 +175,12 @@ EnumImplStrategy::EnumImplStrategy(IRGenModule &IGM,
     NumElements(NumElements) {
 }
 
+StackAddress EnumImplStrategy::allocateStack(IRGenFunction &IGF, SILType T,
+                                             const llvm::Twine &name) const {
+  // Allocate memory on the stack, loading the size from the type metadata for T.
+  return IGF.emitDynamicAlloca(T, name);
+}
+
 void EnumImplStrategy::initializeFromParams(IRGenFunction &IGF,
                                             Explosion &params,
                                             Address dest, SILType T,
@@ -2322,6 +2328,19 @@ namespace {
     }
 
   public:
+    StackAddress allocateStack(IRGenFunction &IGF, SILType T,
+                               const llvm::Twine &name) const override {
+      assert(TIK < Fixed);
+
+      if (ElementsWithNoPayload.size() > 1)
+        return EnumImplStrategy::allocateStack(IGF, T, name);
+
+      llvm::Value *size = emitLoadOfSize(IGF, getPayloadType(IGM, T));
+      size = IGF.Builder.CreateAdd(size, llvm::ConstantInt::get(IGM.SizeTy, 1));
+
+      return IGF.emitDynamicAlloca(IGM.Int8Ty, size, Alignment(16), name);
+    }
+
     void copy(IRGenFunction &IGF, Explosion &src, Explosion &dest,
               Atomicity atomicity) const override {
       assert(TIK >= Loadable);
@@ -2796,10 +2815,17 @@ namespace {
 
     void collectMetadataForOutlining(OutliningMetadataCollector &collector,
                                      SILType T) const override {
+      if (T.getOptionalObjectType()) {
+        auto payloadT = getPayloadType(IGM, T);
+        getPayloadTypeInfo().collectMetadataForOutlining(collector, payloadT);
+        return;
+      }
+
       if (CopyDestroyKind == Normal) {
         auto payloadT = getPayloadType(IGM, T);
         getPayloadTypeInfo().collectMetadataForOutlining(collector, payloadT);
       }
+
       collector.collectTypeMetadataForLayout(T);
     }
 
@@ -5804,6 +5830,15 @@ namespace {
                          IsBitwiseTakable_t bt,
                          IsABIAccessible_t abiAccessible)
       : EnumTypeInfoBase(strategy, irTy, align, pod, bt, abiAccessible) {}
+
+    StackAddress allocateStack(IRGenFunction &IGF, SILType t,
+                               const llvm::Twine &name) const override {
+      auto alloca = Strategy.allocateStack(IGF, t, name);
+
+      IGF.Builder.CreateLifetimeStart(alloca.getAddressPointer());
+      return alloca.withAddress(
+          getAsBitCastAddress(IGF, alloca.getAddressPointer()));
+    }
   };
 
   /// TypeInfo for dynamically-sized enum types.

lib/IRGen/GenEnum.h

@@ -237,7 +237,12 @@ class EnumImplStrategy {
   getBitMaskForNoPayloadElements() const = 0;
 
   /// \group Indirect enum operations
-
+
+  /// Customization point for allocating the enum value on the stack. Only used
+  /// for non-fixed-size enums.
+  virtual StackAddress allocateStack(IRGenFunction &IGF, SILType T,
+                                     const llvm::Twine &name) const;
+
   /// Return the enum case tag for the given value. Payload cases come first,
   /// followed by non-payload cases. Used for the getEnumTag value witness.
   virtual llvm::Value *emitGetEnumTag(IRGenFunction &IGF,

lib/IRGen/NonFixedTypeInfo.h

@@ -56,7 +56,6 @@ class WitnessSizedTypeInfo : public IndirectTypeInfo<Impl, TypeInfo> {
     : super(type, align, pod, bt, IsNotFixedSize, abi,
             SpecialTypeInfoKind::None) {}
 
-private:
   /// Bit-cast the given pointer to the right type and assume it as an
   /// address of this type.
   Address getAsBitCastAddress(IRGenFunction &IGF, llvm::Value *addr) const {

stdlib/public/Darwin/Foundation/NSDictionary.swift

@@ -13,6 +13,17 @@
 @_exported import Foundation // Clang module
 import _SwiftFoundationOverlayShims
 
+@objc internal protocol _NSDictionaryEnumerationHack {
+  @objc(enumerateKeysAndObjectsWithOptions:usingBlock:)
+  func enumerateKeysAndObjects(
+    options: Int,
+    using block: @convention(block) (
+    Unmanaged<AnyObject>,
+    Unmanaged<AnyObject>,
+    UnsafeMutablePointer<UInt8>
+    ) -> Void)
+}
+
 //===----------------------------------------------------------------------===//
 // Dictionaries
 //===----------------------------------------------------------------------===//
@@ -62,64 +73,98 @@ extension Dictionary : _ObjectiveCBridgeable {
   }
 
   public static func _forceBridgeFromObjectiveC(
-    _ d: NSDictionary,
+    _ x: NSDictionary,
     result: inout Dictionary?
   ) {
     if let native = [Key : Value]._bridgeFromObjectiveCAdoptingNativeStorageOf(
-        d as AnyObject) {
+        x as AnyObject) {
       result = native
       return
     }
 
     if _isBridgedVerbatimToObjectiveC(Key.self) &&
        _isBridgedVerbatimToObjectiveC(Value.self) {
-      result = [Key : Value](_cocoaDictionary: d)
-      return
-    }
-
-    if Key.self == String.self {
-      // String and NSString have different concepts of equality, so
-      // string-keyed NSDictionaries may generate key collisions when bridged
-      // over to Swift. See rdar://problem/35995647
-      var dict = Dictionary(minimumCapacity: d.count)
-      d.enumerateKeysAndObjects({ (anyKey: Any, anyValue: Any, _) in
-        let key = Swift._forceBridgeFromObjectiveC(
-          anyKey as AnyObject, Key.self)
-        let value = Swift._forceBridgeFromObjectiveC(
-          anyValue as AnyObject, Value.self)
-        // FIXME: Log a warning if `dict` already had a value for `key`
-        dict[key] = value
-      })
-      result = dict
+      result = [Key : Value](_cocoaDictionary: x)
       return
     }
 
     // `Dictionary<Key, Value>` where either `Key` or `Value` is a value type
     // may not be backed by an NSDictionary.
-    var builder = _DictionaryBuilder<Key, Value>(count: d.count)
-    d.enumerateKeysAndObjects({ (anyKey: Any, anyValue: Any, _) in
-      let anyObjectKey = anyKey as AnyObject
-      let anyObjectValue = anyValue as AnyObject
-      builder.add(
-          key: Swift._forceBridgeFromObjectiveC(anyObjectKey, Key.self),
-          value: Swift._forceBridgeFromObjectiveC(anyObjectValue, Value.self))
-    })
+    let fastSelf = unsafeBitCast(x, to: _NSDictionaryEnumerationHack.self)
+    var builder = _DictionaryBuilder<Key, Value>(count: x.count)
+    fastSelf.enumerateKeysAndObjects(options: 0) { (anyKey, anyValue, stopPtr) in
+      anyKey._withUnsafeGuaranteedRef { (anyKey) in
+        anyValue._withUnsafeGuaranteedRef { (anyValue) in
+          let key = Swift._forceBridgeFromObjectiveC(
+                  anyKey, Key.self)
+          let value = Swift._forceBridgeFromObjectiveC(
+                  anyValue, Value.self)
+          // String and NSString have different concepts of equality, so
+          // string-keyed NSDictionaries may generate key collisions when bridged
+          // over to Swift. See rdar://problem/35995647
+          if Key.self == String.self &&
+                !(unsafeBitCast(key, to: String.self)._isKnownASCII) {
+            // FIXME: Log a warning if `dict` already had a value for `key`
+              builder.add(
+                possibleDuplicateKey: key,
+                value: value)
+          } else {
+            builder.add(
+                key: key,
+                value: value)
+          }
+        }
+      }
+    }
     result = builder.take()
   }
 
   public static func _conditionallyBridgeFromObjectiveC(
     _ x: NSDictionary,
     result: inout Dictionary?
-  ) -> Bool {
-    let anyDict = x as [NSObject : AnyObject]
-    if _isBridgedVerbatimToObjectiveC(Key.self) &&
-       _isBridgedVerbatimToObjectiveC(Value.self) {
-      result = Swift._dictionaryDownCastConditional(anyDict)
-      return result != nil
+    ) -> Bool {
+
+    if let native = [Key : Value]._bridgeFromObjectiveCAdoptingNativeStorageOf(
+      x as AnyObject) {
+      result = native
+      return true
     }
-
-    result = anyDict as? Dictionary
-    return result != nil
+
+    var success = true
+    // `Dictionary<Key, Value>` where either `Key` or `Value` is a value type
+    // may not be backed by an NSDictionary.
+    let fastSelf = unsafeBitCast(x, to: _NSDictionaryEnumerationHack.self)
+    var builder = _DictionaryBuilder<Key, Value>(count: x.count)
+    fastSelf.enumerateKeysAndObjects(options: 0) { (anyKey, anyValue, stopPtr) in
+      anyKey._withUnsafeGuaranteedRef { (anyKey) in
+        anyValue._withUnsafeGuaranteedRef { (anyValue) in
+          guard let key = anyKey as? Key,
+            let value = anyValue as? Value else {
+                stopPtr.pointee = 1
+                success = false
+                return
+          }
+          // String and NSString have different concepts of equality, so
+          // string-keyed NSDictionaries may generate key collisions when bridged
+          // over to Swift. See rdar://problem/35995647
+          if Key.self == String.self &&
+            !(unsafeBitCast(key, to: String.self)._isKnownASCII) {
+            // FIXME: Log a warning if `dict` already had a value for `key`
+            builder.add(
+              possibleDuplicateKey: key,
+              value: value)
+          } else {
+            builder.add(
+              key: key,
+              value: value)
+          }
+        }
+      }
+    }
+    if success {
+      result = builder.take()
+    }
+    return success
   }
 
   @_effects(readonly)

stdlib/public/SwiftShims/CoreFoundationShims.h

@@ -45,7 +45,6 @@ typedef struct {
 } _swift_shims_CFRange;
 
 typedef id _swift_shims_CFStringRef;
-typedef const struct _swift_shims_CFAllocator *_swift_shims_CFAllocatorRef;
 typedef __swift_uint32_t _swift_shims_CFStringEncoding;
 typedef _swift_shims_CFOptionFlags _swift_shims_CFStringCompareFlags;
 typedef _swift_shims_CFIndex _swift_shims_CFComparisonResult;
@@ -80,7 +79,6 @@ _swift_shims_CFIndex _swift_stdlib_CFStringGetLength(
 SWIFT_RUNTIME_STDLIB_API
 __attribute__((ns_returns_retained))
 _swift_shims_CFStringRef _Nonnull _swift_stdlib_CFStringCreateWithSubstring(
-    _swift_shims_CFAllocatorRef _Nullable alloc,
     _swift_shims_CFStringRef _Nonnull str, _swift_shims_CFRange range);
 
 SWIFT_RUNTIME_STDLIB_API
@@ -90,13 +88,11 @@ _swift_shims_UniChar _swift_stdlib_CFStringGetCharacterAtIndex(
 SWIFT_RUNTIME_STDLIB_API
 __attribute__((ns_returns_retained))
 _swift_shims_CFStringRef _Nonnull _swift_stdlib_CFStringCreateCopy(
-    _swift_shims_CFAllocatorRef _Nullable alloc,
     _swift_shims_CFStringRef _Nonnull theString);
 
 SWIFT_RUNTIME_STDLIB_API
 __attribute__((ns_returns_retained))
 _swift_shims_CFStringRef _Nonnull _swift_stdlib_CFStringCreateWithBytes(
-    _swift_shims_CFAllocatorRef _Nullable alloc,
     const __swift_uint8_t *_Nonnull bytes, _swift_shims_CFIndex numBytes,
     _swift_shims_CFStringEncoding encoding,
     _swift_shims_Boolean isExternalRepresentation);

stdlib/public/core/DictionaryBridging.swift

@@ -365,7 +365,7 @@ final internal class _SwiftDeferredNSDictionary<Key: Hashable, Value>
       if stop != 0 { return }
     }
   }
-
+  
   @objc
   internal var count: Int {
     return native.count

stdlib/public/core/DictionaryBuilder.swift

@@ -27,6 +27,13 @@ struct _DictionaryBuilder<Key: Hashable, Value> {
     _target = _NativeDictionary(capacity: count)
     _requestedCount = count
   }
+
+  @inlinable
+  public mutating func add(possibleDuplicateKey newKey: Key, value: Value) {
+    _precondition(_target.count < _requestedCount,
+      "Can't add more members than promised")
+    _target.setValue(value, forKey: newKey, isUnique: true)
+  }
 
   @inlinable
   public mutating func add(key newKey: Key, value: Value) {

stdlib/public/core/String.swift

@@ -842,3 +842,11 @@ extension String {
     }
   }
 }
+
+extension String {
+  //SPI: NSDictionary bridging
+  @inlinable public
+  var _isKnownASCII:Bool {
+    return _guts.isASCII
+  }
+}

stdlib/public/core/StringBridge.swift

@@ -26,7 +26,7 @@ internal typealias _CocoaString = AnyObject
 internal func _stdlib_binary_CFStringCreateCopy(
   _ source: _CocoaString
 ) -> _CocoaString {
-  let result = _swift_stdlib_CFStringCreateCopy(nil, source) as AnyObject
+  let result = _swift_stdlib_CFStringCreateCopy(source) as AnyObject
   return result
 }
 
@@ -274,7 +274,7 @@ extension String {
       return _guts.asSmall.withUTF8 { bufPtr in
         // TODO(String bridging): worth isASCII check for different encoding?
         return _swift_stdlib_CFStringCreateWithBytes(
-            nil, bufPtr.baseAddress._unsafelyUnwrappedUnchecked,
+            bufPtr.baseAddress._unsafelyUnwrappedUnchecked,
             bufPtr.count,
             kCFStringEncodingUTF8, 0)
         as AnyObject

stdlib/public/stubs/FoundationHelpers.mm

@@ -30,7 +30,6 @@
 #define BRIDGE_TYPE(FROM, TO) \
 template <> struct DestType<FROM> { using type = TO; }
 
-BRIDGE_TYPE(_swift_shims_CFAllocatorRef, CFAllocatorRef);
 BRIDGE_TYPE(_swift_shims_CFStringRef, CFStringRef);
 BRIDGE_TYPE(_swift_shims_UniChar *, UniChar *);
 BRIDGE_TYPE(_swift_shims_CFStringEncoding, CFStringEncoding);
@@ -75,12 +74,13 @@ static CFRange cast(_swift_shims_CFRange value) {
   return CFStringGetLength(cast(theString));
 }
 
+extern "C" CFAllocatorRef const kCFAllocatorSystemDefault;
+
 _swift_shims_CFStringRef
 swift::_swift_stdlib_CFStringCreateWithSubstring(
-                                         _swift_shims_CFAllocatorRef alloc,
                                          _swift_shims_CFStringRef str,
                                          _swift_shims_CFRange range) {
-  return cast(CFStringCreateWithSubstring(cast(alloc), cast(str), cast(range)));
+  return cast(CFStringCreateWithSubstring(kCFAllocatorSystemDefault, cast(str), cast(range)));
 }
 
 _swift_shims_CFComparisonResult
@@ -108,17 +108,15 @@ static CFRange cast(_swift_shims_CFRange value) {
 }
 
 _swift_shims_CFStringRef
-swift::_swift_stdlib_CFStringCreateCopy(_swift_shims_CFAllocatorRef alloc,
-                                        _swift_shims_CFStringRef theString) {
-  return cast(CFStringCreateCopy(cast(alloc), cast(theString)));
+swift::_swift_stdlib_CFStringCreateCopy(_swift_shims_CFStringRef theString) {
+  return cast(CFStringCreateCopy(kCFAllocatorSystemDefault, cast(theString)));
 }
 
 _swift_shims_CFStringRef
-swift::_swift_stdlib_CFStringCreateWithBytes(
-    _swift_shims_CFAllocatorRef _Nullable alloc, const uint8_t *bytes,
+swift::_swift_stdlib_CFStringCreateWithBytes(const uint8_t *bytes,
     _swift_shims_CFIndex numBytes, _swift_shims_CFStringEncoding encoding,
     _swift_shims_Boolean isExternalRepresentation) {
-  return cast(CFStringCreateWithBytes(cast(alloc), bytes, numBytes,
+  return cast(CFStringCreateWithBytes(kCFAllocatorSystemDefault, bytes, numBytes,
                                       cast(encoding),
                                       isExternalRepresentation));
 }
@@ -137,7 +135,6 @@ static CFRange cast(_swift_shims_CFRange value) {
 extern "C" CFHashCode CFStringHashCString(const uint8_t *bytes, CFIndex len);
 extern "C" CFHashCode CFStringHashNSString(id str);
 
-
 _swift_shims_CFHashCode
 swift::_swift_stdlib_CFStringHashNSString(id _Nonnull obj) {
   return CFStringHashNSString(obj);