Make KeyPath-based API more like Combine

Sources/KVOSequence/NSObject+KVOSequence.swift

@@ -1,20 +1,30 @@
 import Foundation
 
 public extension NSObject {
-    /// An `AsyncSequence` that produces a new element whenever the observed value changes.
+    /// An `AsyncSequence` that produces a new element whenever the observed value on the subject changes.
     struct KeyValueSequence<Subject, Value> where Subject: NSObject, Value: Sendable {
-        // This is @unchecked Sendable since KeyPath isn't marked as Sendable but
-        // _probably_ should be.
-        public enum KeyPathType: @unchecked Sendable {
-            case `static`(KeyPath<Subject, Value>)
-            case stringy(String)
+        public let subject: Subject
+        public let keyPath: KeyPath<Subject, Value>
+        public let options: NSKeyValueObservingOptions
+
+        public init(
+            subject: Subject,
+            keyPath: KeyPath<Subject, Value>,
+            options: NSKeyValueObservingOptions
+        ) {
+            self.subject = subject
+            self.keyPath = keyPath
+            self.options = options
         }
+    }
 
+    /// An `AsyncSequence` that produces a new element whenever the observed value on the subject changes.
+    struct StringyKeyValueSequence<Subject, Value> where Subject: NSObject, Value: Sendable {
         public let subject: Subject
-        public let keyPath: KeyPathType
+        public let keyPath: String
         public let options: NSKeyValueObservingOptions
 
-        public init(subject: Subject, keyPath: KeyPathType, options: NSKeyValueObservingOptions) {
+        public init(subject: Subject, keyPath: String, options: NSKeyValueObservingOptions) {
             self.subject = subject
             self.keyPath = keyPath
             self.options = options
@@ -27,43 +37,73 @@ public extension NSObjectProtocol {
         for keyPath: KeyPath<Self, Value>,
         options: NSKeyValueObservingOptions = [.initial, .new]
     ) -> NSObject.KeyValueSequence<Self, Value> {
-        NSObject.KeyValueSequence(subject: self, keyPath: .static(keyPath), options: options)
+        NSObject.KeyValueSequence(subject: self, keyPath: keyPath, options: options)
     }
 
     func sequence<Value>(
         of value: Value.Type = Value.self,
         forKeyPath keyPath: String,
         options: NSKeyValueObservingOptions = [.initial, .new]
-    ) -> NSObject.KeyValueSequence<Self, Value> {
-        NSObject.KeyValueSequence(subject: self, keyPath: .stringy(keyPath), options: options)
+    ) -> NSObject.StringyKeyValueSequence<Self, Value> {
+        NSObject.StringyKeyValueSequence(subject: self, keyPath: keyPath, options: options)
     }
 }
 
 // MARK: - AsyncSequence
 
 extension NSObject.KeyValueSequence: AsyncSequence {
+    public typealias Element = Value
+
+    public struct AsyncIterator: AsyncIteratorProtocol {
+        var storage: KeyPathStorage<Subject, Value>
+
+        public func next() async -> Value? {
+            await storage.next()
+        }
+    }
+
+    public func makeAsyncIterator() -> AsyncIterator {
+        let storage = KeyPathStorage<Subject, Value>(
+            subject: subject,
+            keyPath: keyPath,
+            options: options
+        )
+        return AsyncIterator(storage: storage)
+    }
+}
+
+extension NSObject.StringyKeyValueSequence: AsyncSequence {
     public struct Element {
         public var newValue: Value?
         public var oldValue: Value?
     }
 
     public struct AsyncIterator: AsyncIteratorProtocol {
-        var storage: Storage<Subject, Value>
+        var storage: StringyKeyPathStorage<Subject, Value>
 
         public func next() async -> Element? {
             await storage.next()
         }
     }
 
     public func makeAsyncIterator() -> AsyncIterator {
-        let storage = Storage(subject: subject, keyPath: keyPath, options: options)
+        let storage = StringyKeyPathStorage<Subject, Value>(
+            subject: subject,
+            keyPath: keyPath,
+            options: options
+        )
         return AsyncIterator(storage: storage)
     }
 }
 
 // MARK: - Sendable
 
-extension NSObject.KeyValueSequence: Sendable where Subject: Sendable, Value: Sendable {}
+extension NSObject.KeyValueSequence: @unchecked Sendable where Subject: Sendable, Value: Sendable {}
+
+extension NSObject.StringyKeyValueSequence: Sendable where Subject: Sendable, Value: Sendable {}
 
 @available(*, unavailable)
 extension NSObject.KeyValueSequence.AsyncIterator: Sendable {}
+
+@available(*, unavailable)
+extension NSObject.StringyKeyValueSequence.AsyncIterator: Sendable {}

Sources/KVOSequence/Storage.swift

@@ -1,13 +1,12 @@
 import Foundation
 
-final class Storage<Subject, Value>: @unchecked Sendable where Subject: NSObject, Value: Sendable {
-    typealias Element = NSObject.KeyValueSequence<Subject, Value>.Element
-    typealias KeyPathType = NSObject.KeyValueSequence<Subject, Value>.KeyPathType
+final class KeyPathStorage<Subject, Value>: @unchecked Sendable where Subject: NSObject, Value: Sendable {
+    typealias Element = Value
 
     private let subject: Subject
     private let stateMachine: ManagedCriticalState<StateMachine<Subject, Element>>
 
-    init(subject: Subject, keyPath: KeyPathType, options: NSKeyValueObservingOptions) {
+    init(subject: Subject, keyPath: KeyPath<Subject, Value>, options: NSKeyValueObservingOptions) {
         self.subject = subject
         self.stateMachine = ManagedCriticalState(StateMachine())
 
@@ -56,30 +55,104 @@ final class Storage<Subject, Value>: @unchecked Sendable where Subject: NSObject
         }
     }
 
-    private func setupObservation(on keyPath: KeyPathType, options: NSKeyValueObservingOptions) {
-        let token: AnyObject
-
-        switch keyPath {
-        case .static(let keyPath):
-            token = subject.observe(
-                keyPath,
-                options: options,
-                changeHandler: { [weak self] _, change in
-                    let element = Element(newValue: change.newValue, oldValue: change.oldValue)
-                    self?.elementProduced(element)
+    private func setupObservation(
+        on keyPath: KeyPath<Subject, Value>,
+        options: NSKeyValueObservingOptions
+    ) {
+        let token = subject.observe(
+            keyPath,
+            options: options,
+            changeHandler: { [weak self] (_, change: NSKeyValueObservedChange<Value>) in
+                if let value = change.newValue {
+                    self?.elementProduced(value)
                 }
-            )
-        case .stringy(let keyPath):
-            token = StringyObserver<Subject, Value>(
-                subject: subject,
-                keyPath: keyPath,
-                options: options,
-                changeHandler: { [weak self] change in
-                    let element = Element(newValue: change.newValue, oldValue: change.oldValue)
-                    self?.elementProduced(element)
+            }
+        )
+
+        stateMachine.withCriticalRegion { stateMachine in
+            stateMachine.observationCreated(token)
+        }
+    }
+
+    private func elementProduced(_ element: Element) {
+        let action = stateMachine.withCriticalRegion { stateMachine in
+            stateMachine.valueProduced(element)
+        }
+
+        switch action {
+        case .resumeContinuation(let continuation):
+            continuation.resume(returning: element)
+        case .none:
+            break
+        }
+    }
+}
+
+final class StringyKeyPathStorage<Subject, Value>: @unchecked Sendable where Subject: NSObject, Value: Sendable {
+    typealias Element = NSObject.StringyKeyValueSequence<Subject, Value>.Element
+
+    private let subject: Subject
+    private let stateMachine: ManagedCriticalState<StateMachine<Subject, Element>>
+
+    init(subject: Subject, keyPath: String, options: NSKeyValueObservingOptions) {
+        self.subject = subject
+        self.stateMachine = ManagedCriticalState(StateMachine())
+
+        // In order not to miss KVO events, set up the underlying observation immediately.
+        // Any events received before the first call to next() will be buffered using
+        // the state machine.
+        setupObservation(on: keyPath, options: options)
+    }
+
+    func next() async -> Element? {
+        return await withTaskCancellationHandler {
+            let action = self.stateMachine.withCriticalRegion { stateMachine in
+                stateMachine.next()
+            }
+
+            switch action {
+            case .returnValue(let value):
+                return value
+            case .none:
+                break
+            }
+
+            return await withUnsafeContinuation { continuation in
+                let action = self.stateMachine.withCriticalRegion { stateMachine in
+                    stateMachine.nextSuspended(continuation)
+                }
+
+                switch action {
+                case .resumeConsumer(let value):
+                    continuation.resume(returning: value)
+                case .none:
+                    break
                 }
-            )
+            }
+        } onCancel: {
+            let action = self.stateMachine.withCriticalRegion { stateMachine in
+                stateMachine.finish()
+            }
+
+            switch action {
+            case .resumeConsumer(let consumer):
+                consumer.resume(returning: nil)
+            case .none:
+                break
+            }
         }
+    }
+
+    private func setupObservation(on keyPath: String, options: NSKeyValueObservingOptions) {
+        let token = StringyObserver<Subject, Value>(
+            subject: subject,
+            keyPath: keyPath,
+            options: options,
+            changeHandler: { [weak self] change in
+                let element = Element(newValue: change.newValue, oldValue: change.oldValue)
+                self?.elementProduced(element)
+            }
+        )
 
         stateMachine.withCriticalRegion { stateMachine in
             stateMachine.observationCreated(token)

Tests/KVOSequenceTests/StaticKeyPathSequenceTests.swift

@@ -12,26 +12,23 @@ final class StaticKeyPathSequenceTests: XCTestCase {
     }
 
     func testProducesValues() async throws {
-        let sequence = person.sequence(for: \.name, options: [.old, .new])
+        let sequence = person.sequence(for: \.name)
         let iterator = sequence.makeAsyncIterator()
 
         person.name = "Hika"
         var element = await iterator.next()
-        XCTAssertEqual(element?.newValue, "Hika")
+        XCTAssertEqual(element, "Hika")
 
         person.name = "Mai"
         element = await iterator.next()
-        XCTAssertEqual(element?.newValue, "Mai")
-        XCTAssertEqual(element?.oldValue, "Hika")
+        XCTAssertEqual(element, "Mai")
 
         person.name = nil
         element = await iterator.next()
 
         // Unfortunate, but double optional is messy...
-        if case let value?? = element?.newValue {
+        if case let value?? = element {
             XCTFail("Expected newValue to be nil, but was: \(value)")
         }
-
-        XCTAssertEqual(element?.oldValue, "Mai")
     }
 }