Building a Custom Collection with Protocols in Swift [FREE]

Update note : Brody Eller updated this tutorial for Swift 5. Eric Cerney wrote the original.

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!

Note : This tutorial works with Swift 5.0. Previous versions will not compile because of major changes to the Swift standard library.

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.

Note : If you prefer, you can create your own Xcode playground. If you do, delete all the default code to start with an empty playground.

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 .

Adding Add Method

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:

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. 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. 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.

Note : You’ll use precondition throughout this tutorial to ensure that you’re using Bag as you intend. You’ll also use precondition as a sanity check to make sure things work as expected as you add functionality. Doing so incrementally will keep you from accidentally breaking functionality that was working before.

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

Implementing the Remove Method

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.

Adopting Protocols

In Swift, a protocol defines a set of properties and methods that must be implemented in an object that adopts it. To adopt a protocol, simply add a colon after the definition of your class or struct followed by the name of the protocol you’d like to adopt. After you declare your adoption of the protocol, implement the required variables and methods on your object. Once complete, your object conforms to the protocol.

Note : You can learn more about protocols in our Protocol Oriented Programming tutorial .

Here’s an easy example. Currently, Bag objects expose little information on the results sidebar in the Playground.

Add the following code to the end of the playground (outside of the struct) to see Bag in action:

var shoppingCart = Bag<String>()
shoppingCart.add("Banana")
shoppingCart.add("Orange", occurrences: 2)
shoppingCart.add("Banana")
shoppingCart.remove("Orange")

Then press Command-Shift-Enter to execute the playground.

This creates a new Bag with a few pieces of fruit. If you look at the playground debugger, you’ll see the object type without any of its contents.

Adopting CustomStringConvertible

Fortunately, Swift provides the CustomStringConvertible protocol for just this situation! Add the following just after the closing brace of Bag :

extension Bag: CustomStringConvertible {
  var description: String {
    return String(describing: contents)
  }
}

Conforming to CustomStringConvertible requires implementation of a single property named description . This property returns the textual representation of the specific instance.

This is where you would put any logic needed to create a string representing your data. Because Dictionary conforms to CustomStringConvertible , you simply delegate the description call to contents .

Press Command-Shift-Enter to run the playground again.

Take a look at the newly improved debug information for shoppingCart :

Awesome! Now, as you add functionality to Bag , you’ll be able to verify its contents.

Great! You’re on your way as you create powerful collection types that feel native. Next up is initialization.

Creating Initializers

It’s pretty annoying that you have to add each element one at a time. You should be able to initialize your Bag by passing in a collection of objects to add.

Add the following code to the end of the playground (but notice that this will not compile just yet):

let dataArray = ["Banana", "Orange", "Banana"]
let dataDictionary = ["Banana": 2, "Orange": 1]
let dataSet: Set = ["Banana", "Orange", "Banana"]

var arrayBag = Bag(dataArray)
precondition(arrayBag.contents == dataDictionary,
  "Expected arrayBag contents to match \(dataDictionary)")

var dictionaryBag = Bag(dataDictionary)
precondition(dictionaryBag.contents == dataDictionary,
  "Expected dictionaryBag contents to match \(dataDictionary)")

var setBag = Bag(dataSet)
precondition(setBag.contents == ["Banana": 1, "Orange": 1],
  "Expected setBag contents to match \(["Banana": 1, "Orange": 1])")

This is how you might expect to create a Bag . But it won’t compile because you haven’t defined an initializer that takes other collections. Rather than explicitly creating an initialization method for each type, you’ll use generics .

Add the following methods just below totalCount inside the implementation of Bag :

// 1
init() { }

// 2
init<S: Sequence>(_ sequence: S) where
  S.Iterator.Element == Element {
  for element in sequence {
    add(element)
  }
}

// 3
init<S: Sequence>(_ sequence: S) where
  S.Iterator.Element == (key: Element, value: Int) {
  for (element, count) in sequence {
    add(element, occurrences: count)
  }
}

Here’s what you just added:

1. First, you created an empty initializer. You’re required to add this when defining additional init methods.
2. Next, you added an initializer that accepts anything that conforms to the Sequence protocol where the elements of that sequence are the same as the elements of the Bag . This covers both Array and Set types. You iterate over the passed in sequence and add each element one at a time.
3. After this, you added a similar initializer but one that accepts tuples of type (Element, Int) . An example of this is a Dictionary . Here, you iterate over each element in the sequence and add the specified count.

Press Command-Shift-Enter again to run the playground. Notice that the code you added at the bottom earlier now works.

Initializing Collections

These generic initializers enable a much wider variety of data sources for Bag objects. However, they do require you to first create the collection you pass into the initializer.

To avoid this, Swift supplies two protocols that enable initialization with sequence literals. Literals give you a shorthand way to write data without explicitly creating an object.

To see this, first add the following code to the end of your playground: (Note: This, too, will generate errors until you add the needed protocols.)

var arrayLiteralBag: Bag = ["Banana", "Orange", "Banana"]
precondition(arrayLiteralBag.contents == dataDictionary,
  "Expected arrayLiteralBag contents to match \(dataDictionary)")

var dictionaryLiteralBag: Bag = ["Banana": 2, "Orange": 1]
precondition(dictionaryLiteralBag.contents == dataDictionary,
  "Expected dictionaryLiteralBag contents to match \(dataDictionary)")

The code above is an example of initialization using Array and Dictionary literals rather than objects.

Now, to make these work, add the following two extensions just below the CustomStringConvertible extension:

// 1
extension Bag: ExpressibleByArrayLiteral {
  init(arrayLiteral elements: Element...) {
    self.init(elements)
  }
}

// 2
extension Bag: ExpressibleByDictionaryLiteral {
  init(dictionaryLiteral elements: (Element, Int)...) {
    self.init(elements.map { (key: $0.0, value: $0.1) })
  }
}

1. ExpressibleByArrayLiteral is used to create a Bag from an array style literal. Here you use the initializer you created earlier and pass in the elements collection.
2. ExpressibleByDictionaryLiteral does the same but for dictionary style literals. The map converts elements to the named-tuple the initializer expects.

With Bag looking a lot more like a native collection type, it’s time to get to the real magic.

Understanding Custom Collections

You’ve now learned enough to understand what a custom collection actually is: A collection object that you define that conforms to both the Sequence and Collection protocols.

In the last section, you defined an initializer that accepts collection objects conforming to the Sequence protocol. Sequence represents a type that provides sequential, iterated access to its elements. You can think of a sequence as a list of items that let you step over each element one at a time.

Iteration is a simple concept, but this ability provides huge functionality to your object. It allows you to perform a variety of powerful operations like:

• map(_:) : Returns an array of results after transforming each element in the sequence using the provided closure.
• filter(_:) : Returns an array of elements that satisfy the provided closure predicate.
• sorted(by:) : Returns an array of the elements in the sequence sorted based on the provided closure predicate.

This barely scratches the surface. To see all methods available from Sequence , take a look at Apple’s documentation on the Sequence Protocol .

Enforcing Non-destructive Iteration

One caveat: Sequence does not require conforming types to be non-destructive. This means that after iteration, there’s no guarantee that future iterations will start from the beginning. That’s a huge issue if you plan on iterating over your data more than once.

To enforce non-destructive iteration, your object needs to conform to the Collection protocol.

Collection inherits from Indexable and Sequence .

The main difference is that a collection is a sequence you can traverse multiple times and access by index.

You’ll get many methods and properties for free by conforming to Collection . Some examples are:

• isEmpty : Returns a boolean indicating if the collection is empty or not.
• first : Returns the first element in the collection.
• count : Returns the number of elements in the collection.

There are many more available based on the type of elements in the collection. Check them out in Apple’s documentation on the Collection Protocol .

Grab your Bag and adopt these protocols!

Adopting the Sequence Protocol

The most common action performed on a collection type is iterating through its elements. For example, add the following to the end of the playground:

for element in shoppingCart {
  print(element)
}

As with Array and Dictionary , you should be able to loop through a Bag . This won’t compile because currently the Bag type doesn’t conform to Sequence .

Time to fix that now.

Conforming to Sequence

Add the following just after the ExpressibleByDictionaryLiteral extension:

extension Bag: Sequence {
  // 1
  typealias Iterator = DictionaryIterator<Element, Int>

