[stdlib] Set, Dictionary: Fix uniqueness checking for Cocoa indices

Author: lorentey

stdlib/public/core/Dictionary.swift

@@ -1827,6 +1827,14 @@ extension Dictionary.Index {
       _conditionallyUnreachable()
     }
   }
+
+  @inlinable
+  @inline(__always)
+  internal mutating func _isUniquelyReferenced() -> Bool {
+    defer { _fixLifetime(self) }
+    var handle = _asCocoa.handleBitPattern
+    return handle == 0 || _isUnique_native(&handle)
+  }
 #endif
 
   @usableFromInline @_transparent
@@ -1843,14 +1851,27 @@ extension Dictionary.Index {
   }
 
 #if _runtime(_ObjC)
-  @usableFromInline @_transparent
+  @usableFromInline
   internal var _asCocoa: _CocoaDictionary.Index {
-    switch _variant {
-    case .native:
-      _preconditionFailure(
-        "Attempting to access Dictionary elements using an invalid index")
-    case .cocoa(let cocoaIndex):
-      return cocoaIndex
+    @_transparent
+    get {
+      switch _variant {
+      case .native:
+        _preconditionFailure(
+          "Attempting to access Dictionary elements using an invalid index")
+      case .cocoa(let cocoaIndex):
+        return cocoaIndex
+      }
+    }
+    _modify {
+      guard case .cocoa(var cocoa) = _variant else {
+        _preconditionFailure(
+          "Attempting to access Dictionary elements using an invalid index")
+      }
+      let dummy = _HashTable.Index(bucket: _HashTable.Bucket(offset: 0), age: 0)
+      _variant = .native(dummy)
+      yield &cocoa
+      _variant = .cocoa(cocoa)
     }
   }
 #endif

stdlib/public/core/DictionaryBridging.swift

@@ -471,8 +471,8 @@ extension _CocoaDictionary: _DictionaryBuffer {
   @usableFromInline // FIXME(cocoa-index): Should be inlinable
   @_effects(releasenone)
   internal func index(after index: Index) -> Index {
-    var result = index
-    formIndex(after: &result)
+    var result = index.copy()
+    formIndex(after: &result, isUnique: true)
     return result
   }
 
@@ -484,11 +484,9 @@ extension _CocoaDictionary: _DictionaryBuffer {
 
   @usableFromInline // FIXME(cocoa-index): Should be inlinable
   @_effects(releasenone)
-  internal func formIndex(after index: inout Index) {
+  internal func formIndex(after index: inout Index, isUnique: Bool) {
     validate(index)
-    let isUnique = index.isUniquelyReferenced()
     if !isUnique { index = index.copy() }
-    _sanityCheck(index.isUniquelyReferenced())
     let storage = index.storage // FIXME: rdar://problem/44863751
     storage.currentKeyIndex += 1
   }
@@ -636,18 +634,16 @@ extension _CocoaDictionary.Index {
 }
 
 extension _CocoaDictionary.Index {
-  @inlinable
-  internal mutating func isUniquelyReferenced() -> Bool {
-    defer { _fixLifetime(self) }
-    guard _isNativePointer(_storage) else {
-      return false
+  @usableFromInline
+  internal var handleBitPattern: UInt {
+    @_effects(readonly)
+    get {
+      return unsafeBitCast(storage, to: UInt.self)
     }
-    unowned(unsafe) var storage = _bridgeObject(toNative: _storage)
-    return _isUnique_native(&storage)
   }
 
   @usableFromInline
-  internal mutating func copy() -> _CocoaDictionary.Index {
+  internal func copy() -> _CocoaDictionary.Index {
     let storage = self.storage
     return _CocoaDictionary.Index(Storage(
         storage.base,

stdlib/public/core/Set.swift

@@ -1343,6 +1343,14 @@ extension Set.Index {
       _conditionallyUnreachable()
     }
   }
+
+  @inlinable
+  @inline(__always)
+  internal mutating func _isUniquelyReferenced() -> Bool {
+    defer { _fixLifetime(self) }
+    var handle = _asCocoa.handleBitPattern
+    return handle == 0 || _isUnique_native(&handle)
+  }
 #endif
 
   @usableFromInline @_transparent
@@ -1359,14 +1367,27 @@ extension Set.Index {
   }
 
 #if _runtime(_ObjC)
-  @usableFromInline @_transparent
+  @usableFromInline
   internal var _asCocoa: _CocoaSet.Index {
-    switch _variant {
-    case .native:
-      _preconditionFailure(
-        "Attempting to access Set elements using an invalid index")
-    case .cocoa(let cocoaIndex):
-      return cocoaIndex
+    @_transparent
+    get {
+      switch _variant {
+      case .native:
+        _preconditionFailure(
+          "Attempting to access Set elements using an invalid index")
+      case .cocoa(let cocoaIndex):
+        return cocoaIndex
+      }
+    }
+    _modify {
+      guard case .cocoa(var cocoa) = _variant else {
+        _preconditionFailure(
+          "Attempting to access Set elements using an invalid index")
+      }
+      let dummy = _HashTable.Index(bucket: _HashTable.Bucket(offset: 0), age: 0)
+      _variant = .native(dummy)
+      yield &cocoa
+      _variant = .cocoa(cocoa)
     }
   }
 #endif

stdlib/public/core/SetBridging.swift

@@ -345,8 +345,8 @@ extension _CocoaSet: _SetBuffer {
   @usableFromInline // FIXME(cocoa-index): Should be inlinable
   @_effects(releasenone)
   internal func index(after index: Index) -> Index {
-    var result = index
-    formIndex(after: &result)
+    var result = index.copy()
+    formIndex(after: &result, isUnique: true)
     return result
   }
 
@@ -359,9 +359,8 @@ extension _CocoaSet: _SetBuffer {
 
   @usableFromInline // FIXME(cocoa-index): Should be inlinable
   @_effects(releasenone)
-  internal func formIndex(after index: inout Index) {
+  internal func formIndex(after index: inout Index, isUnique: Bool) {
     validate(index)
-    let isUnique = index.isUniquelyReferenced()
     if !isUnique { index = index.copy() }
     let storage = index.storage // FIXME: rdar://problem/44863751
     storage.currentKeyIndex += 1
@@ -472,18 +471,15 @@ extension _CocoaSet.Index {
 }
 
 extension _CocoaSet.Index {
-  @inlinable
-  internal mutating func isUniquelyReferenced() -> Bool {
-    defer { _fixLifetime(self) }
-    guard _isNativePointer(_storage) else {
-      return false
+  @usableFromInline
+  internal var handleBitPattern: UInt {
+    @_effects(readonly)
+    get {
+      return unsafeBitCast(storage, to: UInt.self)
     }
-    unowned(unsafe) var storage = _bridgeObject(toNative: _storage)
-    return _isUnique_native(&storage)
   }
 
-  @usableFromInline
-  internal mutating func copy() -> _CocoaSet.Index {
+  internal func copy() -> _CocoaSet.Index {
     let storage = self.storage
     return _CocoaSet.Index(Storage(
         storage.base,