Building a Custom Collection with Protocols in Swift

Array, Dictionary and Set are commonly used collection types that come bundled in the Swift standard library. But what if they don’t provide everything you need for your app right out of the box? No worries. You can make your own custom collections using protocols from the Swift standard library!

Collections in Swift come with a plethora of handy utilities for iterating through them, filtering them and much more. Instead of using a custom collection, you could add all of business logic to your own code. But this leaves your code bloated, hard to maintain and duplicating what the standard library provides.

Fortunately, Swift provides powerful collection protocols so you can create your own collection types specifically tailored to meet your app’s requirements. You can bring the power of Swift collections simply by implementing these protocols.

In this tutorial, you’re going to build a multiset, otherwise known as a bag, from scratch.

Along the way, you’ll learn how to:

• Adopt these protocols: Hashable, Sequence, Collection, CustomStringConvertible, ExpressibleByArrayLiteral and ExpressibleByDictionaryLiteral.
• Create custom initializations for your collections.
• Improve your custom collections with custom methods.

Time to jump right in!

Getting Started

Start by downloading the project materials using the Download Materials button at the top or bottom of this tutorial. Then open the file Bag.playground in the starter folder.

Creating the Bag Struct

Next, add the following code to your playground:

struct Bag<Element: Hashable> {
}

And just like that, “Papa’s got a brand new bag”!

Your Bag is a generic structure that requires a Hashable element type. Requiring Hashable elements allows you to compare and store unique values with O(1) time complexity. This means that no matter the size of its contents, Bag will perform at constant speeds. Also, notice that you’re using a struct; this enforces value semantics as Swift does for standard collections.

A Bag is like a Set in that it does not store repeated values. The difference is this: A Bag keeps a running count of any repeated values while a Set does not.

Think about it like a shopping list. If you want more than one of something, you don’t list it multiple times. You simply write the number you want next to the item.

To model this, add the following properties to Bag in your playground:

// 1
fileprivate var contents: [Element: Int] = [:]

// 2
var uniqueCount: Int {
  return contents.count
}

// 3
var totalCount: Int {
  return contents.values.reduce(0) { $0 + $1 }
}

These are the basic properties needed for a Bag. Here’s what each does:

1. contents: Uses a Dictionary as the internal data structure. This works great for a Bag because it enforces unique keys which you’ll use to store elements. The value for each element is its count. Notice that you mark this property as fileprivate to hide the inner workings of Bag from the outside world.
2. uniqueCount: Returns the number of unique items, ignoring their individual quantities. For example, a Bag with 4 oranges and 3 apples will return a uniqueCount of 2.
3. totalCount: Returns the total number of items in the Bag. In the example above, totalCount will return 7.

Adding Edit Methods

Now you’ll implement some methods to edit the contents of Bag.

Add the following method below the properties you just added:

// 1
mutating func add(_ member: Element, occurrences: Int = 1) {
  // 2
  precondition(occurrences > 0,
    "Can only add a positive number of occurrences")

  // 3
  if let currentCount = contents[member] {
    contents[member] = currentCount + occurrences
  } else {
    contents[member] = occurrences
  }
}

Here’s what this does:

1. add(_:occurrences:): Provides a way to add elements to the Bag. It takes two parameters: the generic type, Element, and an optional number of occurrences. You mark the method as mutating so you can modify the contents instance variable.
2. precondition(_:_:): Requires greater than 0 occurrences. If this condition is false, execution stops and the String that follows the condition will appear in the playground debugger.
3. This section checks if the element already exists in the bag. If it does, it increments the count. If it doesn’t, it creates a new element.

Now that you have a way to add elements to your Bag instance, you also need a way to remove them.

Add the following method just below add(_:occurrences:):

mutating func remove(_ member: Element, occurrences: Int = 1) {
  // 1
  guard 
    let currentCount = contents[member],
    currentCount >= occurrences 
    else {
      return
  }

  // 2
  precondition(occurrences > 0,
    "Can only remove a positive number of occurrences")

  // 3
  if currentCount > occurrences {
    contents[member] = currentCount - occurrences
  } else {
    contents.removeValue(forKey: member)
  }
}

Notice that remove(_:occurrences:) takes the same parameters as add(_:occurrences:). Here’s how it works:

1. First, it checks that the element exists and that it has at least the number of occurrences the caller is removing. If it doesn’t, the method returns.
2. Next, it makes sure that the number of occurrences to remove is greater than 0.
3. Finally, it checks if the element’s current count is greater than the number of occurrences to remove. If greater, then it sets the element’s new count by subtracting the number of occurrences to remove from the current count. If not greater then currentCount and occurrences are equal and it removes the element entirely.

Right now Bag doesn’t do much. You can’t access its contents and you can’t operate on your collection with any of the useful collection methods like map, filter, and so on.

But all is not lost! Swift provides the tools you need to make Bag into a legitimate collection. You simply need to conform to a few protocols.