How to make an iterable Swift collection type

Sometimes it’s useful to implement custom collection types. For example, you may want to keep the order of the iteration of the objects in the collection hidden from the consumer. Here’s how you can make an iterable Swift collection type with SequenceType.

Resources:

• Swift: How To Conform to the Sequence Protocol (Natasha The Robot)
• Generators in Swift (http://schani.wordpress.com/)
• Swift Sequences (Gordon Fontenot)

A parking lot

Consider the following class ParkingLot. It happens to be a last-in-first-out type parking lot. The implementation of the collection is an Array but it could have been a linked list, or a Dictionary. We want to be able to loop through the cars parked in any given lot using the for…in construct.

class ParkingLot {
   var cars = [Car]()

    func park(car: Car) {
        cars.append(car)
    }
    func unpark() -> Car {
        return cars.removeAtIndex(cars.lastIndex)
    }
}

// Loop through the cars
for car in aParkingLot {
   doSomethingWith(car: car)
}

To achieve this, we implement Swift’s SequenceType protocol on the ParkingLot class.

SequenceType and GeneratorType

The SequenceType protocol definition looks like this:

protocol SequenceType : _Sequence_Type {
    typealias Generator : GeneratorType
    func generate() -> Generator
}

The protocol requires that our class implements the generate method. It should return an object that conforms to the GeneratorType protocol. The GeneratorType protocol definition looks like this (comments removed):

protocol GeneratorType {
    typealias Element
    mutating func next() -> Element?
}

This, in turn, tells us that our GeneratorType needs to implement the next method, returning the next object in the iteration. Element is an alias for the actual type in our collection.

Implementing a GeneratorType may seem a little complicated at first, but here’s a walk-through:

class CarGenerator: GeneratorType {
    // Hold the collection to be iterated through
    private var cars: Array?

    // Hold the index of the next item in the iteration
    private var nextIndex: Int

    // Initialize the generator, passing a reference
    // to the car array
    init(cars: Array) {
        self.cars = cars

        // Set the nextIndex to the index of the
        // last car in the array
        // (since our collection is LIFO)
        nextIndex = cars.count-1
    }

    // Implement the next method, to return nil if
    // there are no more cars, or the next Car.
    // Note that we have specified the return type
    // as an optional "car": Car?.
    func next() -> Car? {

        if (nextIndex < 0) {
            return nil
        }

        // Decrement the index after it has been
        //evaluated
        return self.cars![nextIndex--]
    }
}

Now we can return an instance of CarGenerator in the generate method of ParkingLot:

class ParkingLot: SequenceType {
    private var cars = [Car]()

    func generate() -> CarGenerator {
        return CarGenerator(cars: cars)
    }

    // rest of class
}

Simplified version with GeneratorOf<T>

That seems like an awful lot of boiler-plate though? Actually, there is a simpler way that you can use in most cases. The generate method can return an instance of type GeneratorOf<T> rather than a custom GeneratorType.

Definition of GeneratorOf<T>:

struct GeneratorOf<T> : GeneratorType, SequenceType {
    init(_ next: () -> T?)
    init(_ self_: G)
    mutating func next() -> T?
    func generate() -> GeneratorOf<T>
}

The interesting part of this definition is the line that says init(_ next: () -> T?). What this essentially says it that you can initialize a GeneratorOf<T> by passing to its init method a closure that returns the next object in the iteration. We substitute T for the type we are generating – in our case Car.

Finally:

class ParkingLot: SequenceType {
    private var cars = [Car]()

    func generate() -> GeneratorOf<Car> {
        // keep the index of the next car in the iteration
        var nextIndex = cars.count-1

        // Construct a GeneratorOf<Car> instance,
        // passing a closure that returns the next
        // car in the iteration
        return GeneratorOf<Car> {
            if (nextIndex < 0) {
                return nil
            }
            return cars[nextIndex--]
        }
    }

    // rest of class
}

Notice that the fact that the return type of the closure is an optional is implied from the protocol. We can therefore safely return nil when we are out of cars to iterate through, without specifying it anywhere.