[stdlib] Set, Dictionary: Unify casting operations

Author: lorentey

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,18 +45,15 @@
 public func _dictionaryUpCast<DerivedKey, DerivedValue, BaseKey, BaseValue>(
     _ source: Dictionary<DerivedKey, DerivedValue>
 ) -> Dictionary<BaseKey, BaseValue> {
-  // 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.insertWithGuaranteedCapacity(
-      key: k as! BaseKey,
-      value: v as! BaseValue,
-      allowingDuplicates: allowDuplicates)
-  }
-  return Dictionary(_native: builder)
+  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.
@@ -77,19 +98,15 @@ public func _dictionaryDownCast<BaseKey, BaseValue, DerivedKey, DerivedValue>(
   // 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)
+  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.
@@ -120,21 +137,19 @@ public func _dictionaryDownCastConditional<
 >(
   _ source: Dictionary<BaseKey, BaseValue>
 ) -> Dictionary<DerivedKey, DerivedValue>? {
-
-  // 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.insertWithGuaranteedCapacity(
-      key: k1,
-      value: v1,
-      allowingDuplicates: allowDuplicates)
+  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 Dictionary(_native: builder)
 }

stdlib/public/core/NativeDictionary.swift

@@ -460,28 +460,6 @@ 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.
@@ -507,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

@@ -124,6 +124,7 @@ extension _NativeSet { // Low-level unchecked operations
   }
 
   @_alwaysEmitIntoClient @inlinable // Introduced in 5.1
+  @inline(__always)
   internal func uncheckedAssign(
     at bucket: Bucket,
     to element: __owned Element
@@ -356,33 +357,12 @@ 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.
   @inlinable
   internal func _unsafeInsertNew(_ element: __owned Element) {
-    _internalInvariant(count + 1 <= capacity)
+    _precondition(count < capacity)
     let hashValue = self.hashValue(for: element)
     if _isDebugAssertConfiguration() {
       // In debug builds, perform a full lookup and trap if we detect duplicate
@@ -461,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,25 +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> {
-  // 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 Set(_native: builder)
+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.
@@ -64,18 +86,15 @@ public func _setDownCast<BaseValue, DerivedValue>(_ source: Set<BaseValue>)
   // 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)
+  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.
@@ -102,17 +121,13 @@ internal func _setDownCastConditionalIndirect<SourceValue, TargetValue>(
 public func _setDownCastConditional<BaseValue, DerivedValue>(
   _ source: Set<BaseValue>
 ) -> Set<DerivedValue>? {
-  // 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 {
-    guard let derivedMember = member as? DerivedValue else { return nil }
-    _ = builder.insertWithGuaranteedCapacity(
-      derivedMember,
-      allowingDuplicates: allowDuplicates)
+  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 Set(_native: builder)
 }