  // 2
  func makeIterator() -> Iterator {
    // 3
    return contents.makeIterator()
  }
}

There’s not too much needed to conform to Sequence . In the code above, you:

1. Create a typealias named Iterator as DictionaryIterator . Sequence requires this to know how you iterate your sequence. DictionaryIterator is the type that Dictionary objects use to iterate through their elements. You’re using this type because Bag stores its underlying data in a Dictionary .
2. Define makeIterator() as a method that returns an Iterator for stepping through each element of the sequence.
3. Return an iterator by delegating to makeIterator() on contents , which itself conforms to Sequence .

That’s all you need to make Bag conform to Sequence !

You can now iterate through each element of a Bag and get the count for each object. Add the following to the end of the playground after the previous for-in loop:

for (element, count) in shoppingCart {
  print("Element: \(element), Count: \(count)")
}

Press Command-Shift-Enter to run the playground. Open the playground console and you’ll see the printout of the elements and their count in the sequence.

Viewing Benefits of Sequence

Being able to iterate through a Bag enables many useful methods implemented by Sequence . Add the following to the end of the playground to see some of these in action:

// Find all elements with a count greater than 1
let moreThanOne = shoppingCart.filter { $0.1 > 1 }
moreThanOne
precondition(
  moreThanOne.first!.key == "Banana" && moreThanOne.first!.value == 2,
  "Expected moreThanOne contents to be [(\"Banana\", 2)]")

// Get an array of all elements without their counts
let itemList = shoppingCart.map { $0.0 }
itemList
precondition(
  itemList == ["Orange", "Banana"] ||
    itemList == ["Banana", "Orange"],
  "Expected itemList contents to be [\"Orange\", \"Banana\"] or [\"Banana\", \"Orange\"]")

// Get the total number of items in the bag
let numberOfItems = shoppingCart.reduce(0) { $0 + $1.1 }
numberOfItems
precondition(numberOfItems == 3,
  "Expected numberOfItems contents to be 3")

// Get a sorted array of elements by their count in descending order
let sorted = shoppingCart.sorted { $0.0 < $1.0 }
sorted
precondition(
  sorted.first!.key == "Banana" && moreThanOne.first!.value == 2,
  "Expected sorted contents to be [(\"Banana\", 2), (\"Orange\", 1)]")

Press Command-Shift-Enter to run the playground and see these in action.

These are all useful methods for working with sequences — and you got them practically for free!

Now, you could be content with the way things are with Bag , but where's the fun in that?! You can definitely improve the current Sequence implementation.

Improving Sequence

Currently, you're relying on Dictionary to handle the heavy lifting for you. That's fine because it makes creating powerful collections of your own easy. The problem is that it creates strange and confusing situations for Bag users. For example, it's not intuitive that Bag returns an iterator of type DictionaryIterator .

But Swift comes to the rescue again! Swift provides the type AnyIterator to hide the underlying iterator from the outside world.

Replace the implementation of the Sequence extension with the following:

extension Bag: Sequence {
  // 1
  typealias Iterator = AnyIterator<(element: Element, count: Int)>

  func makeIterator() -> Iterator {
    // 2
    var iterator = contents.makeIterator()

    // 3
    return AnyIterator {
      return iterator.next()
    }
  }
}

In this revised Sequence extension, you:

1. Define Iterator as conforming to AnyIterator instead of DictionaryIterator . Then, as before, you create makeIterator() to return an Iterator .
2. Create iterator by calling makeIterator() on contents . You'll need this variable for the next step.
3. Wrap iterator in a new AnyIterator object to forward its next() method. The next() method is what is called on an iterator to get the next object in the sequence.

Press Command-Shift-Enter to run the playground. You'll notice a couple of errors:

Before, you were using the DictionaryIterator with tuple names of key and value . You've hidden DictionaryIterator from the outside world and renamed the exposed tuple names to element and count .

To fix the errors, replace key and value with element and count respectively. Run the playground now and your precondition blocks will pass just as they did before.

Now no one will know that you're just using Dictionary to the hard work for you!

It's time to bring your Bag home. OK, OK, collect your excitement, it's Collection time! :]

Adopting the Collection Protocol

