merge: optimize task creation and remove sendable constraint

Author: twittemb

Package.swift

@@ -16,9 +16,14 @@ let package = Package(
     .library(name: "_CAsyncSequenceValidationSupport", type: .static, targets: ["AsyncSequenceValidation"]),
     .library(name: "AsyncAlgorithms_XCTest", targets: ["AsyncAlgorithms_XCTest"]),
   ],
-  dependencies: [],
+  dependencies: [
+    .package(url: "https://github.com/apple/swift-collections.git", .upToNextMajor(from: "1.0.3"))
+  ],
   targets: [
-    .target(name: "AsyncAlgorithms"),
+    .target(
+      name: "AsyncAlgorithms",
+      dependencies: [.product(name: "Collections", package: "swift-collections")]
+    ),
     .target(
       name: "AsyncSequenceValidation",
       dependencies: ["_CAsyncSequenceValidationSupport", "AsyncAlgorithms"]),

Sources/AsyncAlgorithms/AsyncChannel.swift

@@ -28,7 +28,7 @@ public final class AsyncChannel<Element: Sendable>: AsyncSequence, Sendable {
     init(_ channel: AsyncChannel<Element>) {
       self.channel = channel
     }
-    
+
     /// Await the next sent element or finish.
     public mutating func next() async -> Element? {
       guard active else {
@@ -116,7 +116,7 @@ public final class AsyncChannel<Element: Sendable>: AsyncSequence, Sendable {
 
   /// Create a new `AsyncChannel` given an element type.
   public init(element elementType: Element.Type = Element.self) { }
-  
+
   func establish() -> Int {
     state.withCriticalRegion { state in
       defer { state.generation &+= 1 }
@@ -152,7 +152,7 @@ public final class AsyncChannel<Element: Sendable>: AsyncSequence, Sendable {
           }
           return UnsafeResumption(continuation: send, success: continuation)
         case .awaiting(var nexts):
-          if nexts.update(with: Awaiting(generation: generation, continuation: continuation)) != nil {
+            if nexts.update(with: Awaiting(generation: generation, continuation: continuation)) != nil {
             nexts.remove(Awaiting(placeholder: generation))
             cancelled = true
           }

Sources/AsyncAlgorithms/AsyncChunksOfCountOrSignalSequence.swift

@@ -55,39 +55,89 @@ extension AsyncSequence {
   }
 }
 
+struct AsyncEitherSequence<Base: AsyncSequence>: AsyncSequence {
+  typealias Element = Either
+  typealias AsyncIterator = Iterator
+
+  enum Either {
+    case first(Base.Element)
+    case second(Base.Element)
+  }
+
+  let base: Base
+  let factory: (Base.Element) -> Either
+
+  func makeAsyncIterator() -> AsyncIterator {
+    Iterator(
+      base: self.base.makeAsyncIterator(),
+      factory: self.factory
+    )
+  }
+
+  struct Iterator: AsyncIteratorProtocol {
+
+    var base: Base.AsyncIterator
+    let factory: (Base.Element) -> Either
+
+    mutating func next() async rethrows -> Element? {
+      guard let element = try await self.base.next() else {
+        return nil
+      }
+      return self.factory(element)
+    }
+  }
+
+}
+
 /// An `AsyncSequence` that chunks elements into collected `RangeReplaceableCollection` instances by either count or a signal from another `AsyncSequence`.
 public struct AsyncChunksOfCountOrSignalSequence<Base: AsyncSequence, Collected: RangeReplaceableCollection, Signal: AsyncSequence>: AsyncSequence, Sendable where Collected.Element == Base.Element, Base: Sendable, Signal: Sendable, Base.AsyncIterator: Sendable, Signal.AsyncIterator: Sendable, Base.Element: Sendable, Signal.Element: Sendable {
 
   public typealias Element = Collected
 
+  enum Either {
+    case first(Base.Element)
+    case second(Signal.Element)
+  }
+
   /// The iterator for a `AsyncChunksOfCountOrSignalSequence` instance.
   public struct Iterator: AsyncIteratorProtocol, Sendable {
     let count: Int?
-    var state: Merge2StateMachine<Base, Signal>
-    init(base: Base.AsyncIterator, count: Int?, signal: Signal.AsyncIterator) {
+    let state: MergeStateMachine<Either>
+    init(base: Base, count: Int?, signal: Signal) {
       self.count = count
-      self.state = Merge2StateMachine(base, terminatesOnNil: true, signal)
+      let eitherBase = base.map { Either.first($0) }
+      let eitherSignal = signal.map { Either.second($0) }
+      self.state = MergeStateMachine(eitherBase, terminatesOnNil: true, eitherSignal)
     }
-    
+
     public mutating func next() async rethrows -> Collected? {
-      var result : Collected?
-      while let next = try await state.next() {
-        switch next {
-          case .first(let element):
-            if result == nil {
-              result = Collected()
-            }
-            result!.append(element)
-            if result?.count == count {
-              return result
-            }
-          case .second(_):
-            if result != nil {
-              return result
-            }
-        }
+      var collected: Collected?
+
+    loop: while true {
+      let next = await state.next()
+
+      switch next {
+        case .termination:
+          break loop
+        case .element(let result):
+          let element = try result._rethrowGet()
+          switch element {
+            case .first(let element):
+              if collected == nil {
+                collected = Collected()
+              }
+              collected!.append(element)
+              if collected?.count == count {
+                return collected
+              }
+            case .second(_):
+              if collected != nil {
+                return collected
+              }
+          }
       }
-      return result
+    }
+      return collected
     }
   }
 
@@ -105,6 +155,6 @@ public struct AsyncChunksOfCountOrSignalSequence<Base: AsyncSequence, Collected:
   }
 
   public func makeAsyncIterator() -> Iterator {
-    return Iterator(base: base.makeAsyncIterator(), count: count, signal: signal.makeAsyncIterator())
+    return Iterator(base: base, count: count, signal: signal)
   }
 }