Merge pull request #19612 from lorentey/opaque-cocoa-indices

[stdlib] Set, Dictionary: Resilient Cocoa indices, vol 2

Author: lorentey

stdlib/public/core/Dictionary.swift

@@ -647,6 +647,11 @@ extension Dictionary: Collection {
     return _variant.index(after: i)
   }
 
+  @inlinable
+  public func formIndex(after i: inout Index) {
+    _variant.formIndex(after: &i)
+  }
+
   /// Returns the index for the given key.
   ///
   /// If the given key is found in the dictionary, this method returns an index
@@ -1350,6 +1355,11 @@ extension Dictionary {
       return _variant.index(after: i)
     }
 
+    @inlinable
+    public func formIndex(after i: inout Index) {
+      _variant.formIndex(after: &i)
+    }
+
     @inlinable
     public subscript(position: Index) -> Element {
       return _variant.key(at: position)
@@ -1456,6 +1466,11 @@ extension Dictionary {
       return _variant.index(after: i)
     }
 
+    @inlinable
+    public func formIndex(after i: inout Index) {
+      _variant.formIndex(after: &i)
+    }
+
     @inlinable
     public subscript(position: Index) -> Element {
       // FIXME(accessors): Provide a _read
@@ -1827,6 +1842,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 +1866,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
@@ -1897,8 +1933,8 @@ extension Dictionary.Index: Comparable {
 }
 
 extension Dictionary.Index: Hashable {
-  @_effects(readonly) // FIXME(cocoa-index): Make inlinable
-  public func hash(into hasher: inout Hasher) {
+  public // FIXME(cocoa-index): Make inlinable
+  func hash(into hasher: inout Hasher) {
   #if _runtime(_ObjC)
     switch _variant {
     case .native(let nativeIndex):

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
   }
 
@@ -483,11 +483,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.storage = index.copy() }
+    if !isUnique { index = index.copy() }
     let storage = index.storage // FIXME: rdar://problem/44863751
     storage.currentKeyIndex += 1
   }
@@ -575,17 +573,34 @@ extension _CocoaDictionary {
   @usableFromInline
   internal struct Index {
     @usableFromInline
-    internal var storage: Storage
+    internal var _object: Builtin.BridgeObject
+    @usableFromInline
+    internal var _storage: Builtin.BridgeObject
+
+    internal var object: AnyObject {
+      @inline(__always)
+      get {
+        return _bridgeObject(toNonTaggedObjC: _object)
+      }
+    }
+
+    internal var storage: Storage {
+      @inline(__always)
+      get {
+        let storage = _bridgeObject(toNative: _storage)
+        return unsafeDowncast(storage, to: Storage.self)
+      }
+    }
 
     internal init(_ storage: Storage) {
-      self.storage = storage
+      self._object = _bridgeObject(fromNonTaggedObjC: storage.base.object)
+      self._storage = _bridgeObject(fromNative: storage)
     }
   }
 }
 
 extension _CocoaDictionary.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
@@ -618,15 +633,21 @@ extension _CocoaDictionary.Index {
 }
 
 extension _CocoaDictionary.Index {
-  @inlinable
-  internal mutating func isUniquelyReferenced() -> Bool {
-    return _isUnique_native(&storage)
+  @usableFromInline
+  internal var handleBitPattern: UInt {
+    @_effects(readonly)
+    get {
+      return unsafeBitCast(storage, to: UInt.self)
+    }
   }
 
   @usableFromInline
-  internal mutating func copy() -> Storage {
+  internal func copy() -> _CocoaDictionary.Index {
     let storage = self.storage
-    return Storage(storage.base, storage.allKeys, storage.currentKeyIndex)
+    return _CocoaDictionary.Index(Storage(
+        storage.base,
+        storage.allKeys,
+        storage.currentKeyIndex))
   }
 }
 
@@ -647,7 +668,7 @@ extension _CocoaDictionary.Index {
   internal var age: Int32 {
     @_effects(readonly)
     get {
-      return _HashTable.age(for: storage.base.object)
+      return _HashTable.age(for: object)
     }
   }
 }

stdlib/public/core/DictionaryVariant.swift

@@ -189,6 +189,20 @@ extension Dictionary._Variant: _DictionaryBuffer {
     }
   }
 
+  @inlinable
+  internal func formIndex(after index: inout Index) {
+    switch self {
+    case .native(let native):
+      index = native.index(after: index)
+#if _runtime(_ObjC)
+    case .cocoa(let cocoa):
+      cocoaPath()
+      let isUnique = index._isUniquelyReferenced()
+      cocoa.formIndex(after: &index._asCocoa, isUnique: isUnique)
+#endif
+    }
+  }
+
   @inlinable
   @inline(__always)
   internal func index(forKey key: Key) -> Index? {

stdlib/public/core/Set.swift

@@ -340,6 +340,11 @@ extension Set: Collection {
     return _variant.index(after: i)
   }
 
+  @inlinable
+  public func formIndex(after i: inout Index) {
+    _variant.formIndex(after: &i)
+  }
+
   // APINAMING: complexity docs are broadly missing in this file.
 
   /// Returns the index of the given element in the set, or `nil` if the
@@ -1343,6 +1348,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 +1372,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
@@ -1418,8 +1444,8 @@ extension Set.Index: Hashable {
   ///
   /// - Parameter hasher: The hasher to use when combining the components
   ///   of this instance.
-  @_effects(readonly) // FIXME(cocoa-index): Make inlinable
-  public func hash(into hasher: inout Hasher) {
+  public // FIXME(cocoa-index): Make inlinable
+  func hash(into hasher: inout Hasher) {
   #if _runtime(_ObjC)
     switch _variant {
     case .native(let nativeIndex):

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
   }
 
@@ -358,11 +358,9 @@ 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.storage = index.copy() }
+    if !isUnique { index = index.copy() }
     let storage = index.storage // FIXME: rdar://problem/44863751
     storage.currentKeyIndex += 1
   }
@@ -412,17 +410,34 @@ extension _CocoaSet {
   @usableFromInline
   internal struct Index {
     @usableFromInline
-    internal var storage: Storage
+    internal var _object: Builtin.BridgeObject
+    @usableFromInline
+    internal var _storage: Builtin.BridgeObject
+
+    internal var object: AnyObject {
+      @inline(__always)
+      get {
+        return _bridgeObject(toNonTaggedObjC: _object)
+      }
+    }
+
+    internal var storage: Storage {
+      @inline(__always)
+      get {
+        let storage = _bridgeObject(toNative: _storage)
+        return unsafeDowncast(storage, to: Storage.self)
+      }
+    }
 
     internal init(_ storage: __owned Storage) {
-      self.storage = storage
+      self._object = _bridgeObject(fromNonTaggedObjC: storage.base.object)
+      self._storage = _bridgeObject(fromNative: 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
@@ -455,15 +470,20 @@ extension _CocoaSet.Index {
 }
 
 extension _CocoaSet.Index {
-  @inlinable
-  internal mutating func isUniquelyReferenced() -> Bool {
-    return _isUnique_native(&storage)
+  @usableFromInline
+  internal var handleBitPattern: UInt {
+    @_effects(readonly)
+    get {
+      return unsafeBitCast(storage, to: UInt.self)
+    }
   }
 
-  @usableFromInline
-  internal mutating func copy() -> Storage {
+  internal func copy() -> _CocoaSet.Index {
     let storage = self.storage
-    return Storage(storage.base, storage.allKeys, storage.currentKeyIndex)
+    return _CocoaSet.Index(Storage(
+        storage.base,
+        storage.allKeys,
+        storage.currentKeyIndex))
   }
 }
 
@@ -484,7 +504,7 @@ extension _CocoaSet.Index {
   internal var age: Int32 {
     @_effects(releasenone)
     get {
-      return _HashTable.age(for: storage.base.object)
+      return _HashTable.age(for: object)
     }
   }
 }