Without further ado, here's the real meat of creating a collection: the Collection protocol! To reiterate, a Collection is a sequence that you can access by index and traverse multiple times.

To adopt Collection , you'll need to provide the following details:

• startIndex and endIndex : Defines the bounds of a collection and exposes starting points for transversal.
• subscript (position:) : Enables access to any element within the collection using an index. This access should run in O(1) time complexity.
• index(after:) : Returns the index immediately after the passed in index.

You're only four details away from having a working collection. You got this; it's in the Bag !

Add the following code just after the Sequence extension:

extension Bag: Collection {
  // 1
  typealias Index = DictionaryIndex<Element, Int>

  // 2
  var startIndex: Index {
    return contents.startIndex
  }

  var endIndex: Index {
    return contents.endIndex
  }

  // 3
  subscript (position: Index) -> Iterator.Element {
    precondition(indices.contains(position), "out of bounds")
    let dictionaryElement = contents[position]
    return (element: dictionaryElement.key,
      count: dictionaryElement.value)
  }

  // 4
  func index(after i: Index) -> Index {
    return contents.index(after: i)
  }
}

This is fairly straightforward. Here, you:

1. Declare the Index type defined in Collection as DictionaryIndex . You'll pass these indices through to contents .
2. Return the start and end indices from contents .
3. Use a precondition to enforce valid indices. You return the value from contents at that index as a new tuple.
4. Return the value of index(after:) called on contents .

By simply adding these properties and methods, you've created a fully functional collection!

Testing Your Collection

Add the following code to the end of the playground to test some of the new functionality:

// Get the first item in the bag
let firstItem = shoppingCart.first
precondition(
  (firstItem!.element == "Orange" && firstItem!.count == 1) ||
  (firstItem?.element == "Banana" && firstItem?.count == 2),
  "Expected first item of shopping cart to be (\"Orange\", 1) or (\"Banana\", 2)")

// Check if the bag is empty
let isEmpty = shoppingCart.isEmpty
precondition(isEmpty == false,
  "Expected shopping cart to not be empty")

// Get the number of unique items in the bag
let uniqueItems = shoppingCart.count
precondition(uniqueItems == 2,
  "Expected shoppingCart to have 2 unique items")

// Find the first item with an element of "Banana"
let bananaIndex = shoppingCart.indices.first { 
  shoppingCart[$0].element == "Banana"
}!
let banana = shoppingCart[bananaIndex]
precondition(banana.element == "Banana" && banana.count == 2,
  "Expected banana to have value (\"Banana\", 2)")

Once again, run the playground. Awesome!

Cue the moment where you're feeling pretty good about what you've done, but sense that a "but wait, you can do better" comment is coming... Well, you're right! You can do better. There's still some Dictionary smell leaking from your Bag .

Improving Collection

Bag is back to showing too much of its inner workings. Users of Bag need to use DictionaryIndex objects to access elements within the collection.

You can easily fix this. Add the following after the Collection extension:

// 1
struct BagIndex<Element: Hashable> {
  // 2
  fileprivate let index: DictionaryIndex<Element, Int>

  // 3
  fileprivate init(
    _ dictionaryIndex: DictionaryIndex<Element, Int>) {
    self.index = dictionaryIndex
  }
}

In the code above, you:

1. Define a new generic type, BagIndex . Like Bag , this requires a generic type that's Hashable for use with dictionaries.
2. Make the underlying data for this index type a DictionaryIndex object. BagIndex is really just a wrapper that hides its true index from the outside world.
3. Create an initializer that accepts a DictionaryIndex to store.

Now you need to think about the fact that Collection requires Index to be comparable to allow comparing two indexes to perform operations. Because of this, BagIndex needs to adopt Comparable .

Add the following extension just after BagIndex :

extension BagIndex: Comparable {
  static func ==(lhs: BagIndex, rhs: BagIndex) -> Bool {
    return lhs.index == rhs.index
  }

  static func <(lhs: BagIndex, rhs: BagIndex) -> Bool {
    return lhs.index < rhs.index
  }
}

The logic here is simple; you're using the equivalent methods of DictionaryIndex to return the correct value.

Updating BagIndex

