Associated Types in Swift

Swift’s powerful generics system allows developers to write flexible and reusable code. One of the core features of generics is associated types, which are used in protocols to define placeholder types that get specified later when the protocol is adopted by a class, struct, or enum. Understanding associated types is crucial when working with Swift’s protocols, and in this blog, we’ll explore their concept, usage, and real-world examples.

What Are Associated Types?

Associated types in Swift are placeholders that define types that will be provided when a protocol is adopted. They allow protocols to work with generic types without specifying what those types are upfront. When you declare an associated type, you’re essentially saying, “I don’t care what the type is right now, but the conforming type will specify it later.”

Syntax

You declare an associated type in a protocol using the associatedtype keyword:

protocol SomeProtocol {
    associatedtype Item
    func doSomething(with item: Item)
}

In this protocol, Item is an associated type. Any type that conforms to SomeProtocol will need to define what Item refers to.

Why Use Associated Types?

Associated types make protocols more flexible by allowing them to work with generic types. Instead of defining a protocol that works with one specific type, you can define a protocol that works with any type, which will be specified later when the protocol is implemented.

Basic Example

Let’s start with a simple example. Suppose we want to create a protocol for a container that can hold items of any type. We don’t know the type of the items at the time of protocol declaration, but the conforming type will specify it later.

protocol Container {
    associatedtype Item
    var items: [Item] { get set }
    
    func add(item: Item)
    func getItem(at index: Int) -> Item
}

Here, Container is a protocol that defines an associated type Item. The protocol also defines an array items that holds elements of type Item and functions to add and retrieve items from the container.

Now, let’s conform to this protocol using a StringContainer that holds strings:

struct StringContainer: Container {
    var items: [String] = []
    
    func add(item: String) {
        items.append(item)
    }
    
    func getItem(at index: Int) -> String {
        return items[index]
    }
}

In this case, the associated type Item is specified as String when the protocol is adopted by StringContainer. We can now use the StringContainer to store strings:

var myContainer = StringContainer()
myContainer.add(item: "Hello")
myContainer.add(item: "World")

print(myContainer.getItem(at: 0))  // Output: Hello

Using Associated Types with Multiple Types

Associated types are not limited to single types; you can create more complex examples where protocols have multiple associated types.

For instance, let’s say we want to create a protocol for a dictionary-like data structure that maps keys to values.

protocol KeyValueStore {
    associatedtype Key
    associatedtype Value
    
    func setValue(_ value: Value, forKey key: Key)
    func getValue(forKey key: Key) -> Value?
}

Here, the protocol KeyValueStore defines two associated types: Key and Value. This protocol can be used to build any dictionary-like structure that maps any type of key to any type of value.

We can now create a struct that conforms to this protocol:

struct DictionaryStore: KeyValueStore {
    var storage: [String: Int] = [:]
    
    func setValue(_ value: Int, forKey key: String) {
        storage[key] = value
    }
    
    func getValue(forKey key: String) -> Int? {
        return storage[key]
    }
}

Here, the associated types Key and Value are specified as String and Int, respectively. Now, we can use this store to map strings to integers:

var store = DictionaryStore()
store.setValue(42, forKey: "Answer")
print(store.getValue(forKey: "Answer") ?? "Not found")  // Output: 42

Associated Types with Constraints

You can add constraints to associated types by specifying that they must conform to a certain protocol. This is useful when you want to restrict the types used in a protocol to specific kinds of types.

For example, let’s modify the Container protocol to ensure that the items it holds must conform to the Equatable protocol:

protocol Container {
    associatedtype Item: Equatable
    var items: [Item] { get set }
    
    func add(item: Item)
    func contains(item: Item) -> Bool
}

By using Item: Equatable, we specify that Item must conform to the Equatable protocol. Now, any type that conforms to Container must use a type that supports equality comparison.

Here’s how we can adopt this modified Container protocol:

struct IntContainer: Container {
    var items: [Int] = []
    
    func add(item: Int) {
        items.append(item)
    }
    
    func contains(item: Int) -> Bool {
        return items.contains(item)
    }
}

Since Int conforms to Equatable, this implementation works fine:

var intContainer = IntContainer()
intContainer.add(item: 5)
intContainer.add(item: 10)

print(intContainer.contains(item: 5))  // Output: true
print(intContainer.contains(item: 7))  // Output: false

Associated Types and Type Erasure

One of the challenges with associated types is that they make it harder to work with protocols in a generic way. Since associated types are not concrete until a type conforms to the protocol, you can’t use the protocol as a type by itself. This can be solved using type erasure.

For example, if we try to use the Container protocol as a type directly, Swift won’t allow it because the associated type is not specified:

// This won't compile because `Container` has an associated type.
var container: Container

To work around this, we use type erasure to create a concrete type that hides the associated type, allowing us to use the protocol as a type. A common way to do this is by creating a wrapper around the protocol:

struct AnyContainer<T: Equatable>: Container {
    private let _add: (T) -> Void
    private let _contains: (T) -> Bool
    private var _items: [T]
    
    var items: [T] {
        get { return _items }
        set { _items = newValue }
    }
    
    init<C: Container>(_ container: C) where C.Item == T {
        _items = container.items
        _add = container.add
        _contains = container.contains
    }
    
    func add(item: T) {
        _add(item)
    }
    
    func contains(item: T) -> Bool {
        return _contains(item)
    }
}

Now, you can use AnyContainer as a concrete type:

let container = AnyContainer(IntContainer())

Real-World Example: Swift’s Collection Protocol

Swift’s standard library uses associated types extensively. One prominent example is the Collection protocol, which defines how types like arrays, sets, and dictionaries behave.

protocol Collection {
    associatedtype Element
    var count: Int { get }
    func element(at index: Int) -> Element
}

Collection defines an associated type Element, which represents the type of the items stored in the collection. When you adopt Collection, you specify what Element will be, allowing the protocol to be used with any type of collection.

Conclusion

Associated types are a powerful feature in Swift’s generics system, allowing protocols to define placeholder types that are specified when the protocol is adopted. They make protocols flexible and adaptable to different types, enabling developers to write highly reusable and efficient code. Whether you’re defining containers, key-value stores, or working with Swift’s standard library protocols, associated types are key to mastering Swift’s protocol-oriented programming.

Tags: Swift