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)
    }
}

@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
    }
}