Is it possible to make an Iterator that yelds?

I'm wondering if now that we have a form of coroutines in the language (with async/await) if there is a way to implement a Collection Iterator (sync) that can just yeld on every next call instead of having to keep internal state.

So we could implement for example an iterator that goes over a nested array like this:

AnyIterator {
            for y in self.indices {
                for x in self[y].indices {
                    yield Point(x: x, y: y)
                }
            }
}

I know this doesn't work out of the box but I'm wondering if we can vend Task.yield() to do it.

1 Like

you can try AsyncStream or AsyncThrowingStream

But that's for async sequences. I want just a normal sync iterator to use in a normal for in loop. It's just that on the implementation side I would like to just yield instead of having to keep all the internal state myself.

very quick & dirty:

var taskResult: Int?
var taskDone = false

func foo() {
    Task {
        for i in 0 ..< 10 {
            taskResult = i
            await Task.yield()
        }
        taskDone = true
    }
}

func bar() {
    Task {
        print("loop start")
        while !taskDone {
            if let result = taskResult {
                print("Generator result \(result)")
                taskResult = nil
            }
            await Task.yield()
        }
        print("loop end")
    }
}

func test() {
    foo()
    bar()
}

output:

loop start
Generator result 0
Generator result 1
Generator result 2
Generator result 3
Generator result 4
Generator result 5
Generator result 6
Generator result 7
Generator result 8
Generator result 9
loop end

That's not really an Iterator tho. You can't provide an Iterator that Swift understand and can be used in a for loop with that code. Because the next() protocol requirement is synchronous of course.

I don't want to confuse anybody, I just mentioned the Task.yield() because is the only user level API we have that uses that word. But I guess what I'm looking for is to be able to use coroutines/generators without using the async machinery.

Synchronous coroutines are, I believe, being planned for in a "whenever we can find the time for it" sort of way. If I remember correctly, that's why _read & _modify use yield to give up their storage. Sadly, I don't think there are any immediate plans to add them in.

You could hack together something like this to "simulate synchronous coroutines / generators" (though it requires Dispatch for the semaphore, and the finished variable should be made an atomic variable). Also there is no way to cancel the iterator.

import Dispatch

class Generator<T>: @unchecked Sendable {
    let anyGenerator: AnyGenerator<T>

    init(_ anyGenerator: AnyGenerator<T>) {
        self.anyGenerator = anyGenerator
    }

    func yield(_ value: T) async {
        anyGenerator.nextValue = value
        await withUnsafeContinuation { (continuation: UnsafeContinuation<Void, Never>) in
            anyGenerator.continuation = continuation
            anyGenerator.semaphore.signal()
        }
    }
}

@globalActor
actor AnyGeneratorActor {
    static let shared: AnyGeneratorActor = AnyGeneratorActor()
}

final class AnyGenerator<T>: Sequence {
    let semaphore = DispatchSemaphore(value: 0)
    var continuation: UnsafeContinuation<Void, Never>?
    var nextValue: T?
    var finished = false

    public init(_ block: @AnyGeneratorActor @Sendable @escaping (Generator<T>) async -> Void) {
        Task.detached { [self] in
            let generator = Generator(self)
            await block(generator)
            nextValue = nil
            finished = true
            continuation = nil
            semaphore.signal()
        }
    }

    func makeIterator() -> some IteratorProtocol {
        return Iterator(self)
    }

    public class Iterator<T>: @unchecked Sendable, IteratorProtocol {
        private let anyGenerator: AnyGenerator<T>

        init(_ anyGenerator: AnyGenerator<T>) {
            self.anyGenerator = anyGenerator
        }

        public func next() -> T? {
            if anyGenerator.finished { 
                return nil
            }
            anyGenerator.semaphore.wait()
            let returnValue = anyGenerator.nextValue
            anyGenerator.continuation?.resume()
            return returnValue
        }
    }
}

It can be used like this in synchronous code

let generator = AnyGenerator<(Int, Int)> { generator in
    for x in 0..<10 {
        for y in 0..<10 {
            await generator.yield((x,y))
        }
    }
}

for point in generator {
    print(point)
}
1 Like