Explain Swift Iterators

Question

There’s very little up-to-date guidance on how to make generators in Swift (or iterators as they’re apparently called in Swift), especially if you are new to the language. Why are there so many generator types like AnyIterator and UnfoldSequence? Why doesn’t the following code, which should yield from a sequence of either individual Ints or Arrays of Ints, work?

func chain(_ segments: Any...) -> AnyIterator<Int>{
    return AnyIterator<Int> {
        for segment in segments {
            switch segment {
            case let segment as Int:
                return segment
            case let segment as [Int]:
                for i in segment {
                    return i
                }
            default:
                return nil
            }
        }
        return nil
    }
}

let G = chain(array1, 42, array2)
while let g = G.next() {
    print(g)
}

The way I understand it, AnyIterator is supposed to take the closure in the {}s and turn it into the .next() method in the returned generator, but it doesn’t seem to be working. Or should I be using UnfoldSequence like in this question instead. I’m very confused.

Answer 1

Yes, the next() method of AnyIterator calls the given closure. And in your code, that closure returns the same first element on each call, because it does not remember what elements have been returned already.

If Swift had a yield statement like Python or C# then things would be easier: you could yield segment or yield i and are done.

But – unfortunately? – Swift has no yield statement, which means that the closure must explicitly manage some state in order to resume the iteration with the next element on each call.

One possibility would be to maintain two indices, one for the current segment and one for the current element within a segment if that is an array:

func chain(_ segments: Any...) -> AnyIterator<Int> {
    var currentSegment = 0 // index of current segment
    var currentElement = 0 // index of current element within current segment
    return AnyIterator<Int> {
        while currentSegment < segments.count {
            let next = segments[currentSegment]
            switch next {
            case let value as Int:
                currentSegment += 1
                return value
            case let segment as [Int]:
                if currentElement < segment.count {
                    let val = segment[currentElement]
                    currentElement += 1
                    return val
                }
                currentSegment += 1
                currentElement = 0
            default:
                return nil
            }
        }
        return nil
    }
}

This can be generalized to arbitrarily nested arrays:

func chain(_ segments: Any...) -> AnyIterator<Int> {
    var stack: [(Any, Int)] = [(segments, 0)]
    return AnyIterator<Int> {
        while let (next, idx) = stack.popLast() {
            switch next {
            case let value as Int:
                return value
            case let segments as [Any]:
                if idx < segments.count {
                    stack.append((segments, idx + 1))
                    stack.append((segments[idx], 0))
                }
            default:
                return nil
            }
        }
        return nil
    }
}

The still-to-be-processed arrays are on the stack together with their current index. The arrays themselves are not modified, so that the copy is cheap.

Example:

let G = chain([1, 2, [3]], 4, [5, 6, [], 7])
while let g = G.next() {
    print(g)
}
// 1 2 3 4 5 6 7

See also Implementing recursive generator for simple tree structure in Swift for more approaches to recursively enumerate a tree-like structure.