Use Sequence and Collection extensions in Swift to keep your code cleaner

In this blog post I am going to show how you can use Sequence and Collection extension to keep your upper layer code cleaner and more readable. I attended try! Swift conference in Tokyo Japan at the beginning of March this year, and in my opinion one of the best talks there was Soroush Khanlou's talk with topic: "Everything you ever wanted to know about Sequence and Collection". That talk gave me an idea to blog about Sequence and Collection extensions and how to use them in your project. While Khanlou went really deep explaining how the Sequence and Collection works and the theory behind it, I'll try to keep it a bit more on the concrete example side. I'll briefly go through the theory part, but to get the whole picture I really encourage you to check the try! Swift talk.

What is a Sequence

Sequence is basically a list of elements, finite or infinite, that you can iterate through only once. Sometimes you can iterate it through more than once, but it is not guaranteed. Also, by conforming to the Sequence protocol you will get all the higher order functions to your struct such as map, flatMap, reduce etc. Sequence protocol looks like this:

protocol Sequence {
    associatedtype Iterator: IteratorProtocol
    func makeIterator() -> Iterator
}

Sequence has associatedtype Iterator which is the type of IteratorProtocol. In addition to that it also has makeIterator function to create the iterator, which returns the same associatedtype as defined in Iterator. Next, let's look at the Iterator protocol:

protocol Iterator {
    associatedtype Element
    mutating func next() -> Element?
}

Iterator protocol has an associated type Element and also a function called next which returns the Element. You can think of the Element type as a Generic (or template if you come from C++) as it defines the type of the elements that the Sequence holds. The next function always returns the next element in the Sequence unless the end of the Sequence is reached, in which case it returns a nil.

Sequence extension for common good

Next, let's play around a little bit and create a struct called Restaurant:

struct Restaurant {
    enum FoodType {
        case ramen
        case sushi
    }

    var name: String
    var openingTime: Date
    var closingTime: Date
    var type: FoodType
}

So we have a restaurant which has a name, opening and closing time and also a variable which tells what kind of food it serves. As you can see it can be either one of my favourites in Tokyo: ramen or sushi! Next, let's define some data that we can later play around with:

struct RestaurantHolder {
    static var restaurants: [Restaurant] {
        let date = Date()
        return [
            Restaurant(name: "Ichiraku Ramen",
                       openingTime: date.addingTimeInterval(-3600),
                       closingTime: date.addingTimeInterval(3600),
                       type: .ramen),
            Restaurant(name: "Fu-unji",
                       openingTime: date.addingTimeInterval(-3600),
                       closingTime: date.addingTimeInterval(3600),
                       type: .ramen),
            Restaurant(name: "Maka Sushi",
                       openingTime: date.addingTimeInterval(3600),
                       closingTime: date.addingTimeInterval(-3600),
                       type: .sushi),
            Restaurant(name: "Lempi Sushi",
                       openingTime: date.addingTimeInterval(-3600),
                       closingTime: date.addingTimeInterval(3600),
                       type: .sushi)
        ]
    }
}

So we defined a RestaurantHolder struct which has a static variable "restaurants" that holds our data. There are two ramen restaurants: "Ichiraku Ramen" and "Fu-unji" and two sushi places: "Maka Sushi" and "Lempi Sushi". The data also defines that all the other restaurants except for "Maka Sushi" are currently open. AddingTimeInterval takes seconds as a parameter so by the time we run the code "Maka Sushi" was closed an hour ago and will be open again after one hour.

Let's say that we have an application that lists restaurants and this is the data model that we are using. Now, we need to provide information to the user about how many restaurants are open at the moment. We can do it by using a filter like this:

let date = Date()
let openRestaurantsCount = RestaurantHolder.restaurants.filter { ($0.openingTime <= date && date < $0.closingTime) }.count
print(openRestaurantsCount) // 2

We get the restaurant objects from the array, filter them by defining that "date" is between the openingTime and the closingTime. When the resulted array is returned we call its count method. This is good but the problem is that we end up creating an array that we only use for a brief moment to call count to it and then toss it away.

Let's try the same thing by creating an extension for Sequence:

extension Sequence {
    func count(_ isValid: ((Iterator.Element) -> Bool)) -> Int {
        var validItems = 0
        for element in self {
            if isValid(element) {
                validItems += 1
            }
        }
        return validItems
    }
}

