[stdlib] Fix Set/Dictionary casting issues

stdlib/public/core/DictionaryCasting.swift

@@ -12,6 +12,30 @@
 
 //===--- Compiler conversion/casting entry points for Dictionary<K, V> ----===//
 
+extension Dictionary {
+  @_alwaysEmitIntoClient @inlinable // Introduced in 5.1
+  @inline(__always)
+  internal init?<C: Collection>(
+    _mapping source: C,
+    allowingDuplicates: Bool,
+    transform: (C.Element) -> (key: Key, value: Value)?
+  ) {
+    var target = _NativeDictionary<Key, Value>(capacity: source.count)
+    if allowingDuplicates {
+      for member in source {
+        guard let (key, value) = transform(member) else { return nil }
+        target._unsafeUpdate(key: key, value: value)
+      }
+    } else {
+      for member in source {
+        guard let (key, value) = transform(member) else { return nil }
+        target._unsafeInsertNew(key: key, value: value)
+      }
+    }
+    self.init(_native: target)
+  }
+}
+
 /// Perform a non-bridged upcast that always succeeds.
 ///
 /// - Precondition: `BaseKey` and `BaseValue` are base classes or base `@objc`
@@ -21,12 +45,15 @@
 public func _dictionaryUpCast<DerivedKey, DerivedValue, BaseKey, BaseValue>(
     _ source: Dictionary<DerivedKey, DerivedValue>
 ) -> Dictionary<BaseKey, BaseValue> {
-  var builder = _DictionaryBuilder<BaseKey, BaseValue>(count: source.count)
-
-  for (k, v) in source {
-    builder.add(key:k as! BaseKey, value: v as! BaseValue)
-  }
-  return builder.take()
+  return Dictionary(
+    _mapping: source,
+    // String and NSString have different concepts of equality, so
+    // NSString-keyed Dictionaries may generate key collisions when "upcasted"
+    // to String. See rdar://problem/35995647
+    allowingDuplicates: (BaseKey.self == String.self)
+  ) { k, v in
+    (k as! BaseKey, v as! BaseValue)
+  }!
 }
 
 /// Called by the casting machinery.
@@ -66,7 +93,20 @@ public func _dictionaryDownCast<BaseKey, BaseValue, DerivedKey, DerivedValue>(
       _immutableCocoaDictionary: source._variant.asNative.bridged())
   }
 #endif
-  return _dictionaryDownCastConditional(source)!
+
+  // Note: We can't delegate this call to _dictionaryDownCastConditional,
+  // because we rely on as! to generate nice runtime errors when the downcast
+  // fails.
+
+  return Dictionary(
+    _mapping: source,
+    // String and NSString have different concepts of equality, so
+    // NSString-keyed Dictionaries may generate key collisions when downcasted
+    // to String. See rdar://problem/35995647
+    allowingDuplicates: (DerivedKey.self == String.self)
+  ) { k, v in
+    (k as! DerivedKey, v as! DerivedValue)
+  }!
 }
 
 /// Called by the casting machinery.
