KVO in swift concurrency

AsyncKVOSequence.swift

import Foundation

struct AsyncKVOSequence<Object: NSObject, Property>: AsyncSequence {
    typealias Element = NSKeyValueObservedChange<Property>
    let storage: Storage
    init(_ object: Object,
         _ keyPath: KeyPath<Object, Property>,
         _ options: NSKeyValueObservingOptions = []) {
        self.storage = Storage(
            object: object,
            keyPath: keyPath,
            options: options)
    }

    func makeAsyncIterator() -> AsyncIterator {
        let (stream, continuation) = AsyncStream.makeStream(of: Element.self, bufferingPolicy: .unbounded)
        storage.startObservation(with: continuation)
        var iterator = stream.makeAsyncIterator()
        return .init { await iterator.next() }
    }

    struct AsyncIterator: AsyncIteratorProtocol {
        let nextElement: () async -> Element?

        mutating func next() async -> Element? {
            await nextElement()
        }
    }

    final class Storage {
        weak var object: Object?
        let keyPath: KeyPath<Object, Property>
        let options: NSKeyValueObservingOptions

        init(object: Object, keyPath: KeyPath<Object, Property>, options: NSKeyValueObservingOptions) {
            self.object = object
            self.keyPath = keyPath
            self.options = options
        }

        func startObservation(with continuation: AsyncStream<Element>.Continuation) {
            guard let object else {
                continuation.finish()
                return
            }
            DeallocWatcher.of(object).append {
                continuation.finish()
            }
            let observation = object.observe(keyPath, options: options) { object, change in
                continuation.yield(change)
            }
            continuation.onTermination = { reason in
                guard case .cancelled = reason else { return }
                observation.invalidate()
            }
        }
    }
}

private final class DeallocWatcher {
    private var callbacks: [@Sendable () -> Void]

    private init(callbacks: [@Sendable () -> Void] = []) {
        self.callbacks = callbacks
    }

    func append(callback: @escaping @Sendable () -> Void) {
        self.callbacks.append(callback)
    }

    deinit {
        for callback in callbacks { callback() }
    }

    private static var associatedKey: UInt8 = 0

    static func of(_ object: NSObject) -> DeallocWatcher {
        objc_sync_enter(object)
        defer { objc_sync_exit(object) }

        if let watcher = objc_getAssociatedObject(object, &associatedKey) as? DeallocWatcher {
            return watcher
        }

        let watcher = DeallocWatcher()
        objc_setAssociatedObject(object, &associatedKey, watcher, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
        return watcher
    }
}