When creating collections of objects or values in Swift, we usually use data structures provided by the standard library - such as Array, Dictionary and Set. While those three cover most use cases, sometimes creating a custom wrapper collection can enable you to make your code more predictable and less prone to errors.

This week, let's take a look at how we as app developers can define such custom collections in Swift, and how - combined with the power of enums - it can let us create some pretty nice APIs for ourselves.

Removing optionals

Like we took a look at in "Handling non-optional optionals in Swift", reducing the need to use optionals when the values you're looking for are actually required can really help us avoid bugs and make our code easier to work with.

The problem with collections in general, is that you usually can't make a guarantee whether they contain a certain value, and you therefor tend to end up with lots of optionals and logic that requires to unwrap them one way or another.

Let's say that we're building an app for a grocery store, and we want to have a UI that lets the user display all products by category. To create a model for such a UI, we might use a Dictionary, which uses Category as its key type and [Product] as its value type, like this:

let products: [Category : [Product]] = [
    .dairy: [
        Product(name: "Milk", category: .dairy),
        Product(name: "Butter", category: .dairy)
    ],
    .vegetables: [
        Product(name: "Cucumber", category: .vegetables),
        Product(name: "Lettuce", category: .vegetables)
    ]
]

While the above works, it will require us to write code like this in order to - for example - only display all dairy products:

if let dairyProducts = products[.dairy] {
    guard !dairyProducts.isEmpty else {
        renderEmptyView()
        return
    }

    render(dairyProducts)
} else {
    renderEmptyView()
}

That's fine, but it could be nicer. Inserting new products, however, becomes much more of a hassle:

class ShoppingCart {
    private(set) var products = [Category : [Product]]()

    func add(_ product: Product) {
        if var productsInCategory = products[product.category] {
            productsInCategory.append(product)
            products[product.category] = productsInCategory
        } else {
            products[product.category] = [product]
        }
    }
}

The good news is that we can make both of the above examples much nicer and cleaner by creating our own custom collection. And the even better news is that - thanks to Swift's protocol oriented design - creating such a collection is actually quite easy!

To be a collection

All collections in the Swift standard library conform to the Collection protocol, which in turn inherits from the Sequence protocol. By making a custom collection conform to these two protocols, it can take advantage of all the standard collection operations - such as iterating & filtering - completely for free.

Let's start by defining the base of our custom ProductCollection, that will enable us to deal with products and categories in a much nicer way.

struct ProductCollection {
    typealias DictionaryType = [Category : [Product]]

    // Underlying, private storage, that is the same type of dictionary
    // that we previously was using at the call site
    private var products = DictionaryType()

    // Enable our collection to be initialized with a dictionary
    init(products: DictionaryType) {
        self.products = products
    }
}

Next, we'll make it conform to Collection by implementing the protocol requirements. Most of what we'll do is to simply forward calls to the underlying products dictionary, and let that do the "heavy lifting":

extension ProductCollection: Collection {
    // Required nested types, that tell Swift what our collection contains
    typealias Index = DictionaryType.Index
    typealias Element = DictionaryType.Element

    // The upper and lower bounds of the collection, used in iterations
    var startIndex: Index { return products.startIndex }
    var endIndex: Index { return products.endIndex }

    // Required subscript, based on a dictionary index
    subscript(index: Index) -> Iterator.Element {
        get { return products[index] }
    }

    // Method that returns the next index when iterating
    func index(after i: Index) -> Index {
        return products.index(after: i)
    }
}

The above code is using Swift 4, which makes defining custom collections a lot simpler, thanks to improvements in generic constraints (we'll look more closely into those improvements and how to use type constraints in a future post).

We now have a custom collection that can be used just as one of the built-in ones. We can, for example, iterate through it:

for (category, productsInCategory) in products {
    ...
}

Or use an operation like map on it:

let categories = productCollection.map { $0.key }

Custom collection APIs

Now that we have laid the ground work for our collection, let's start adding some APIs to it that will enable us to make our product handling code a lot nicer. We'll start with a custom subscript overload that lets us get or set an array of products without having to deal with optionals:

extension ProductCollection {
    subscript(category: Category) -> [Product] {
        get { return products[category] ?? [] }
        set { products[category] = newValue }
    }
}

Let's also add a convenience API to easily insert a new Product into our collection:

extension ProductCollection {
    mutating func insert(_ product: Product) {
        var productsInCategory = self[product.category]
        productsInCategory.append(product)
        self[product.category] = productsInCategory
    }
}

We can now go back to our original product handling code, and update it to be much nicer. For reading:

let dairyProducts = products[.dairy]

if dairyProducts.isEmpty {
    renderEmptyView()
} else {
    render(dairyProducts)
}

And writing:

class ShoppingCart {
    private(set) var products = ProductCollection()

    func add(product: Product) {
        products.insert(product)
    }
}

Becoming expressible by a literal

OK, time for the bonus round! Since our custom collection is basically just a wrapper around a Dictionary, we can easily add support for initializing one using a dictionary literal. Doing that will enable us to write code like this:

let products: ProductCollection = [
    .dairy: [
        Product(name: "Milk", category: .dairy),
        Product(name: "Butter", category: .dairy)
    ],
    .vegetables: [
        Product(name: "Cucumber", category: .vegetables),
        Product(name: "Lettuce", category: .vegetables)
    ]
]

Pretty cool! This is not only useful for reducing verbosity in our production code, but will also make setting up product collection mocks in our tests a lot simpler.

All we have to do to make the above happen is to conform to ExpressibleByDictionaryLiteral, which requires us to implement an initializer that takes a literal, like this:

extension ProductCollection: ExpressibleByDictionaryLiteral {
    typealias Key = Category
    typealias Value = [Product]

    init(dictionaryLiteral elements: (Category, [Product])...) {
        for (category, productsInCategory) in elements {
            products[category] = productsInCategory
        }
    }
}

Conclusion

Using custom collections can be a really powerful tool to handle groups of values in a more predictable and easy-to-use way. While it probably shouldn't always be your go-to solution as soon as you're dealing with multiple values, in the right situations it can really help you write cleaner code.

Understanding how things like collections work under the hood can also be really helpful when debugging, or to give you insight as to how code dealing with collections can be optimized. And what better way to learn more about collections than building your own? 😄

What do you think? Do you already use custom collections or will you try making one? Let me know, along with any questions, comments or feedback you might have - either here in the comment section below, or on Twitter @johnsundell.

Thanks for reading! 🚀