Understanding protocol associated types and their constraints

Learn to love the associatedtype keyword

Paul Hudson    April 8th 2018    @twostraws

Protocols with associated types are one of the most powerful, expressive, and flexible features of Swift, but they can complicate your code if you aren’t careful. One of the best ways to control their complexity is to add constraints using them, which restricts how they are used.

To demonstrate this we’re going to build an app architecture a bit like iTunes on macOS, but we’ll start off nice and simple then work our way up.

To begin with, this protocol declares an associated type called ItemType, and says that whatever type conforms to this protocol must provide an array of that item type:

protocol Screen {
    associatedtype ItemType
    var items: [ItemType] { get set }
}

Associated types are effectively holes in a protocol, so when you make a type conform to this protocol it must specify what ItemType actually means.

To continue the example, we’ll create a MainScreen class that conforms to Screen. This will tell Swift that ItemType is a string, which means it must also provide an items string array:

class MainScreen: Screen {
    typealias ItemType = String
    var items = [String]()
}

Swift is smart, though: the typealias line explicitly tells Swift that ItemType is a string, but Swift is able to figure that out because items is an array of strings and therefore ItemType must be a string.

So, we can make the struct a little simpler:

class MainScreen: Screen {
    var items = [String]()
}

Because the protocol knows it can expect an array of something – even though it doesn’t know what the something is – you can write protocol extensions to manipulate the array however you want.

Sponsor Hacking with Swift and reach the world's largest Swift community!

Adding a simple constraint

Our code so far uses a regular associated type, which means that ItemType can be anything at all – a string, an array of strings, an array of arrays of arrays of arrays of strings, an Int, or any custom type you design.

We can use associated type constraints to limit what ItemType can be. This could be done using an existing protocol such as Numeric (any kind of number) or Collection (array, dictionary, and so on, but here we’re going to specify our own custom protocol:

protocol Item {
    init(filename: String)
}

And here are simple Movie and Song structs that conform to Item:

struct Movie: Item {
    var filename: String
}

struct Song: Item {
    var filename: String
}

In case you’re not familiar, Movie and Song conform to Item because Swift structs automatically get a memberwise initializer to satisfy the init(filename: String) requirement.

We can put that to work immediately. We have a MainScreen class that right now can work with any kind of ItemType through the Screen protocol, but we’re now able to be more specific by saying that ItemType can be anything that conforms to Item:

protocol Screen {
    associatedtype ItemType: Item
    var items: [ItemType] { get set }
}

This gives Swift a lot more knowledge about our code – whatever stores items now knows that each item must have a filename attached

As a result of this change, our MainScreen class no longer conforms to Screen – Swift knows that strings can’t be initialized from filenames, so we’re going to make it using our Movie struct instead:

class MainScreen: Screen {
    var items = [Movie]()
}

Adding recursive constraints

You might have been wondering why I made MainScreen a class rather than a struct. Well, Swift 4.1 introduced the ability to make associated type constraints recursive: an associated type can be constrained so that it must also be of the same protocol.

Although this works really well, and lets us express rules that were previously impossible, it does require the use of classes as you’ll see in a moment.

First, we’re going to extend the Screen protocol: in order to conform to this protocol each type must say what kind of child screen it will contain, so that our app forms a tree-like structure:

protocol Screen {
    associatedtype ItemType: Item
    associatedtype ChildScreen: Screen
    var items: [ItemType] { get set }
    var childScreens: [ChildScreen] { get set }
}

We just created a recursive associated type constraints: part of the protocol depends on itself. Such a constraint wasn’t possible before Swift 4.1, so you had to implement extra protocols to create the correct restriction.

To conform to that updated protocol, we can change MainScreen so that it has a childScreens property like this:

class MainScreen: Screen {
    var items = [Movie]()
    var childScreens = [MainScreen]()
}

So, one main screen can show other main screens as children, which can then show more main screens.

This is where the use of class rather than struct becomes important: if a struct has a property of its own type, it will become an infinitely sized struct – object A might contain object B, which might contain object C, and so on. Classes don’t have this restriction because their memory layout is different: a class can contain instances of itself without becoming infinitely sized.

Using where clauses

You’ve seen how we can add constraints to associated types (“ItemType can be anything that conforms to Item”), and also how those constraints can be recursive (“Each Screen must have a ChildScreen can be anything that is itself a Screen"), but Swift lets us go further with where clauses.

Using our current protocols we can implement a hierarchy of screens in our app: there’s one main screen for movies, several category screens (“Action”, “Family”, “Sci-Fi”, etc) and lots of detail screens showing individual movies or other movies that are related to them:

class MainScreen: Screen {
    var items = [Movie]()
    var childScreens = [CategoryScreen]()
}

class CategoryScreen: Screen {
    var items = [Movie]()
    var childScreens = [DetailScreen]()
}

class DetailScreen: Screen {
    var items = [Movie]()
    var childScreens = [DetailScreen]()
}

But this isn’t strict enough, which means you can make mistakes while coding. Anything Swift can do to highlight mistakes while you compile is a good idea, which is why restricting your associated types is so useful.

To see the problem yourself, try changing the CategoryScreen class to this:

class CategoryScreen: Screen {
    var items = [Song]()
    var childScreens = [DetailScreen]()
}

That doesn’t make sense in our hierarchy: we’ve gone from a movie main screen to a song category screen, then down to a movie detail screen – such a thing ought not to be possible, but Swift allows it because all our current criteria are met.

The solution here is to add a where clause to the Screen protocol to make clear that the child screens must be of the same item type as the current screen. I think this reads really naturally in Swift:

associatedtype ChildScreen: Screen where ChildScreen.ItemType == ItemType

That says “whatever fills the ChildScreen hole must also be a Screen, but only one that has the same ItemType as whatever was used to fill our ItemType hole.

In full, it looks like this:

protocol Screen {
    associatedtype ItemType: Item
    associatedtype ChildScreen: Screen where ChildScreen.ItemType == ItemType
    var items: [ItemType] { get set }
    var childScreens: [ChildScreen] { get set }
}

With that change you’ll need to make CategoryScreen back to storing movies, because Swift can now catch the error at compile time – much better!

For the finishing touch, let’s add generics

At this point I hope you’re starting to see the usefulness of associated types and their constraints. However, our current code has a rather serious flaw: we’ve hard-coded Movie into MainScreen, CategoryScreen, and DetailScreen, which means our little iTunes clone isn’t able to work with songs, audio books, and other content.

Fortunately, this can be resolved really cleanly using generics: we can make all three of our screens generic so they can support any type of Item, and both the protocol and generic will work together to make sure all our requirements are met.

Here’s how that looks in code:

class MainScreen<T: Item>: Screen {
    var items = [T]()
    var childScreens = [CategoryScreen<T>]()
}

class CategoryScreen<T: Item>: Screen {
    var items = [T]()
    var childScreens = [DetailScreen<T>]()
}

class DetailScreen<T: Item>: Screen {
    var items = [T]()
    var childScreens = [DetailScreen<T>]()
}

Thanks to Swift’s type-checking system, we now have strict rules in place:

• Each Item must be able to be loaded from a filename.
• Our screens can work with any type of data, as long as it is some sort of Item.
• Each Screen has an array of child screens, where each child screen must be a screen that has the same item type.

That last rule is enforced both through our use of generic and through the protocol, so even if you created a new type without generics you would still be protected by the protocol.

If you’re keen to learn more about how Swift can help you write better, safer code, you should read my article: How Swift keypaths let us write more natural code.

Sponsor Hacking with Swift and reach the world's largest Swift community!