Implementing a custom asynchronous sequence in Swift

Imagine I want to create a function that, given an array of numbers, computes the square, cube, and fourth power of each number in an asynchronous fashion and returns a flattened, asynchronous sequence of all these results.

So, for example, for the input array [2, 3, 4], it should return an AsyncSequence instance yielding the elements [4, 8, 16, 9, 27, 81, 16, 64, 256].

Then let's say, instead of computing x^2, x^3, x^4, I would like it to compute x, x^2, x^3, ..., x^k where k is sort of a random integer that can be different for every x and is not known beforehand (its value comes to be known only as the powers are being computed). How would I implement such a pattern?

Thanks a lot to Rob for providing the basic idea on how to implement something like this.

I wrote it in the following way:

func powers(of numbers: [Int]) -> AsyncStream<Int> {
    return AsyncStream<Int> { continuation in
        Task {
            for number in numbers {
                for await power in Powers(of: number) {
                    continuation.yield(power)
                }
            }
            continuation.finish()
        }
    }
}

struct Powers: AsyncSequence {
    init(of base: Int) {
        self.base = base
    }
    
    func makeAsyncIterator() -> PowersIterator {
        return PowersIterator(base: self.base)
    }
    
    let base: Int
    typealias Element = Int
}

struct PowersIterator: AsyncIteratorProtocol {
    mutating func next() async -> Int? {
        if !self.shouldFinish() {
            try? await Task.sleep(nanoseconds: 1_000_000_000)
            defer {
                self.exponent += 1
            }
            return power(self.base, self.exponent)
        } else {
            return nil
        }
    }
    
    private func shouldFinish() -> Bool {
        return Int.random(in: 1...10) == 1
    }
    
    private func power(_ base: Int, _ exponent: UInt) -> Int {
        return (0..<exponent).reduce(1) { power, _ in power * base }
    }
    
    var exponent = UInt(1)
    let base: Int
}

It can be invoked using this code:

Task {
    let numbers = [1, 2, 3, 4, 5]
    for await power in powers(of: numbers) {
        print(power, terminator: " ")
    }
}

Possible output:

1 1 1 2 4 8 16 32 64 3 9 27 4 16 64 256 5 25 125 625 3125 15625

The solution is a little more complex than it ought to be. But that's of course because I actually wanted to compute something that would really need to be computed asynchronously and has the same computational structure as this example. That is also the reason for why I created a separate async sequence for computing the powers.

If this helps anyone out, I'll be glad.