Merge pull request #2417 from swiftix/new-indexing-model-regression-fixes

Fix some Set-related performance regressions introduced by the new in…

Author: swiftix

stdlib/public/core/HashedCollections.swift.gyb

@@ -229,6 +229,11 @@ internal protocol _HashStorage {
   @discardableResult
   mutating func updateValue(_ value: Value, forKey key: Key) -> Value?
 
+  @discardableResult
+  mutating func insert(
+    _ value: Value, forKey key: Key
+  ) -> (inserted: Bool, memberAfterInsert: Value)
+
   @discardableResult
   mutating func remove(at index: Index) -> SequenceElement
 
@@ -577,12 +582,7 @@ public struct Set<Element : Hashable> :
   public mutating func insert(
     _ newMember: Element
   ) -> (inserted: Bool, memberAfterInsert: Element) {
-    // FIXME: optimize this to do just one lookup
-    if let i = index(of: newMember) {
-      return (inserted: false, memberAfterInsert: self[i])
-    }
-    _variantStorage.updateValue(newMember, forKey: newMember)
-    return (inserted: true, memberAfterInsert: newMember)
+    return _variantStorage.insert(newMember, forKey: newMember)
   }
 
   /// Inserts `newMember` unconditionally.
@@ -2730,6 +2730,14 @@ struct _Native${Self}Storage<${TypeParametersDecl}> :
       "don't call mutating methods on _Native${Self}Storage")
   }
 
+  @discardableResult
+  internal mutating func insert(
+    _ value: Value, forKey key: Key
+  ) -> (inserted: Bool, memberAfterInsert: Value) {
+    _sanityCheckFailure(
+      "don't call mutating methods on _Native${Self}Storage")
+  }
+
   @discardableResult
   internal mutating func remove(at index: Index) -> SequenceElement {
     _sanityCheckFailure(
@@ -3421,6 +3429,13 @@ internal struct _Cocoa${Self}Storage : _HashStorage {
     _sanityCheckFailure("cannot mutate NS${Self}")
   }
 
+  @discardableResult
+  internal mutating func insert(
+    _ value: Value, forKey key: Key
+  ) -> (inserted: Bool, memberAfterInsert: Value) {
+    _sanityCheckFailure("cannot mutate NS${Self}")
+  }
+
   @discardableResult
   internal mutating func remove(at index: Index) -> SequenceElement {
     _sanityCheckFailure("cannot mutate NS${Self}")
@@ -3843,6 +3858,61 @@ internal enum _Variant${Self}Storage<${TypeParametersDecl}> : _HashStorage {
     }
   }
 
+  internal mutating func nativeInsert(
+    _ value: Value, forKey key: Key
+  ) -> (inserted: Bool, memberAfterInsert: Value) {
+    var (i, found) = asNative._find(key, startBucket: asNative._bucket(key))
+
+    if found {
+%if Self == 'Set':
+      return (inserted: false, memberAfterInsert: asNative.key(at: i.offset))
+%elif Self == 'Dictionary':
+      return (inserted: false, memberAfterInsert: asNative.value(at: i.offset))
+%end
+    }
+
+    let minCapacity = NativeStorage.minimumCapacity(
+      minimumCount: asNative.count + 1,
+      maxLoadFactorInverse: asNative.maxLoadFactorInverse)
+
+    let (_, capacityChanged) = ensureUniqueNativeStorage(minCapacity)
+    if capacityChanged {
+      i = asNative._find(key, startBucket: asNative._bucket(key)).pos
+    }
+
+%if Self == 'Set':
+    asNative.initializeKey(key, at: i.offset)
+    asNative.count += 1
+%elif Self == 'Dictionary':
+    asNative.initializeKey(key, value: value, at: i.offset)
+    asNative.count += 1
+%end
+
+    return (inserted: true, memberAfterInsert: value)
+  }
+
+  @discardableResult
+  internal mutating func insert(
+    _ value: Value, forKey key: Key
+  ) -> (inserted: Bool, memberAfterInsert: Value) {
+
+    if _fastPath(guaranteedNative) {
+      return nativeInsert(value, forKey: key)
+    }
+
+    switch self {
+    case .native:
+      return nativeInsert(value, forKey: key)
+    case .cocoa(let cocoaStorage):
+#if _runtime(_ObjC)
+      migrateDataToNativeStorage(cocoaStorage)
+      return nativeInsert(value, forKey: key)
+#else
+      _sanityCheckFailure("internal error: unexpected cocoa ${Self}")
+#endif
+    }
+  }
+
   /// - parameter idealBucket: The ideal bucket for the element being deleted.
   /// - parameter offset: The offset of the element that will be deleted.
   /// Precondition: there should be an initialized entry at offset.