AsyncThrowingStream instead of custom type

KMPNativeCoroutinesAsync/AsyncSequence.swift

@@ -6,103 +6,70 @@
 //
 
 import Dispatch
+import Foundation
 import KMPNativeCoroutinesCore
 
-/// Wraps the `NativeFlow` in a `NativeFlowAsyncSequence`.
-/// - Parameter nativeFlow: The native flow to collect.
-/// - Returns: A `NativeFlowAsyncSequence` that yields the collected values.
+/// Converts a NativeFlow to an ``AsyncThrowingStream``
 public func asyncSequence<Output, Failure: Error, Unit>(
-    for nativeFlow: @escaping NativeFlow<Output, Failure, Unit>
-) -> NativeFlowAsyncSequence<Output, Error, Unit> {
-    return NativeFlowAsyncSequence(nativeFlow: nativeFlow)
-}
-
-public struct NativeFlowAsyncSequence<Output, Failure: Error, Unit>: AsyncSequence {
-    public typealias Element = Output
+  for nativeFlow: @escaping NativeFlow<Output, Failure, Unit>
+) -> AsyncThrowingStream<Output, Error> {
+  AsyncThrowingStream { continuation in
+    let lock = NSRecursiveLock()
 
-    var nativeFlow: NativeFlow<Output, Failure, Unit>
-
-    public class Iterator: AsyncIteratorProtocol, @unchecked Sendable {
-
-        private let semaphore = DispatchSemaphore(value: 1)
-        private var nativeCancellable: NativeCancellable<Unit>?
-        private var item: (Output, () -> Unit)? = nil
-        private var result: Failure?? = Optional.none
-        private var cancellationError: Failure? = nil
-        private var continuation: UnsafeContinuation<Output?, Error>? = nil
-
-        init(nativeFlow: NativeFlow<Output, Failure, Unit>) {
-            nativeCancellable = nativeFlow({ item, next, unit in
-                self.semaphore.wait()
-                defer { self.semaphore.signal() }
-                if let continuation = self.continuation {
-                    continuation.resume(returning: item)
-                    self.continuation = nil
-                    return next()
-                } else {
-                    self.item = (item, next)
-                    return unit
-                }
-            }, { error, unit in
-                self.semaphore.wait()
-                defer { self.semaphore.signal() }
-                self.result = Optional.some(error)
-                if let continuation = self.continuation {
-                    if let error = error {
-                        continuation.resume(throwing: error)
-                    } else {
-                        continuation.resume(returning: nil)
-                    }
-                    self.continuation = nil
-                }
-                self.nativeCancellable = nil
-                return unit
-            }, { cancellationError, unit in
-                self.semaphore.wait()
-                defer { self.semaphore.signal() }
-                self.cancellationError = cancellationError
-                if let continuation = self.continuation {
-                    continuation.resume(returning: nil)
-                    self.continuation = nil
-                }
-                self.nativeCancellable = nil
-                return unit
-            })
+    let callbacks = FlowCallbacks<Output, Failure, Unit>(
+      onItem: { item, next, unit in
+        lock.lock()
+        defer { lock.unlock() }
+        guard !Task.isCancelled else {
+          continuation.finish(throwing: CancellationError())
+          return unit
         }
-
-        public func next() async throws -> Output? {
-            return try await withTaskCancellationHandler {
-                try await withUnsafeThrowingContinuation { continuation in
-                    self.semaphore.wait()
-                    defer { self.semaphore.signal() }
-                    if let (item, next) = self.item {
-                        continuation.resume(returning: item)
-                        _ = next()
-                        self.item = nil
-                    } else if let result = self.result {
-                        if let error = result {
-                            continuation.resume(throwing: error)
-                        } else {
-                            continuation.resume(returning: nil)
-                        }
-                    } else if self.cancellationError != nil {
-                        continuation.resume(throwing: CancellationError())
-                    } else {
-                        guard self.continuation == nil else {
-                            fatalError("Concurrent calls to next aren't supported")
-                        }
-                        self.continuation = continuation
-                    }
-                }
-            } onCancel: {
-                _ = nativeCancellable?()
-                nativeCancellable = nil
-            }
+        continuation.yield(item)
+        return next()
+      },
+      onComplete: { error, unit in
+        lock.lock()
+        defer { lock.unlock() }
+        guard !Task.isCancelled else {
+          continuation.finish(throwing: CancellationError())
+          return unit
         }
-    }
+        if let error = error {
+          continuation.finish(throwing: error)
+        } else {
+          continuation.finish()
+        }
+        return unit
+      },
+      onCancelled: { error, unit in
+        lock.lock()
+        defer { lock.unlock() }
+        guard !Task.isCancelled else {
+          continuation.finish(throwing: CancellationError())
+          return unit
+        }
+        continuation.finish(throwing: error)
+        return unit
+      }
+    )
 
-    public func makeAsyncIterator() -> Iterator {
-        return Iterator(nativeFlow: nativeFlow)
+    let cancellable = nativeFlow(
+      callbacks.onItem,
+      callbacks.onComplete,
+      callbacks.onCancelled
+    )
+
+    continuation.onTermination = { @Sendable _ in
+      lock.lock()
+      defer { lock.unlock() }
+      _ = cancellable()
     }
+  }
 }
 
+/// A struct containing the callbacks for a native Kotlin Flow.
+struct FlowCallbacks<Output, Failure: Error, Unit> {
+  let onItem: NativeCallback2<Output, () -> Unit, Unit>
+  let onComplete: NativeCallback<Failure?, Unit>
+  let onCancelled: NativeCallback<Failure, Unit>
+}

Package.resolved

@@ -0,0 +1,16 @@
+{
+  "object": {
+    "pins": [
+      {
+        "package": "RxSwift",
+        "repositoryURL": "https://github.com/ReactiveX/RxSwift.git",
+        "state": {
+          "branch": null,
+          "revision": "5dd1907d64f0d36f158f61a466bab75067224893",
+          "version": "6.9.0"
+        }
+      }
+    ]
+  },
+  "version": 1
+}