Inside the extension block we defined a function called "count". It takes an isValid function as a parameter and that functions parameter is defined as Iterator.Element. Now, if you recall the IteratorProtocol which we went through in the beginning of the post, you might remember that the Element works as a Generic in Swift. So basically the isValid function takes a condition as a parameter, and returns true or false. Sequence calls this function to all its elements and increases the validItems count for every element that passes. The count function eventually returns an Int which is the number passed items.

Sequence also implements the "for...in" loop so we can use it here. We go through the Sequence, take the elements inside and run them through the isValid function. If the element passes the condition we add +1 to the number of valid items. Eventually the validItems is returned and it holds the number of restaurants open at the moment. Now let's see how we can call this function:

let openRestaurantsCount2 = RestaurantHolder.restaurants.count { ($0.openingTime <= date && date < $0.closingTime) }
print(openRestaurantsCount2) // 2

This solution is better since we don't create an extra copy of the array and the function name is more definitive and tells us more about what it does. Now that we are using it only for counting the open restaurants the name could be even more specific (countRestaurantsCurrentlyOpen etc..), but the idea here is that count can be used to all Sequences through out the application. Credits for the count function goes to Khanlou since the code is from his presentation. But what if we wanted something specific only for the Sequence or Collection that holds restaurant information? Next, let's define an extension for a Collection that holds elements of a specific type.

Collection extension for specific type of elements

Collection is the protocol which Array, Set and Dictionary, which you might have used in your swift application, conform to. Let's continue with the same problem: we need to get the number of restaurants open at the moment. Since Collection conforms to the Sequence protocol we could also do this inside the Sequence extension, but we know that restaurants are stored inside a Collection so let's go with a Collection extension:

extension Collection where Iterator.Element == Restaurant {
    func countOpenRestaurants(at: Date) -> Int {
        var count = 0
        for element in self {
            if element.openingTime <= at && at < element.closingTime {
                count += 1
            }
        }
        return count
    }
}

Now when we look at the Collection extension we can see that the "Iterator.Element" is defined as "Restaurant". This means that the functions inside this extension block are only available for Collections that hold elements of type "Restaurant". Now the function countOpenRestaurants(at:) gets the Date parameter which is the specific moment in time that we want know the restaurants' opening status. Inside the function we define the start count to be zero and then go through all the elements. Since the "Iterator.Element" is defined as "Restaurant" we can easily refer to the openingTime and closingTime variables using the element parameter. Every time the conditions are met, we increase the count and finally return it, after the end of the Collection is reached.

let openRestaurantsCount3 = RestaurantHolder.restaurants.numberOfOpenRestaurants(at: Date())
print(openRestaurantsCount3) // 2

Now when we call the function it is very hard to misunderstand what we are trying to do here. Furthermore, since the condition is moved inside the numberOfRestaurants function, calling it looks a lot tidier.

The last example covers filtering the Collection to only get the types of restaurants that satisfy our condition. We want to be able to list all restaurants where we can eat ramen. Let's define another function inside the same Collection extension called getRestaurantsOf(type:):

func getRestaurantsOf(type: Restaurant.FoodType) -> [Iterator.Element] {
    var restaurants: [Iterator.Element] = []
    for element in self {
        if element.type == type {
            restaurants.append(element)
        }
    }
    return restaurants
}

Now this is starting to look pretty familiar. The function takes the Restaurant.FoodType as a parameter, appends the valid elements to the restaurants array which it then returns at the end of the function.

let ramenRestaurants = RestaurantHolder.restaurants.getRestaurantsOf(type: .ramen)

Now if we debug the code, "ramenRestaurants" would hold "Icharaku Ramen" & "Fu-unji" inside it.

I know I already said last example, but there is one more thing. Since all requests can be chained we can easily get the count of open sushi restaurants like this:

let openSushiRestaurants = restaurants.getRestaurantsOf(type: .sushi)
                                      .numberOfOpenRestaurants(at: Date()) // 1

And that was all I wanted to cover today. I hope you liked what you just read and if there is something that you like to ask or comment please use the comment field below. Thanks for reading and I hope you have a great day my friend!

 

Jussi Suojanen

Jussi Suojanen
Swifty SW Developer.

Join the conversation