Grouping elements of a Sequence in Swift

Lately, I’ve been looking for ways to share some of the things I’ve learned in a different, more entertaining way. I’m pretty happy with the results so far. It really made me hungry for experimenting with drawing, painting, recording videos, podcasting. The problem is, I’m still pretty bad at it and that’s why I won’t show you the result of my last 30 minutes spent in Pixelmator.

Use more generics

Oops.

Today I had to work on a relatively common task: grouping a list of view models into separate sections and display them in a UITableView. Somehow that kind of task always ends up with all the other one that doesn’t need to be made generic or reusable easily. After all, it never takes more than 5 to 10 lines of code. Let’s make it both generic and reusable anyway.

What do we want

For this post, I’m going to reuse a piece of code I’m using in Synchronizable’s tests (more about that project later). We have a list of heroes and we want to group them by brand:

enum Brand: String {
    case Marvel
    case DC
    case Valiant
}

struct Hero {
    let name: String
    let brand: Brand
}

let heroes: [Hero] = [
    Hero(name: "Spider-Man", brand: .Marvel),
    Hero(name: "Batman", brand: .DC),
    Hero(name: "Zephyr", brand: .Valiant),
    Hero(name: "Green-Lantern", brand: .DC),
    Hero(name: "Dr Strange", brand: .Marvel),
    Hero(name: "Blue Beetle", brand: .DC)
]

In a world where we wouldn’t mind solving this problem over and over again, we would do it like this:

var groupedHeroes = [Brand: [Hero]]()

heroes.forEach { hero in
    if groupedHeroes[hero.brand] == nil {
        groupedHeroes[hero.brand] = []
    }

    groupedHeroes[hero.brand]?.append(hero)
}

Pretty straightforward. Note that I’m not saying something like that should be made available as a Pod or an iOS framework. That’s way too simple. It doesn’t hurt to have this kind of snippet around, though.

About Hashable

Let’s step back a little and talk about Dictionaries in Swift. A Dictionary is an association between a key and a value. In our situation, the key is the brand and the value is a collection of heroes. You can store pretty much anything in a dictionary, the only restriction is having a key that conforms to the Hashable protocol as we need unique keys. As the documentation states:

Many types in the standard library conform to Hashable: strings, integers, floating-point and Boolean values, and even sets provide a hash value by default. Your own custom types can be Hashable as well.

Making it generic

The Swift standard library has a bunch of methods we can draw inspiration from to make our own groupBy method. The index(where:) method available in the Collection type is a great one: it takes a closure that will allow us to filter a sequence and then retrieve the index. We’ll use a similar approach to retrieve the key we want the sequence to be grouped by.

We want our groupBy method to leverage Swift’s type system so we don’t need to cast the resulting dictionary. Fortunately, that’s exactly what generics constraints are for. That way we can reuse the code we wrote earlier, make it generic and available as an extension of the Sequence type which Array inherits from!

extension Sequence {
    // Using a `typealias` because it's shorter to write `E`
    // Think of it as a shortcut
    typealias E = Iterator.Element

    // Declaring a `K` generic that we'll use as the type of the key
    // for the resulting dictionary. The only restriction is having
    // it conforming to the `Hashable` protocol 
    func groupBy<K: Hashable>(handler: (E) -> K) -> [K: [E]] {
        // Creating the resulting dictionary
        var grouped = [K: [E]]()

        // Iterating over our elements
        self.forEach { item in
            // Retrieving the key based on the current item
            let key = handler(item)

            if grouped[key] == nil {
                grouped[key] = []
            }
            grouped[key]?.append(item)
        }

        return grouped
    }
}

So yeah. If generics are not solving your problems, make sure you’re using enough generics.