Async Zipped

AsyncZipped.swift

/*
 
 Make sure to compile this with the following flags:
 
 -Xfrontend -warn-concurrency -Xfrontend -enable-actor-data-race-checks
 
*/

extension AsyncIteratorProtocol {
    func newAndNext() async throws -> (Self, Element)? {
        var copy = self
        if let n = try await copy.next() {
            return (copy, n)
        } else {
            return nil
        }
    }
}

struct AsyncZipped<S1: AsyncSequence, S2: AsyncSequence>: AsyncSequence
   where S1: Sendable, S2: Sendable, S1.AsyncIterator: Sendable, S2.AsyncIterator: Sendable
{
    var left: S1
    var right: S2
    
    typealias Element = (S1.Element, S2.Element)
    
    func makeAsyncIterator() -> AsyncIterator {
        AsyncIterator(left: left.makeAsyncIterator(), right: right.makeAsyncIterator())
    }
    
    struct AsyncIterator: AsyncIteratorProtocol {
        var left: S1.AsyncIterator
        var right: S2.AsyncIterator
        
        mutating func next() async throws -> Element? {
            let l0 = left
            let r0 = right
            async let x = try await l0.newAndNext()
            async let y = try await r0.newAndNext()
            switch try await (x,y) {
            case let ((newLeft, l)?, (newRight, r)?):
                left = newLeft
                right = newRight
                return (l, r)
            default:
                return nil
            }
        }
    }
}

extension AsyncSequence where Self: Sendable, AsyncIterator: Sendable {
    func zip<Other: AsyncSequence>(_ other: Other) -> AsyncZipped<Self, Other>
    where Other: Sendable, Other.AsyncIterator: Sendable
    {
        AsyncZipped(left: self, right: other)
    }
}

You can't mutate the iterators inside a child task. This is my workaround, if you have a better idea please let me know!

newAndNext btw needs to return the iterator even in the nil and thrown cases since the iterator may have state that indicates it is at the end of an iteration.

@phausler I guess the current approach works as long as consumers never call next again, but of course, that's not a valid assumption. Will change it. I guess an even better approach is to somehow set a done flag in the combining iterator, just so that we don't unnecessarily consume the other iterator.

@phausler AsyncIterator does need to keep track of when one iterator reaches the end so that it doesn't needlessly retrieve the next elements from the other iterator, but this doesn't require changing newAndNext:

mutating func next() async throws -> Element? {
    
    if reachedEnd {
        return nil
    }

    let l0 = left
    let r0 = right
    async let x = try await l0.newAndNext()
    async let y = try await r0.newAndNext()
    switch try await (x,y) {
        case let ((newLeft, l)?, (newRight, r)?):
            left = newLeft
            right = newRight
            return (l, r)
        default:
            reachedEnd = true
            return nil
    }
}