Swift sequences: The art of being lazy

When creating lists and sequences of objects in Swift, most of the time we use the Array data structure. While arrays have a huge benefit in being easy to use, and something more or less every programmer on the planet knows about, they do require you to create all of your elements up front.

When dealing with smaller datasets, where each member isn’t very expensive to construct, this is not a problem. However, when this is not true, you can get some pretty big performance benefits from implementing your own lazily evaluated sequence, instead of using an array.

Let’s say we want to load a sequence of models from a local database:

class ModelLoader {
    func loadAllModels() -> [Model] { … }
}

Potentially our database could contain a large set of records, and for each model we need to hit the disk to actually load its data, so we don’t want to load everything at once. To make this happen, we’re going to replace the array return type with our own custom sequence.

Thanks to Swift’s protocol-oriented nature, defining your own sequences is quite easy. It also — thanks to protocol extensions — gives you access to all the APIs that you can use on the standard library-provided sequences like Array, Set or Dictionary, without having to write any code for it.

We start by creating a struct for our sequence, and make it conform to the Sequence protocol:

struct ModelSequence: Sequence {
    func makeIterator() -> ModelIterator {
        return ModelIterator()
    }
}

As you can see, all we need to do to conform to Sequence, is to be able to act as a factory for creating iterators. An iterator is what Swift actually uses to iterate over our sequence, like in a for-loop or forEach() call.

For our iterator, we’re going to keep loading a model from disk, until one couldn’t be found anymore, in which case we’ll return nil. Returning nil from an iterator’s next() method signals to Swift that the sequence has come to an end and the iteration will stop.

We initialize our iterator with a Database that we’ll use to load each model. We default to using a shared instance of the database, but to facilitate testing, we enable dependency injection of it as well.

struct ModelIterator: IteratorProtocol {
    private let database: Database
    private var index = 0

    init(database: Database = .shared) {
        self.database = database
    }

    mutating func next() -> Model? {
        let model = database.model(at: index)
        index += 1
        return model
    }
}

That’s it! Now we have a lazily evaluated sequence, that loads each model ad-hoc when it’s needed 🎉. We can now easily use our sequence whenever we want to iterate over all models in our database. For example, we can now search our database without having to load all of its records up front:

func searchForModel(matching query: String) -> Model? {
    for model in ModelSequence() {
        if model.title.contains(query) {
            return model
        }
    }

    return nil
}

The nice thing about the code above, is that as soon as we’ve found a match, we can simply exit out of the iteration by returning, preventing further database records from being loaded.

OK, time for the bonus round! As we’ve just seen, implementing your own custom sequences and iterators in Swift is quite easy. But for when you are really lazy, the standard library’s AnySequence type has a closure-based API that you can use to quickly implement simple sequences, like this:

class ModelLoader {
    func loadAllModels() -> AnySequence {
        return AnySequence { () -> AnyIterator in
            var index = 0

            return AnyIterator {
                let model = database.model(at: index)
                index += 1
                return model
            }
        }
    }
}

I’ve personally started to use custom sequences in a lot of different situations in my code. I find it not only more performant in cases like above, but also easier to debug as you can simply step through your iteration code to find any issues.

What do you think? Do you find the ability to define sequences useful in Swift?

Feel free to reach out to me on Twitter if you have any questions, suggestions or feedback. I’d also love to hear from you if you have any topic that you’d like me to cover in an upcoming post.

Thanks for reading 🚀

Improving Swift compile times

Type erasure using closures in Swift