Improving your protocols with associated types

One more awesome aspect of protocol-oriented programming is the use of associated types. An associated type is a generic, non-existing type that can be used in your protocol like any type that does exist. The real type of this generic is determined by the compiler based on the context it's used in. This description is abstract, and you might not immediately understand why or how an associated type can benefit your protocols. After all, aren't protocols themselves a very flexible way to make several unrelated objects fit certain criteria based on the protocols they conform to?

To illustrate and discover the use of associated types, you will expand your animal kingdom a bit. What you should do is give the herbivores an eat method and an array to keep track of the plants they've eaten, as follows:

protocol HerbivoreType {
  var plantsEaten: [PlantType] { get set }

  mutating func eat(plant: PlantType)
}

extension HerbivoreType {
  mutating func eat(plant: PlantType) {
    plantsEaten.append(plant)
  }
}

This code looks fine at first sight. A herbivore eats plants, and this is established by this protocol. The PlantType protocol is defined as follows:

protocol PlantType {   
  var latinName: String { get }   
} 

Let's define two different plant types and an animal that will be used to demonstrate the problem with the preceding code:

struct Grass: PlantType{
  var latinName = "Poaceae"
}

struct Pine: PlantType{
  var latinName = "Pinus"
}

struct Cow: HerbivoreType {
  var plantsEaten = [PlantType]()
} 

There shouldn't be a big surprise here. Let's continue with creating a Cow instance and feed it Pine:

var cow = Cow()   
let pine = Pine()   
cow.eat(plant: pine)

This doesn't really make sense. Cows don't eat pines; they eat grass! We need some way to limit this cow's food intake because this approach isn't going to work. Currently, you can feed HerbivoreType animals anything that's considered a plant. You need some way to limit the types of food your cows are given. In this case, you should restrict the FoodType to Grass only, without having to define the eat(plant:) method for every plant type you might want to feed a HerbivoreType.

The problem you're facing now is that all HerbivoreType animals mainly eat one plant type, and not all plant types are a good fit for all herbivores. This is where associated types are a great solution. An associated type for the HerbivoreType protocol can constrain the PlantType that a certain herbivore can eat to a single type that is defined by the HerbivoreType itself. Let's see what this looks like:

protocol HerbivoreType {
  // 1
  associatedtype Plant: PlantType

  var plantsEaten: [Plant] { get set }

  mutating func eat(plant: Plant)
}

extension HerbivoreType {
  mutating func eat(plant: Plant) {
    // 2
    print("eating a (plant.latinName)")
    plantsEaten.append(plant)
  }
} 

The first highlighted line associates the generic Plant type, which doesn't exist as a real type, with the protocol. A constraint has been added to Plant to ensure that it's a PlantType.

The second highlighted line demonstrates how the Plant associated type is used as a PlantType. The plant type itself is merely an alias for any type that conforms to PlantType and is used as the type of object we use for plantsEaten and the eat methods. Let's redefine the Cow struct to see this associated type in action:

struct Cow: HerbivoreType {   
  var plantsEaten = [Grass]()
}

Instead of making plantsEaten a PlantType array, it's now defined as an array of Grass. In the protocol and the definition, the type of plant is now Grass. The compiler understands this because the plantsEaten array is defined as [Grass]. Let's define a second HerbivoreType that eats a different type of PlantType:

struct Carrot: PlantType {   
  let latinName = "Daucus carota"   
}   

struct Rabbit: HerbivoreType {   
  var plantsEaten = [Carrot]()   
} 

If you try to feed a Cow some carrots, or if you attempt to feed the Rabbit a Pine, the compiler will throw errors. The reason for this is that the associated type constraint allows you to define the type of Plant in each struct separately.

One side note about associated types is that it's not always possible for the compiler to correctly infer the real type for an associated type. In our current example, this would happen if we didn't have the plantsEaten array in the protocol. The solution would be to define a typealias on objects that conform to HerbivoreType so that the compiler understands which type Plant represents:

protocol HerbivoreType {   
  associatedtype Plant: PlantType   

  mutating func eat(plant: Plant)   
}   

struct Cow: HerbivoreType {   
  typealias Plant = Grass   
} 

Associated types can be really powerful when used correctly, but sometimes using them can also cause you a lot of headaches because of the amount of inferring the compiler has to do. If you forget a few tiny steps, the compiler can quickly lose track of what you're trying to do, and the error messages aren't always the most unambiguous messages. Keep this in mind when you're using associated types, and try to make sure that you're as explicit as possible about the type you're looking to be associated. Sometimes, adding a typealias to give the compiler a helping hand is better than trying to get the compiler to infer everything on its own correctly.

This type of flexibility is not limited to protocols. You can also add generic properties to functions, classes, structs, and enums. Let's see how this works and how it can make your code extremely flexible.