Now you're ready to update Bag to use BagIndex . Replace the Collection extension with the following:

extension Bag: Collection {
  // 1
  typealias Index = BagIndex<Element>

  var startIndex: Index {
    // 2.1
    return BagIndex(contents.startIndex)
  }

  var endIndex: Index {
    // 2.2
    return BagIndex(contents.endIndex)
  }

  subscript (position: Index) -> Iterator.Element {
    precondition((startIndex ..< endIndex).contains(position),
      "out of bounds")
    // 3
    let dictionaryElement = contents[position.index]
    return (element: dictionaryElement.key,
      count: dictionaryElement.value)
  }

  func index(after i: Index) -> Index {
    // 4
    return Index(contents.index(after: i.index))
  }
}

Each numbered comment marks a change. Here's what they are:

1. Replaces the Index type from DictionaryIndex to BagIndex .
2. Creates a new BagIndex from contents for both startIndex and endIndex .
3. Uses the index property of BagIndex to access and return an element from contents .
4. Gets the DictionaryIndex value from contents using the property of BagIndex and creates a new BagIndex using this value.

That's it! Users are back to knowing nothing about how you store the data. You also have the potential for much greater control of index objects.

Before wrapping this up, there's one more important topic to cover. With the addition of index-based access, you can now index a range of values in a collection. Time for you to take a look at how a slice works with collections.

Using Slices

A slice is a view into a subsequence of elements within a collection. It lets you perform actions on a specific subsequence of elements without making a copy.

A slice stores a reference to the base collection you create it from. Slices share indices with their base collection, keeping references to the start and end indices to mark the subsequence range. Slices have an O(1) space complexity because they directly reference their base collection.

To see how this works, add the following code to the end of the playground:

// 1
let fruitBasket = Bag(dictionaryLiteral:
  ("Apple", 5), ("Orange", 2), ("Pear", 3), ("Banana", 7))

// 2
let fruitSlice = fruitBasket.dropFirst()

// 3
if let fruitMinIndex = fruitSlice.indices.min(by:
  { fruitSlice[$0] > fruitSlice[$1] }) {
  // 4
  let basketElement = fruitBasket[fruitMinIndex]
  let sliceElement = fruitSlice[fruitMinIndex]
  precondition(basketElement == sliceElement,
    "Expected basketElement and sliceElement to be the same element")
}

Run the playground again.

In the code above, you:

1. Create a fruit basket made up of four different fruits.
2. Remove the first type of fruit. This actually just creates a new slice view into the fruit basket excluding the first element you removed, instead of creating a whole new Bag object. You'll notice in the results bar that the type here is Slice<Bag<String>> .
3. Find the index of the least occurring fruit in those that remain.
4. Prove that you're able to use the index from both the base collection as well as the slice to retrieve the same element, even though you calculated the index from the slice.

Note : Slices may seem a little less useful for hash-based collections like Dictionary and Bag because their order isn't defined in any meaningful way. An Array , on the other hand, is an excellent example of a collection type where slices play a huge role in performing subsequence operations.

Congratulations — you're now a collection pro! You can celebrate by filling your Bag with your own custom prizes. :]

Where to Go From Here?

You can download the complete playground with all the code in this tutorial using the Download Materials button at the top or bottom of the tutorial.

In this tutorial, you learned how to make a custom collection in Swift. You added conformance to Sequence , Collection , CustomStringConvertible , ExpressibleByArrayLiteral , ExpressibleByDictionaryLiteral and you created your own index type.

If you'd like to view or contribute to a more complete Bag implementation, check out the Swift Algorithm Club implementation as well as the Foundation implementation, NSCountedSet .

These are just a taste of all the protocols Swift provides to create robust and useful collection types. If you'd like to read about some not covered here, check out the following:

• Array Relationships
• Dictionary Relationships
• BidirectionalCollection Relationships

You can also check out more information about Protocols in Swift and learn more about adopting common protocols available in the Swift standard library.

Finally, be sure to read our article on Protocol-Oriented Programming in Swift !

I hope you enjoyed this tutorial! Building your own custom collection definitely comes in handy, and gives you a better understanding of Swift's standard collection types.

If you have any comments or questions, feel free to join in the forum discussion below!