Swift, Create Custom Collection

Swift provides a powerful standard library with built-in collection types like `Array`, `Dictionary`, and `Set`. However, in some cases, creating custom collections can make your code more predictable, less error-prone, and enable you to define more expressive APIs. In this article, we’ll explore how to create custom collections in Swift and leverage the power of enums to design clean APIs.

The Challenge of Optionals in Collections

One common issue when working with collections is dealing with optionals. Standard collections don’t guarantee the presence of a specific value, leading to a lot of optional unwrapping and complex logic. Let’s consider an example where we want to categorize and display products in a grocery store:

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)
 ]
]

To display dairy products, we have to write code like this:

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

Adding new products to this structure is also cumbersome:

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]
 }
 }
}

We can simplify and improve these scenarios by creating a custom collection in Swift.

Creating a Custom Collection

All Swift collections conform to the `Collection` protocol, which inherits from the `Sequence` protocol. By making our custom collection conform to these protocols, we can take advantage of various standard collection operations. Let’s start by defining our `ProductCollection`:

struct ProductCollection {
 private var products = [Category: [Product]]()
}

Next, we’ll make it conform to the `Collection` protocol by implementing the required protocol methods:

extension ProductCollection: Collection {
 typealias Index = Dictionary<Category, [Product]>.Index
 typealias Element = Dictionary<Category, [Product]>.Element
var startIndex: Index { return products.startIndex }
 var endIndex: Index { return products.endIndex }
subscript(index: Index) -> Element {
 get { return products[index] }
 }
func index(after i: Index) -> Index {
 return products.index(after: i)
 }
}

Now, our custom collection can be used just like any built-in collection. We can iterate through it or use operations like `map`:

for (category, productsInCategory) in productCollection {
 // …
}
let categories = productCollection.map { $0.key }

Adding Custom Collection APIs

Let’s further enhance our custom collection by adding some custom APIs to make product handling more convenient. First, we’ll create a custom subscript that allows us to get or set an array of products without dealing with optionals:

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

We can also add a convenience method 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
 }
}

Now, we can significantly improve our original code for both reading and writing products:

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

class ShoppingCart {
 private(set) var products = ProductCollection()
func add(product: Product) {
 products.insert(product)
 }
}

Expressible by Dictionary Literal

As a bonus, we can make our custom collection even more expressive by allowing it to be initialized using a dictionary literal. This feature enables 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)
 ]
]

To achieve this, we need to conform to the `ExpressibleByDictionaryLiteral` protocol by implementing an initializer that accepts a dictionary literal:

extension ProductCollection: ExpressibleByDictionaryLiteral {
 typealias Key = Category
 typealias Value = [Product]
init(dictionaryLiteral elements: (Category, [Product])…) {
 for (category, productsInCategory) in elements {
 products[category] = productsInCategory
 }
 }
}

Conclusion

Creating custom collections in Swift can greatly simplify your code