[stdlib] Turn _CocoaSet.Index into a COW value type with opaque storage

This enables making the struct @_fixed_layout, which should win us back some of the performance lost by hiding Cocoa indexing internals.

Author: lorentey

stdlib/public/core/Set.swift

@@ -1428,7 +1428,7 @@ extension Set.Index: Hashable {
     case .cocoa(let cocoaIndex):
       _cocoaPath()
       hasher.combine(1 as UInt8)
-      hasher.combine(cocoaIndex.currentKeyIndex)
+      hasher.combine(cocoaIndex.storage.currentKeyIndex)
     }
   #else
     hasher.combine(_asNative.bucket.offset)

stdlib/public/core/SetBridging.swift

@@ -328,33 +328,43 @@ extension _CocoaSet: _SetBuffer {
   internal var startIndex: Index {
     @_effects(releasenone)
     get {
-      return Index(self, startIndex: ())
+      let allKeys = _stdlib_NSSet_allObjects(self.object)
+      return Index(Index.Storage(self, allKeys, 0))
     }
   }
 
   @usableFromInline // FIXME(cocoa-index): Should be inlinable
   internal var endIndex: Index {
     @_effects(releasenone)
     get {
-      return Index(self, endIndex: ())
+      let allKeys = _stdlib_NSSet_allObjects(self.object)
+      return Index(Index.Storage(self, allKeys, allKeys.value))
     }
   }
 
   @usableFromInline // FIXME(cocoa-index): Should be inlinable
   @_effects(releasenone)
-  internal func index(after i: Index) -> Index {
-    var i = i
-    formIndex(after: &i)
-    return i
+  internal func index(after index: Index) -> Index {
+    var result = index
+    formIndex(after: &result)
+    return result
+  }
+
+  internal func validate(_ index: Index) {
+    _precondition(index.storage.base.object === self.object,
+      "Invalid index")
+    _precondition(index.storage.currentKeyIndex < index.storage.allKeys.value,
+      "Attempt to access endIndex")
   }
 
   @usableFromInline // FIXME(cocoa-index): Should be inlinable
   @_effects(releasenone)
-  internal func formIndex(after i: inout Index) {
-    _precondition(i.base.object === self.object, "Invalid index")
-    _precondition(i.currentKeyIndex < i.allKeys.value,
-      "Cannot increment endIndex")
-    i.currentKeyIndex += 1
+  internal func formIndex(after index: inout Index) {
+    validate(index)
+    let isUnique = index.isUniquelyReferenced()
+    if !isUnique { index.storage = index.copy() }
+    let storage = index.storage // FIXME: rdar://problem/44863751
+    storage.currentKeyIndex += 1
   }
 
   @usableFromInline // FIXME(cocoa-index): Should be inlinable
@@ -369,16 +379,13 @@ extension _CocoaSet: _SetBuffer {
     }
 
     let allKeys = _stdlib_NSSet_allObjects(object)
-    var keyIndex = -1
     for i in 0..<allKeys.value {
       if _stdlib_NSObject_isEqual(element, allKeys[i]) {
-        keyIndex = i
-        break
+        return Index(Index.Storage(self, allKeys, i))
       }
     }
-    _sanityCheck(keyIndex >= 0,
-        "Key was found in fast path, but not found later?")
-    return Index(self, allKeys, keyIndex)
+    _sanityCheckFailure(
+      "An NSSet member wasn't listed amongst its enumerated contents")
   }
 
   @inlinable
@@ -401,9 +408,22 @@ extension _CocoaSet: _SetBuffer {
 }
 
 extension _CocoaSet {
-  // FIXME(cocoa-index): Overhaul and make @_fixed_layout
+  @_fixed_layout
   @usableFromInline
   internal struct Index {
+    @usableFromInline
+    internal var storage: Storage
+
+    internal init(_ storage: __owned Storage) {
+      self.storage = storage
+    }
+  }
+}
+
+extension _CocoaSet.Index {
+  // FIXME(cocoa-index): Try using an NSEnumerator to speed this up.
+  @usableFromInline
+  internal class Storage {
     // Assumption: we rely on NSDictionary.getObjects when being
     // repeatedly called on the same NSDictionary, returning items in the same
     // order every time.
@@ -422,18 +442,6 @@ extension _CocoaSet {
     /// Index into `allKeys`
     internal var currentKeyIndex: Int
 
-    internal init(_ base: __owned _CocoaSet, startIndex: ()) {
-      self.base = base
-      self.allKeys = _stdlib_NSSet_allObjects(base.object)
-      self.currentKeyIndex = 0
-    }
-
-    internal init(_ base: __owned _CocoaSet, endIndex: ()) {
-      self.base = base
-      self.allKeys = _stdlib_NSSet_allObjects(base.object)
-      self.currentKeyIndex = allKeys.value
-    }
-
     internal init(
       _ base: __owned _CocoaSet,
       _ allKeys: __owned _HeapBuffer<Int, AnyObject>,
@@ -446,15 +454,28 @@ extension _CocoaSet {
   }
 }
 
+extension _CocoaSet.Index {
+  @inlinable
+  internal mutating func isUniquelyReferenced() -> Bool {
+    return _isUnique_native(&storage)
+  }
+
+  @usableFromInline
+  internal mutating func copy() -> Storage {
+    let storage = self.storage
+    return Storage(storage.base, storage.allKeys, storage.currentKeyIndex)
+  }
+}
+
 extension _CocoaSet.Index {
   @usableFromInline // FIXME(cocoa-index): Make inlinable
   @nonobjc
   internal var element: AnyObject {
     @_effects(readonly)
     get {
-      _precondition(currentKeyIndex < allKeys.value,
+      _precondition(storage.currentKeyIndex < storage.allKeys.value,
         "Attempting to access Set elements using an invalid index")
-      return allKeys[currentKeyIndex]
+      return storage.allKeys[storage.currentKeyIndex]
     }
   }
 
@@ -463,7 +484,7 @@ extension _CocoaSet.Index {
   internal var age: Int32 {
     @_effects(releasenone)
     get {
-      return _HashTable.age(for: base.object)
+      return _HashTable.age(for: storage.base.object)
     }
   }
 }
@@ -472,19 +493,19 @@ extension _CocoaSet.Index: Equatable {
   @usableFromInline // FIXME(cocoa-index): Make inlinable
   @_effects(readonly)
   internal static func == (lhs: _CocoaSet.Index, rhs: _CocoaSet.Index) -> Bool {
-    _precondition(lhs.base.object === rhs.base.object,
+    _precondition(lhs.storage.base.object === rhs.storage.base.object,
       "Comparing indexes from different sets")
-    return lhs.currentKeyIndex == rhs.currentKeyIndex
+    return lhs.storage.currentKeyIndex == rhs.storage.currentKeyIndex
   }
 }
 
 extension _CocoaSet.Index: Comparable {
   @usableFromInline // FIXME(cocoa-index): Make inlinable
   @_effects(readonly)
   internal static func < (lhs: _CocoaSet.Index, rhs: _CocoaSet.Index) -> Bool {
-    _precondition(lhs.base.object === rhs.base.object,
+    _precondition(lhs.storage.base.object === rhs.storage.base.object,
       "Comparing indexes from different sets")
-    return lhs.currentKeyIndex < rhs.currentKeyIndex
+    return lhs.storage.currentKeyIndex < rhs.storage.currentKeyIndex
   }
 }