Iterators and lazy sequences

Every sequence so far has been fully built before anyone used it: an array of all its elements, in memory. CS61A’s Lazy Evaluation and Generators lectures ask what happens if elements are only computed when someone asks for them. The answer includes sequences with no end.

Iterators

An iterator hands out the elements of a sequence one at a time. Its one method, next(), returns the next element, or nil when there are no more:

var iterator = [10, 20].makeIterator()
print(iterator.next()!)          // 10
print(iterator.next()!)          // 20
print(iterator.next() as Any)    // nil

An iterator remembers where it is, so each call changes it; that is why it is a var. A for-in loop is shorthand for exactly this: make an iterator, call next() until it returns nil.

Writing your own

A type becomes a sequence by conforming to the Sequence protocol, which the previous lesson introduced. The simplest way is to be your own iterator, conforming to IteratorProtocol as well:

struct Countdown: Sequence, IteratorProtocol {
    var current: Int

    mutating func next() -> Int? {
        guard current > 0 else {
            return nil
        }
        defer { current -= 1 }
        return current
    }
}

for n in Countdown(current: 3) {
    print(n)    // 3, 2, 1
}
print(Countdown(current: 5).map { $0 * 10 })   // [50, 40, 30, 20, 10]

defer runs its block just before the function returns, so next() returns the current value and then counts down. Once Countdown is a Sequence, it gets map, filter, contains, and everything else for free.

Generators from closures

Python’s generators are functions that yield values one at a time. Swift has no yield for this, but a closure that captures its state does the same job. AnyIterator wraps such a closure as an iterator:

func naturals() -> AnyIterator<Int> {
    var n = 0
    return AnyIterator {
        n += 1
        return n
    }
}

let numbers = naturals()
print(numbers.next()!, numbers.next()!, numbers.next()!)   // 1 2 3

This is Lesson 5’s counter, returning its values through the iterator protocol. The sequence never ends: next() never returns nil. That is fine, as long as nobody asks for all of it.

For the common case of “start here, compute the next from the previous”, the standard library has sequence(first:next:). It ends when the closure returns nil:

let powers = sequence(first: 1) { $0 * 2 }
print(Array(powers.prefix(5)))   // [1, 2, 4, 8, 16]

let hailstone = sequence(first: 6) { n in
    n == 1 ? nil : (n % 2 == 0 ? n / 2 : 3 * n + 1)
}
print(Array(hailstone))

Lazy sequences

map and filter on an array are eager: they compute every element immediately and return a new array. Adding .lazy makes them compute each element only when it is needed:

let doubled = [1, 2, 3].lazy.map { x -> Int in
    print("map", x)
    return x * 2
}
print(doubled.first!)

Laziness is what makes infinite sequences useful. 1... is the range of all integers from 1 upward; filtering it eagerly would never finish, but a lazy filter can be asked for just the first few matches:

func isPrime(_ n: Int) -> Bool {
    n > 1 && !(2..<n).contains { $0 * $0 <= n && n % $0 == 0 }
}

let primes = (1...).lazy.filter(isPrime)
print(Array(primes.prefix(6)))   // [2, 3, 5, 7, 11, 13]

prefix(6) is lazy too; only Array(...) finally asks for elements, and it stops after six.

let evens = (0...).lazy.map { $0 * 2 }.filter { $0 % 3 != 0 }
print(Array(evens.prefix(4)))

struct Countdown: Sequence, IteratorProtocol {
    var current: Int
    mutating func next() -> Int? {
        if current == 0 {
            return nil
        }
        current -= 1
        return current + 1
    }
}
print(Array(Countdown(current: 3)))

Lab 16: Sequences on demand

Q1: A range with a step

Write a sequence Steps that yields start, start + step, … up to but not including end. Assume step is positive.

struct Steps: Sequence, IteratorProtocol {
    var current: Int
    let end: Int
    let step: Int

    init(from start: Int, to end: Int, by step: Int) {
        precondition(step > 0, "step must be positive")
        current = start
        self.end = end
        self.step = step
    }

    mutating func next() -> Int? {
        guard current < end else {
            return nil
        }
        defer { current += step }
        return current
    }
}

assert(Array(Steps(from: 0, to: 10, by: 3)) == [0, 3, 6, 9])
assert(Array(Steps(from: 5, to: 5, by: 1)) == [])
assert(Steps(from: 1, to: 100, by: 2).reduce(0, +) == 2500)

Q2: Fibonacci forever

Write fibonacci(), returning an AnyIterator<Int> that produces 0, 1, 1, 2, 3, 5, … without end.

func fibonacci() -> AnyIterator<Int> {
    // your code here
}

func fibonacci() -> AnyIterator<Int> {
    var current = 0
    var next = 1
    return AnyIterator {
        defer { (current, next) = (next, current + next) }
        return current
    }
}

let fibs = fibonacci()
assert((1...8).map { _ in fibs.next()! } == [0, 1, 1, 2, 3, 5, 8, 13])

Q3: The first perfect squares that are even

Using a lazy sequence of all positive integers, return the first n perfect squares that are even.

func evenSquares(_ n: Int) -> [Int] {
    // your code here
}

func evenSquares(_ n: Int) -> [Int] {
    Array((1...).lazy.map { $0 * $0 }.filter { $0 % 2 == 0 }.prefix(n))
}

assert(evenSquares(4) == [4, 16, 36, 64])
assert(evenSquares(0) == [])

What’s next

Next lesson: efficiency. How to count the work a program does, and how memoization turns Lesson 8’s exponential fib into a fast one.