[stdlib] Fix Set/Dictionary casting issues

- Fix Set/Dictionary up/downcasting with String keys.
- Improve error handling.

Author: lorentey

stdlib/public/core/DictionaryCasting.swift

@@ -21,12 +21,18 @@
 public func _dictionaryUpCast<DerivedKey, DerivedValue, BaseKey, BaseValue>(
     _ source: Dictionary<DerivedKey, DerivedValue>
 ) -> Dictionary<BaseKey, BaseValue> {
-  var builder = _DictionaryBuilder<BaseKey, BaseValue>(count: source.count)
-
+  // String and NSString have different concepts of equality, so
+  // NSString-keyed Dictionaries may generate key collisions when "upcasted"
+  // to String. See rdar://problem/35995647
+  let allowDuplicates = (BaseKey.self == String.self)
+  var builder = _NativeDictionary<BaseKey, BaseValue>(capacity: source.count)
   for (k, v) in source {
-    builder.add(key:k as! BaseKey, value: v as! BaseValue)
+    builder.insertWithGuaranteedCapacity(
+      key: k as! BaseKey,
+      value: v as! BaseValue,
+      allowingDuplicates: allowDuplicates)
   }
-  return builder.take()
+  return Dictionary(_native: builder)
 }
 
 /// Called by the casting machinery.
@@ -66,7 +72,24 @@ 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.
+
+  // String and NSString have different concepts of equality, so
+  // NSString-keyed Dictionaries may generate key collisions when downcasted
+  // to String. See rdar://problem/35995647
+  let allowDuplicates = (DerivedKey.self == String.self)
+  var builder = _NativeDictionary<DerivedKey, DerivedValue>(
+    capacity: source.count)
+  for (k, v) in source {
+    builder.insertWithGuaranteedCapacity(
+      key: k as! DerivedKey,
+      value: v as! DerivedValue,
+      allowingDuplicates: allowDuplicates)
+  }
+  return Dictionary(_native: builder)
 }
 
 /// Called by the casting machinery.
@@ -98,11 +121,20 @@ public func _dictionaryDownCastConditional<
   _ source: Dictionary<BaseKey, BaseValue>
 ) -> Dictionary<DerivedKey, DerivedValue>? {
 
-  var builder = _DictionaryBuilder<DerivedKey, DerivedValue>(count: source.count)
+  // String and NSString have different concepts of equality, so
+  // NSString-keyed Dictionaries may generate key collisions when downcasted
+  // to String. See rdar://problem/35995647
+  let allowDuplicates = (DerivedKey.self == String.self)
+
+  var builder = _NativeDictionary<DerivedKey, DerivedValue>(
+    capacity: 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)
+    builder.insertWithGuaranteedCapacity(
+      key: k1,
+      value: v1,
+      allowingDuplicates: allowDuplicates)
   }
-  return builder.take()
+  return Dictionary(_native: builder)
 }

stdlib/public/core/NativeDictionary.swift

@@ -460,6 +460,28 @@ internal func KEY_TYPE_OF_DICTIONARY_VIOLATES_HASHABLE_REQUIREMENTS(
 }
 
 extension _NativeDictionary { // Insertions
+  /// Insert a new element into uniquely held storage.  Storage must be uniquely
+  /// referenced with adequate capacity.  If `allowingDuplicates` is false,
+  /// `element` must not be already present in the dictionary.
+  @_alwaysEmitIntoClient @inlinable // Introduced in 5.1
+  internal mutating func insertWithGuaranteedCapacity(
+    key: __owned Key,
+    value: __owned Value,
+    allowingDuplicates: Bool
+  ) {
+    _precondition(count < capacity)
+    if !allowingDuplicates {
+      _unsafeInsertNew(key: key, value: value)
+      return
+    }
+    let (bucket, found) = find(key)
+    if found {
+      (_values + bucket.offset).pointee = value
+    } else {
+      _insert(at: bucket, key: key, value: value)
+    }
+  }
+
   /// Insert a new element into uniquely held storage.
   /// Storage must be uniquely referenced with adequate capacity.
   /// The `key` must not be already present in the Dictionary.

stdlib/public/core/NativeSet.swift

@@ -117,10 +117,20 @@ 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
+  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
@@ -346,6 +356,27 @@ internal func ELEMENT_TYPE_OF_SET_VIOLATES_HASHABLE_REQUIREMENTS(
 }
 
 extension _NativeSet { // Insertions
+  /// Insert a new element into uniquely held storage.  Storage must be uniquely
+  /// referenced with adequate capacity.  If `allowingDuplicates` is false,
+  /// `element` must not be already present in the set.
+  @_alwaysEmitIntoClient @inlinable // Introduced in 5.1
+  internal mutating func insertWithGuaranteedCapacity(
+    _ element: __owned Element,
+    allowingDuplicates: Bool
+  ) {
+    _precondition(count < capacity)
+    if !allowingDuplicates {
+      _unsafeInsertNew(element)
+      return
+    }
+    let (bucket, found) = find(element)
+    if found {
+      uncheckedAssign(at: bucket, to: element)
+    } else {
+      _unsafeInsertNew(element, at: bucket)
+    }
+  }
+
   /// Insert a new element into uniquely held storage.
   /// Storage must be uniquely referenced with adequate capacity.
   /// The `element` must not be already present in the Set.

stdlib/public/core/SetCasting.swift

@@ -19,11 +19,18 @@
 @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)
+  // String and NSString have different concepts of equality, so Set<NSString>
+  // may generate key collisions when "upcasted" to Set<String>.
+  // See rdar://problem/35995647
+  let allowDuplicates = (BaseValue.self == String.self)
+
+  var builder = _NativeSet<BaseValue>(capacity: source.count)
+  for member in source {
+    _ = builder.insertWithGuaranteedCapacity(
+      member as! BaseValue,
+      allowingDuplicates: allowDuplicates)
   }
-  return builder.take()
+  return Set(_native: builder)
 }
 
 /// Called by the casting machinery.
@@ -54,7 +61,21 @@ 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.
+
+  // String and NSString have different concepts of equality, so
+  // NSString-keyed Sets may generate key collisions when downcasted
+  // to String. See rdar://problem/35995647
+  let allowDuplicates = (DerivedValue.self == String.self)
+
+  var builder = _NativeSet<DerivedValue>(capacity: source.count)
+  for member in source {
+    builder.insertWithGuaranteedCapacity(
+      member as! DerivedValue,
+      allowingDuplicates: allowDuplicates)
+  }
+  return Set(_native: builder)
 }
 
 /// Called by the casting machinery.
@@ -81,13 +102,17 @@ internal func _setDownCastConditionalIndirect<SourceValue, TargetValue>(
 public func _setDownCastConditional<BaseValue, DerivedValue>(
   _ source: Set<BaseValue>
 ) -> Set<DerivedValue>? {
-  var result = Set<DerivedValue>(minimumCapacity: source.count)
+  // String and NSString have different concepts of equality, so
+  // NSString-keyed Sets may generate key collisions when downcasted
+  // to String. See rdar://problem/35995647
+  let allowDuplicates = (DerivedValue.self == String.self)
+
+  var builder = _NativeSet<DerivedValue>(capacity: source.count)
   for member in source {
-    if let derivedMember = member as? DerivedValue {
-      result.insert(derivedMember)
-      continue
-    }
-    return nil
+    guard let derivedMember = member as? DerivedValue else { return nil }
+    _ = builder.insertWithGuaranteedCapacity(
+      derivedMember,
+      allowingDuplicates: allowDuplicates)
   }
-  return result
+  return Set(_native: builder)
 }