Merge pull request #19602 from lorentey/dict-iterators

[stdlib] Dictionary.Keys, .Values: Implement custom iterators

Author: lorentey

stdlib/public/core/Dictionary.swift

@@ -1518,6 +1518,66 @@ extension Dictionary {
   }
 }
 
+extension Dictionary.Keys {
+  @_fixed_layout
+  public struct Iterator: IteratorProtocol {
+    @usableFromInline
+    internal var _base: Dictionary<Key, Value>.Iterator
+
+    @inlinable
+    internal init(_ base: Dictionary<Key, Value>.Iterator) {
+      self._base = base
+    }
+
+    @inlinable
+    public mutating func next() -> Key? {
+#if _runtime(_ObjC)
+      if case .cocoa(let cocoa) = _base._variant {
+        _base._cocoaPath()
+        guard let cocoaKey = cocoa.nextKey() else { return nil }
+        return _forceBridgeFromObjectiveC(cocoaKey, Key.self)
+      }
+#endif
+      return _base._asNative.nextKey()
+    }
+  }
+
+  @inlinable
+  public func makeIterator() -> Iterator {
+    return Iterator(_variant.makeIterator())
+  }
+}
+
+extension Dictionary.Values {
+  @_fixed_layout
+  public struct Iterator: IteratorProtocol {
+    @usableFromInline
+    internal var _base: Dictionary<Key, Value>.Iterator
+
+    @inlinable
+    internal init(_ base: Dictionary<Key, Value>.Iterator) {
+      self._base = base
+    }
+
+    @inlinable
+    public mutating func next() -> Value? {
+#if _runtime(_ObjC)
+      if case .cocoa(let cocoa) = _base._variant {
+        _base._cocoaPath()
+        guard let (_, cocoaValue) = cocoa.next() else { return nil }
+        return _forceBridgeFromObjectiveC(cocoaValue, Value.self)
+      }
+#endif
+      return _base._asNative.nextValue()
+    }
+  }
+
+  @inlinable
+  public func makeIterator() -> Iterator {
+    return Iterator(_variant.makeIterator())
+  }
+}
+
 extension Dictionary: Equatable where Value: Equatable {
   @inlinable
   public static func == (lhs: [Key: Value], rhs: [Key: Value]) -> Bool {

stdlib/public/core/DictionaryBridging.swift

@@ -705,7 +705,7 @@ extension _CocoaDictionary.Iterator: IteratorProtocol {
   internal typealias Element = (key: AnyObject, value: AnyObject)
 
   @usableFromInline
-  internal func next() -> Element? {
+  internal func nextKey() -> AnyObject? {
     if itemIndex < 0 {
       return nil
     }
@@ -733,6 +733,12 @@ extension _CocoaDictionary.Iterator: IteratorProtocol {
     let itemsPtr = _UnmanagedAnyObjectArray(itemsPtrUP)
     let key: AnyObject = itemsPtr[itemIndex]
     itemIndex += 1
+    return key
+  }
+
+  @usableFromInline
+  internal func next() -> Element? {
+    guard let key = nextKey() else { return nil }
     let value: AnyObject = base.object.object(forKey: key)!
     return (key, value)
   }

stdlib/public/core/NativeDictionary.swift

@@ -653,6 +653,18 @@ extension _NativeDictionary.Iterator: IteratorProtocol {
   @usableFromInline
   internal typealias Element = (key: Key, value: Value)
 
+  @inlinable
+  internal mutating func nextKey() -> Key? {
+    guard let index = iterator.next() else { return nil }
+    return base.uncheckedKey(at: index)
+  }
+
+  @inlinable
+  internal mutating func nextValue() -> Value? {
+    guard let index = iterator.next() else { return nil }
+    return base.uncheckedValue(at: index)
+  }
+
   @inlinable
   internal mutating func next() -> Element? {
     guard let index = iterator.next() else { return nil }