Add nonthrowing scalar flatMapLatest and lazy Just factories

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 - Just: release consumed values from iterators so nested `flatMapLatest` subscriptions do not retain removed feature objects (https://github.com/sideeffect-io/AsyncExtensions/issues/35).
 
+- FlatMapLatest: add nonthrowing asynchronous scalar transformations and lazy `AsyncJustSequence` factories; returned `Result.failure` values do not terminate collection (https://github.com/sideeffect-io/AsyncExtensions/issues/45).
+
 - Breaking: use Swift Async Algorithms for `Sequence.async` and two- or three-input `zip`/`merge`, removing import ambiguities. Add the `AsyncAlgorithms` product dependency and import when migrating these APIs. Preserve variadic `zip`/`merge` and the explicit `AsyncLazySequence` constructor.
 - Breaking: rename the buffered Date timer from `AsyncTimerSequence` to `AsyncBufferedTimerSequence` to avoid ambiguity with Apple's clock-based timer.
 - SwiftPM: require Swift 5.8 or later for the Swift Async Algorithms test dependency.

README.md

@@ -124,6 +124,24 @@ or replay consumer receives the latest stored state, followed by subsequent send
 
 More operators and extensions are to come. Pull requests are of course welcome.
 
+## Asynchronous single-value transformations
+
+`AsyncJustSequence(factory:)` creates its value lazily when an iterator first advances. The factory runs once per independent iterator; returning `nil` finishes without emitting a value. Cancellation before iteration skips the factory, and cancellation during the factory discards its eventual result.
+
+The scalar `flatMapLatest` overload accepts a nonthrowing async transformation. It cancels the previous transformation when a new input arrives. When both the upstream and transformation are nonthrowing, collection needs no `try`:
+
+```swift
+let doubled = AsyncJustSequence(21).flatMapLatest { value async -> Int in
+  value * 2
+}
+
+for await value in doubled {
+  print(value) // 42
+}
+```
+
+A transformation can return `Result` to represent a failure as an ordinary element and keep collecting later inputs. A throwing upstream still requires `try` and propagates its error; the nonthrowing transformation does not suppress upstream failures.
+
 ## Subscription lifetime
 
 An `AsyncJustSequence` iterator releases its stored value after emitting it. An inner sequence can continue producing values after the object that provided it is released.

Sources/Creators/AsyncJustSequence.swift

@@ -18,31 +18,50 @@ public struct AsyncJustSequence<Element>: AsyncSequence {
   public typealias Element = Element
   public typealias AsyncIterator = Iterator
 
-  let element: Element?
+  enum Source {
+    case value(Element?)
+    case factory(@Sendable () async -> Element?)
+  }
+
+  let source: Source
 
   public init(_ element: Element?) {
-    self.element = element
+    self.source = .value(element)
+  }
+
+  /// Creates a single value lazily for each iterator. Returning nil finishes without a value.
+  public init(factory: @Sendable @escaping () async -> Element?) {
+    self.source = .factory(factory)
   }
 
   public func makeAsyncIterator() -> AsyncIterator {
-    Iterator(element: self.element)
+    Iterator(source: self.source)
   }
 
   public struct Iterator: AsyncIteratorProtocol {
-    let pendingElement: ManagedCriticalState<Element?>
+    let pendingSource: ManagedCriticalState<Source?>
 
-    init(element: Element?) {
-      self.pendingElement = ManagedCriticalState(element)
+    init(source: Source) {
+      self.pendingSource = ManagedCriticalState(source)
     }
 
     public mutating func next() async -> Element? {
-      self.pendingElement.withCriticalRegion { pendingElement in
-        defer { pendingElement = nil }
-        return Task.isCancelled ? nil : pendingElement
+      let source = self.pendingSource.withCriticalRegion { pendingSource -> Source? in
+        defer { pendingSource = nil }
+        return pendingSource
+      }
+      guard !Task.isCancelled, let source = source else { return nil }
+      switch source {
+        case .value(let element):
+          return element
+        case .factory(let factory):
+          let element = await factory()
+          return Task.isCancelled ? nil : element
       }
     }
   }
 }
 
 extension AsyncJustSequence: Sendable where Element: Sendable {}
+extension AsyncJustSequence.Source: Sendable where Element: Sendable {}
 extension AsyncJustSequence.Iterator: Sendable where Element: Sendable {}

Sources/Operators/AsyncSequence+FlatMapLatest.swift

@@ -74,9 +74,21 @@ public extension AsyncSequence {
   ///
   /// - note: This operator is a combination of `map` and `switchToLatest`.
   /// - Returns: An async sequence emitting the value of the latest inner async sequence.
+  // A sequence-returning transform should use the flattening overload rather than emit a sequence.
+  @_disfavoredOverload
   func flatMapLatest<Output>(
     _ transform: @Sendable @escaping (Element) async throws -> Output
   ) -> AsyncSwitchToLatestSequence<AsyncThrowingMapSequence<Self, AsyncThrowingJustSequence<Output>>> where Element: Sendable {
     self.map { element in AsyncThrowingJustSequence { try await transform(element) } }.switchToLatest()
   }
+
+  /// Transforms each element asynchronously, cancelling the previous transformation when a new
+  /// element arrives. A nonthrowing upstream and transformation can be iterated without `try`.
+  /// Returning a `Result.failure` is an ordinary value and does not terminate the sequence.
+  @_disfavoredOverload
+  func flatMapLatest<Output>(
+    _ transform: @Sendable @escaping (Element) async -> Output
+  ) -> AsyncSwitchToLatestSequence<AsyncMapSequence<Self, AsyncJustSequence<Output>>> where Element: Sendable {
+    self.map { element in AsyncJustSequence<Output> { await transform(element) } }.switchToLatest()
+  }
 }

Tests/Creators/AsyncJustFactoryTests.swift

@@ -0,0 +1,93 @@
+@testable import AsyncExtensions
+import XCTest
+
+struct JustFactoryGate: Sendable {
+  struct State: Sendable {
+    var continuation: UnsafeContinuation<Void, Never>?
+    var opened = false
+  }
+  let state = ManagedCriticalState(State())
+  let suspended: XCTestExpectation
+
+  func wait() async {
+    await withUnsafeContinuation { continuation in
+      let opened = state.withCriticalRegion { state -> Bool in
+        guard !state.opened else { return true }
+        state.continuation = continuation
+        return false
+      }
+      if opened { continuation.resume() }
+      suspended.fulfill()
+    }
+  }
+
+  func open() {
+    let continuation = state.withCriticalRegion { state -> UnsafeContinuation<Void, Never>? in
+      state.opened = true
+      defer { state.continuation = nil }
+      return state.continuation
+    }
+    continuation?.resume()
+  }
+}
+
+final class AsyncJustFactoryTests: XCTestCase {
+  func test_factory_is_lazy_and_runs_once_per_independent_iterator() async {
+    let calls = ManagedCriticalState(0)
+    let sequence = AsyncJustSequence<Int> {
+      calls.withCriticalRegion { $0 += 1 }
+      return 42
+    }
+    XCTAssertEqual(calls.criticalState, 0)
+    var first = sequence.makeAsyncIterator()
+    var independent = sequence.makeAsyncIterator()
+    let firstValue = await first.next()
+    let pastEnd = await first.next()
+    let independentValue = await independent.next()
+    XCTAssertEqual(firstValue, 42)
+    XCTAssertNil(pastEnd)
+    XCTAssertEqual(independentValue, 42)
+    XCTAssertEqual(calls.criticalState, 2)
+  }
+
+  func test_nil_factory_result_finishes_without_emitting() async {
+    var iterator = AsyncJustSequence<Int> { nil }.makeAsyncIterator()
+    let value = await iterator.next()
+    XCTAssertNil(value)
+  }
+
+  func test_already_cancelled_iterator_does_not_invoke_factory() async {
+    let gate = JustFactoryGate(suspended: expectation(description: "Paused before iteration"))
+    let calls = ManagedCriticalState(0)
+    let task = Task {
+      await gate.wait()
+      var iterator = AsyncJustSequence<Int> {
+        calls.withCriticalRegion { $0 += 1 }
+        return 42
+      }.makeAsyncIterator()
+      return await iterator.next()
+    }
+    await fulfillment(of: [gate.suspended], timeout: 2)
+    task.cancel()
+    gate.open()
+    let value = await task.value
+    XCTAssertNil(value)
+    XCTAssertEqual(calls.criticalState, 0)
+  }
+
+  func test_cancellation_discards_late_noncooperative_factory_value() async {
+    let gate = JustFactoryGate(suspended: expectation(description: "Factory suspended"))
+    let task = Task {
+      var iterator = AsyncJustSequence<Int> {
+        await gate.wait()
+        return 42
+      }.makeAsyncIterator()
+      return await iterator.next()
+    }
+    await fulfillment(of: [gate.suspended], timeout: 2)
+    task.cancel()
+    gate.open()
+    let value = await task.value
+    XCTAssertNil(value)
+  }
+}

Tests/Operators/AsyncNonThrowingFlatMapLatestTests.swift

@@ -0,0 +1,88 @@
+@testable import AsyncExtensions
+import XCTest
+
+final class AsyncNonThrowingFlatMapLatestTests: XCTestCase {
+  func test_sequence_returning_transform_still_flattens_without_an_expected_type() async {
+    let sequence = AsyncJustSequence(21).flatMapLatest { value async -> AsyncJustSequence<Int> in
+      AsyncJustSequence(value * 2)
+    }
+    var values = [Int]()
+    for await value in sequence { values.append(value) }
+    XCTAssertEqual(values, [42])
+  }
+
+  func test_optional_scalar_nil_is_an_emitted_value() async {
+    let sequence = AsyncJustSequence(21).flatMapLatest { _ async -> Int? in nil }
+    var values = [Int?]()
+    for await value in sequence { values.append(value) }
+    XCTAssertEqual(values, [nil])
+  }
+
+  func test_new_input_cancels_previous_transform_and_discards_its_late_value() async {
+    let source = AsyncBufferedChannel<Int>()
+    let gate = JustFactoryGate(suspended: expectation(description: "First transformation suspended"))
+    let cancelled = expectation(description: "First transformation cancelled")
+    let receivedLatest = expectation(description: "Newest result received")
+    let sequence = source.flatMapLatest { value async -> Int in
+      if value == 1 {
+        await withTaskCancellationHandler {
+          await gate.wait()
+        } onCancel: {
+          cancelled.fulfill()
+        }
+      }
+      return value
+    }
+    let task = Task {
+      var values = [Int]()
+      for await value in sequence {
+        values.append(value)
+        receivedLatest.fulfill()
+      }
+      return values
+    }
+    source.send(1)
+    await fulfillment(of: [gate.suspended], timeout: 2)
+    source.send(2)
+    await fulfillment(of: [cancelled], timeout: 2)
+    gate.open()
+    await fulfillment(of: [receivedLatest], timeout: 2)
+    source.finish()
+    let values = await task.value
+    XCTAssertEqual(values, [2])
+  }
+
+  func test_throwing_upstream_still_propagates_its_error() async {
+    let sequence = AsyncFailSequence<Int>(MockError(code: 45)).flatMapLatest { value async -> Int in value }
+    do {
+      for try await _ in sequence { XCTFail("The upstream should fail") }
+      XCTFail("Expected the upstream error")
+    } catch {
+      XCTAssertEqual(error as? MockError, MockError(code: 45))
+    }
+  }
+
+  func test_nonthrowing_scalar_transform_can_be_iterated_without_try() async {
+    let sequence = AsyncJustSequence(21).flatMapLatest { value async -> Int in value * 2 }
+    var values = [Int]()
+    for await value in sequence { values.append(value) }
+    XCTAssertEqual(values, [42])
+  }
+
+  func test_nonthrowing_transform_can_return_results_without_terminating_the_sequence() async {
+    let source = AsyncBufferedChannel<Int>()
+    let sequence = source.flatMapLatest { value async -> Result<Int, MockError> in
+      value == 1 ? .failure(MockError(code: 1)) : .success(value)
+    }
+    var iterator = sequence.makeAsyncIterator()
+    source.send(1)
+    let failure = await iterator.next()
+    XCTAssertEqual(failure, .failure(MockError(code: 1)))
+    source.send(2)
+    let success = await iterator.next()
+    XCTAssertEqual(success, .success(2))
+    source.finish()
+    let finished = await iterator.next()
+    XCTAssertNil(finished)
+  }
+}