@@ -97,12 +137,19 @@ public func _dictionaryDownCastConditional<
 >(
   _ source: Dictionary<BaseKey, BaseValue>
 ) -> Dictionary<DerivedKey, DerivedValue>? {
-
-  var builder = _DictionaryBuilder<DerivedKey, DerivedValue>(count: source.count)
-  for (k, v) in source {
-    guard let k1 = k as? DerivedKey, let v1 = v as? DerivedValue
-    else { return nil }
-    builder.add(key: k1, value: v1)
+  return Dictionary(
+    _mapping: source,
+    // String and NSString have different concepts of equality, so
+    // NSString-keyed Dictionaries may generate key collisions when downcasted
+    // to String. See rdar://problem/35995647
+    allowingDuplicates: (DerivedKey.self == String.self)
+  ) { k, v in
+    guard
+      let key = k as? DerivedKey,
+      let value = v as? DerivedValue
+    else {
+      return nil
+    }
+    return (key, value)
   }
-  return builder.take()
 }

stdlib/public/core/NativeDictionary.swift

@@ -485,6 +485,28 @@ extension _NativeDictionary { // Insertions
     _storage._count &+= 1
   }
 
+  /// Insert a new element into uniquely held storage, replacing an existing
+  /// value (if any).  Storage must be uniquely referenced with adequate
+  /// capacity.
+  @_alwaysEmitIntoClient @inlinable // Introduced in 5.1
+  internal mutating func _unsafeUpdate(
+    key: __owned Key,
+    value: __owned Value
+  ) {
+    let (bucket, found) = find(key)
+    if found {
+      // Note that we also update the key here. This method is used to handle
+      // collisions arising from equality transitions during bridging, and in
+      // that case it is desirable to keep values paired with their original
+      // keys. This is not how `updateValue(_:, forKey:)` works.
+      (_keys + bucket.offset).pointee = key
+      (_values + bucket.offset).pointee = value
+    } else {
+      _precondition(count < capacity)
+      _insert(at: bucket, key: key, value: value)
+    }
+  }
+
   /// Insert a new entry into uniquely held storage.
   /// Storage must be uniquely referenced.
   /// The `key` must not be already present in the Dictionary.

stdlib/public/core/NativeSet.swift

@@ -117,10 +117,21 @@ extension _NativeSet { // Low-level unchecked operations
   @inline(__always)
   internal func uncheckedInitialize(
     at bucket: Bucket,
-    to element: __owned Element) {
+    to element: __owned Element
+  ) {
     _internalInvariant(hashTable.isValid(bucket))
     (_elements + bucket.offset).initialize(to: element)
   }
+
+  @_alwaysEmitIntoClient @inlinable // Introduced in 5.1
+  @inline(__always)
+  internal func uncheckedAssign(
+    at bucket: Bucket,
+    to element: __owned Element
+  ) {
+    _internalInvariant(hashTable.isOccupied(bucket))
+    (_elements + bucket.offset).pointee = element
+  }
 }
 
 extension _NativeSet { // Low-level lookup operations
@@ -430,6 +441,21 @@ extension _NativeSet { // Insertions
     _unsafeInsertNew(element, at: bucket)
     return nil
   }
+
+  /// Insert an element into uniquely held storage, replacing an existing value
+  /// (if any).  Storage must be uniquely referenced with adequate capacity.
+  @_alwaysEmitIntoClient @inlinable // Introduced in 5.1
+  internal mutating func _unsafeUpdate(
+    with element: __owned Element
+  ) {
+    let (bucket, found) = find(element)
+    if found {
+      uncheckedAssign(at: bucket, to: element)
+    } else {
+      _precondition(count < capacity)
+      _unsafeInsertNew(element, at: bucket)
+    }
+  }
 }
 
 extension _NativeSet {

stdlib/public/core/SetCasting.swift

@@ -12,18 +12,47 @@
 
 //===--- Compiler conversion/casting entry points for Set<Element> --------===//
 
+extension Set {
+  @_alwaysEmitIntoClient @inlinable // Introduced in 5.1
+  @inline(__always)
+  internal init?<C: Collection>(
+    _mapping source: C,
+    allowingDuplicates: Bool,
+    transform: (C.Element) -> Element?
+  ) {
+    var target = _NativeSet<Element>(capacity: source.count)
+    if allowingDuplicates {
+      for m in source {
+        guard let member = transform(m) else { return nil }
+        target._unsafeUpdate(with: member)
+      }
+    } else {
+      for m in source {
+        guard let member = transform(m) else { return nil }
+        target._unsafeInsertNew(member)
+      }
+    }
+    self.init(_native: target)
+  }
+}
+
 /// Perform a non-bridged upcast that always succeeds.
 ///
 /// - Precondition: `BaseValue` is a base class or base `@objc`
 ///   protocol (such as `AnyObject`) of `DerivedValue`.
 @inlinable
-public func _setUpCast<DerivedValue, BaseValue>(_ source: Set<DerivedValue>)
-  -> Set<BaseValue> {
-  var builder = _SetBuilder<BaseValue>(count: source.count)
-  for x in source {
-    builder.add(member: x as! BaseValue)
-  }
-  return builder.take()
+public func _setUpCast<DerivedValue, BaseValue>(
+  _ source: Set<DerivedValue>
+) -> Set<BaseValue> {
+  return Set(
+    _mapping: source,
+    // String and NSString have different concepts of equality, so Set<NSString>
+    // may generate key collisions when "upcasted" to Set<String>.
+    // See rdar://problem/35995647
+    allowingDuplicates: (BaseValue.self == String.self)
+  ) { member in
+    (member as! BaseValue)
+  }!
 }
 
 /// Called by the casting machinery.
@@ -54,7 +83,18 @@ public func _setDownCast<BaseValue, DerivedValue>(_ source: Set<BaseValue>)
     return Set(_immutableCocoaSet: source._variant.asNative.bridged())
   }
 #endif
-  return _setDownCastConditional(source)!
+  // We can't just delegate to _setDownCastConditional here because we rely on
+  // `as!` to generate nice runtime errors when the downcast fails.
+
+  return Set(
+    _mapping: source,
+    // String and NSString have different concepts of equality, so
+    // NSString-keyed Sets may generate key collisions when downcasted
+    // to String. See rdar://problem/35995647
+    allowingDuplicates: (DerivedValue.self == String.self)
+  ) { member in
+    (member as! DerivedValue)
+  }!
 }
 
 /// Called by the casting machinery.
@@ -81,13 +121,13 @@ internal func _setDownCastConditionalIndirect<SourceValue, TargetValue>(
 public func _setDownCastConditional<BaseValue, DerivedValue>(
   _ source: Set<BaseValue>
 ) -> Set<DerivedValue>? {
-  var result = Set<DerivedValue>(minimumCapacity: source.count)
-  for member in source {
-    if let derivedMember = member as? DerivedValue {
-      result.insert(derivedMember)
-      continue
-    }
-    return nil
+  return Set(
+    _mapping: source,
+    // String and NSString have different concepts of equality, so
+    // NSString-keyed Sets may generate key collisions when downcasted
+    // to String. See rdar://problem/35995647
+    allowingDuplicates: (DerivedValue.self == String.self)
+  ) { member in
+    member as? DerivedValue
   }
-  return result
 }

validation-test/stdlib/Dictionary.swift

@@ -3473,6 +3473,21 @@ DictionaryTestSuite.test("BridgedFromObjC.Nonverbatim.StringEqualityMismatch") {
   expectTrue(v == 42 || v == 23)
 }
 
+DictionaryTestSuite.test("Upcast.StringEqualityMismatch") {
+  // Upcasting from NSString to String keys changes their concept of equality,
+  // resulting in two equal keys, one of which should be discarded by the
+  // downcast. (Along with its associated value.)
+  // rdar://problem/35995647
+  let d: Dictionary<NSString, NSObject> = [
+    "cafe\u{301}": 1 as NSNumber,
+    "café": 2 as NSNumber,
+  ]
+  expectEqual(d.count, 2)
+  let d2 = d as Dictionary<String, NSObject>
+  expectEqual(d2.count, 1)
+}
+
+
 DictionaryTestSuite.test("BridgedFromObjC.Verbatim.OptionalDowncastFailure") {
   let nsd = NSDictionary(
     objects: [1 as NSNumber, 2 as NSNumber, 3 as NSNumber],

validation-test/stdlib/DictionaryTrapsObjC.swift

@@ -284,7 +284,12 @@ DictionaryTraps.test("ForcedVerbatimBridge.Value")
   }
 }
 
-DictionaryTraps.test("Downcast1") {
+DictionaryTraps.test("Downcast.Verbatim")
+  .skip(.custom(
+    { _isFastAssertConfiguration() },
+    reason: "this trap is not guaranteed to happen in -Ounchecked"))
+  .crashOutputMatches("Could not cast value of type")
+  .code {
   let d: Dictionary<NSObject, NSObject> = [ TestObjCKeyTy(10): NSObject(),
                                             NSObject() : NSObject() ]
   let d2: Dictionary<TestObjCKeyTy, NSObject> = _dictionaryDownCast(d)
@@ -296,10 +301,11 @@ DictionaryTraps.test("Downcast1") {
   for (_, _) in d2 { }
 }
 
-DictionaryTraps.test("Downcast2")
+DictionaryTraps.test("Downcast.NonVerbatimBridged")
   .skip(.custom(
     { _isFastAssertConfiguration() },
     reason: "this trap is not guaranteed to happen in -Ounchecked"))
+  .crashOutputMatches("Could not cast value of type")
   .code {
   let d: Dictionary<NSObject, NSObject> = [ TestObjCKeyTy(10): NSObject(),
                                             NSObject() : NSObject() ]