Collection grouping simplified

We've all probably encountered a situation where we needed to group an array of elements by some value, while still keeping the original order of elements. For example, when you need to group people by their age.

struct Person {
    let name: String
    let age: Int
}

let people = [
    Person(name: "Mark", age: 25),
    Person(name: "Ann", age: 24),
    Person(name: "John", age: 25)
]

Swift 4.0 introduced the following dictionary functionality which allowed grouping:

let grouped = Dictionary(grouping: people) { $0.age }.values

/// Result
[
    [
        Person(name: "Ann", age: 24)
    ],
    [
        Person(name: "Mark", age: 25),
        Person(name: "John", age: 25)
    ]
]

It looks good, right? Unfortunately, not really. Since dictionaries are unordered collections of key-value associations, order of inner arrays might not always be preserved - as seen in the example, array of people aged 24 comes before those aged 25, while it should be reversed since "Mark" is the first person in the starting list.

Why don't we add some syntactic sugar, as well as make it always sort properly by using the power of swift extensions and generics:

extension Collection {

    func groupBy<GroupingType: Hashable>(key: (Element) -> (GroupingType)) -> [[Element]] {
        var groups: [GroupingType: [Element]] = [:]
        var groupsOrder: [GroupingType] = []

        forEach { element in
            let key = key(element)

            if case nil = groups[key]?.append(element) {
                groups[key] = [element]
                groupsOrder.append(key)
            }
        }

        return groupsOrder.map { groups[$0] ?? [] }
    }

}

By doing this, our code becomes much cleaner and declarative, as you can see below:

let grouped = people.groupBy { $0.age }

/// Result
[
    [
        Person(name: "Mark", age: 25),
        Person(name: "John", age: 25)
    ],
    [
        Person(name: "Ann", age: 24)
    ]
]

Happy grouping!