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

[stdlib] Set, Dictionary: Make Cocoa indices resilient for now

Author: lorentey

stdlib/public/core/Dictionary.swift

@@ -1897,7 +1897,7 @@ extension Dictionary.Index: Comparable {
 }
 
 extension Dictionary.Index: Hashable {
-  @inlinable
+  @_effects(readonly) // FIXME(cocoa-index): Make inlinable
   public func hash(into hasher: inout Hasher) {
   #if _runtime(_ObjC)
     switch _variant {
@@ -1907,7 +1907,7 @@ extension Dictionary.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/DictionaryBridging.swift

@@ -450,33 +450,50 @@ extension _CocoaDictionary: _DictionaryBuffer {
   @usableFromInline
   internal typealias Value = AnyObject
 
-  @inlinable
+  @usableFromInline // FIXME(cocoa-index): Should be inlinable
   internal var startIndex: Index {
-    return Index(self, startIndex: ())
+    @_effects(releasenone)
+    get {
+      let allKeys = _stdlib_NSDictionary_allKeys(self.object)
+      return Index(Index.Storage(self, allKeys, 0))
+    }
   }
 
-  @inlinable
+  @usableFromInline // FIXME(cocoa-index): Should be inlinable
   internal var endIndex: Index {
-    return Index(self, endIndex: ())
+    @_effects(releasenone)
+    get {
+      let allKeys = _stdlib_NSDictionary_allKeys(self.object)
+      return Index(Index.Storage(self, allKeys, allKeys.value))
+    }
   }
 
-  @inlinable
-  internal func index(after i: Index) -> Index {
-    var i = i
-    formIndex(after: &i)
-    return i
+  @usableFromInline // FIXME(cocoa-index): Should be inlinable
+  @_effects(releasenone)
+  internal func index(after index: Index) -> Index {
+    var result = index
+    formIndex(after: &result)
+    return result
   }
 
-  @usableFromInline
+  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
+  @usableFromInline // FIXME(cocoa-index): Should be inlinable
+  @_effects(releasenone)
   internal func index(forKey key: Key) -> Index? {
     // Fast path that does not involve creating an array of all keys.  In case
     // the key is present, this lookup is a penalty for the slow path, but the
@@ -487,16 +504,13 @@ extension _CocoaDictionary: _DictionaryBuffer {
     }
 
     let allKeys = _stdlib_NSDictionary_allKeys(object)
-    var keyIndex = -1
     for i in 0..<allKeys.value {
       if _stdlib_NSObject_isEqual(key, 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 NSDictionary key wassn't listed amongst its enumerated contents")
   }
 
   @inlinable
@@ -516,28 +530,28 @@ extension _CocoaDictionary: _DictionaryBuffer {
     return object.object(forKey: key)
   }
 
-  @inlinable
-  @inline(__always)
+  @usableFromInline // FIXME(cocoa-index): Should be inlinable
+  @_effects(releasenone)
   internal func lookup(_ index: Index) -> (key: Key, value: Value) {
-    _precondition(index.base.object === self.object, "Invalid index")
-    let key: Key = index.allKeys[index.currentKeyIndex]
-    let value: Value = index.base.object.object(forKey: key)!
+    _precondition(index.storage.base.object === self.object, "Invalid index")
+    let key: Key = index.storage.allKeys[index.storage.currentKeyIndex]
+    let value: Value = index.storage.base.object.object(forKey: key)!
     return (key, value)
   }
 
-  @inlinable
-  @inline(__always)
+  @usableFromInline // FIXME(cocoa-index): Make inlinable
+  @_effects(releasenone)
   func key(at index: Index) -> Key {
-    _precondition(index.base.object === self.object, "Invalid index")
+    _precondition(index.storage.base.object === self.object, "Invalid index")
     return index.key
   }
 
-  @inlinable
-  @inline(__always)
+  @usableFromInline // FIXME(cocoa-index): Make inlinable
+  @_effects(releasenone)
   func value(at index: Index) -> Value {
-    _precondition(index.base.object === self.object, "Invalid index")
-    let key = index.allKeys[index.currentKeyIndex]
-    return index.base.object.object(forKey: key)!
+    _precondition(index.storage.base.object === self.object, "Invalid index")
+    let key = index.storage.allKeys[index.storage.currentKeyIndex]
+    return index.storage.base.object.object(forKey: key)!
   }
 }
 
@@ -557,45 +571,40 @@ extension _CocoaDictionary {
 }
 
 extension _CocoaDictionary {
-  @_fixed_layout // FIXME(sil-serialize-all)
+  @_fixed_layout
   @usableFromInline
   internal struct Index {
-    // Assumption: we rely on NSDictionary.getObjects when being
+    @usableFromInline
+    internal var storage: Storage
+
+    internal init(_ storage: Storage) {
+      self.storage = 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
     // order every time.
     // Similarly, the same assumption holds for NSSet.allObjects.
 
     /// A reference to the NSDictionary, which owns members in `allObjects`,
     /// or `allKeys`, for NSSet and NSDictionary respectively.
-    @usableFromInline // FIXME(sil-serialize-all)
     internal let base: _CocoaDictionary
     // FIXME: swift-3-indexing-model: try to remove the cocoa reference, but
     // make sure that we have a safety check for accessing `allKeys`.  Maybe
     // move both into the dictionary/set itself.
 
     /// An unowned array of keys.
-    @usableFromInline // FIXME(sil-serialize-all)
     internal var allKeys: _HeapBuffer<Int, AnyObject>
 
     /// Index into `allKeys`
-    @usableFromInline // FIXME(sil-serialize-all)
     internal var currentKeyIndex: Int
 
-    @inlinable // FIXME(sil-serialize-all)
-    internal init(_ base: __owned _CocoaDictionary, startIndex: ()) {
-      self.base = base
-      self.allKeys = _stdlib_NSDictionary_allKeys(base.object)
-      self.currentKeyIndex = 0
-    }
-
-    @inlinable // FIXME(sil-serialize-all)
-    internal init(_ base: __owned _CocoaDictionary, endIndex: ()) {
-      self.base = base
-      self.allKeys = _stdlib_NSDictionary_allKeys(base.object)
-      self.currentKeyIndex = allKeys.value
-    }
-
-    @inlinable // FIXME(sil-serialize-all)
     internal init(
       _ base: __owned _CocoaDictionary,
       _ allKeys: __owned _HeapBuffer<Int, AnyObject>,
@@ -610,44 +619,62 @@ extension _CocoaDictionary {
 
 extension _CocoaDictionary.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 _CocoaDictionary.Index {
+  @usableFromInline // FIXME(cocoa-index): Make inlinable
   @nonobjc
   internal var key: AnyObject {
-    _precondition(currentKeyIndex < allKeys.value,
-      "Attempting to access Dictionary elements using an invalid index")
-    return allKeys[currentKeyIndex]
+    @_effects(readonly)
+    get {
+      _precondition(storage.currentKeyIndex < storage.allKeys.value,
+        "Attempting to access Dictionary elements using an invalid index")
+      return storage.allKeys[storage.currentKeyIndex]
+    }
   }
 
-  @usableFromInline
+  @usableFromInline // FIXME(cocoa-index): Make inlinable
   @nonobjc
   internal var age: Int32 {
-    @_effects(releasenone)
+    @_effects(readonly)
     get {
-      return _HashTable.age(for: base.object)
+      return _HashTable.age(for: storage.base.object)
     }
   }
 }
 
 extension _CocoaDictionary.Index: Equatable {
-  @inlinable
+  @usableFromInline // FIXME(cocoa-index): Make inlinable
+  @_effects(readonly)
   internal static func == (
     lhs: _CocoaDictionary.Index,
     rhs: _CocoaDictionary.Index
   ) -> Bool {
-    _precondition(lhs.base.object === rhs.base.object,
+    _precondition(lhs.storage.base.object === rhs.storage.base.object,
       "Comparing indexes from different dictionaries")
-    return lhs.currentKeyIndex == rhs.currentKeyIndex
+    return lhs.storage.currentKeyIndex == rhs.storage.currentKeyIndex
   }
 }
 
 extension _CocoaDictionary.Index: Comparable {
-  @inlinable
+  @usableFromInline // FIXME(cocoa-index): Make inlinable
+  @_effects(readonly)
   internal static func < (
     lhs: _CocoaDictionary.Index,
     rhs: _CocoaDictionary.Index
   ) -> Bool {
-    _precondition(lhs.base.object === rhs.base.object,
+    _precondition(lhs.storage.base.object === rhs.storage.base.object,
       "Comparing indexes from different dictionaries")
-    return lhs.currentKeyIndex < rhs.currentKeyIndex
+    return lhs.storage.currentKeyIndex < rhs.storage.currentKeyIndex
   }
 }
 

stdlib/public/core/Set.swift

@@ -1418,7 +1418,7 @@ extension Set.Index: Hashable {
   ///
   /// - Parameter hasher: The hasher to use when combining the components
   ///   of this instance.
-  @inlinable
+  @_effects(readonly) // FIXME(cocoa-index): Make inlinable
   public func hash(into hasher: inout Hasher) {
   #if _runtime(_ObjC)
     switch _variant {
@@ -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)