Building a Custom Collection in Swift

Arrays, dictionaries, and sets are all examples of commonly used collection types; they come bundled with the Swift standard library.

But what if they don’t provide everything you need for your application right out of the box?

One common solution is to use an Array or Dictionary with a bunch of business logic to keep your data organized. This can be problematic as it’s often unintuitive and hard to maintain.

That’s where creating your own collection type becomes relevant. In this article, you’ll create a custom collection using Swift’s powerful collection protocols.

In the end, you’ll have a useful custom collection data type with all the functionality of a Swift provided collection.

Getting Started

In this tutorial, you’re going to be building a multiset (bag) from scratch.

A bag is like a set in that it stores objects with no repeated values. In a set, duplicate object are ignored. A bag, on the other hand, keeps a running count for each object.

A great example of where this would come in handy is a shopping list. You’d want a list of unique grocery items with an associated quantity for each. Rather than adding duplicate item instances, you’d expect to increment the existing item’s quantity.

Before jumping into the collection protocols, you’ll first create the basic implementation of a Bag.

First, create a new playground: in Xcode, select File\New\Playground… and name the playground Bag. You can select either platform since this tutorial is platform-agnostic and only focuses on the Swift language.

Click Next, choose a convenient location to save the playground and click Create.

Next, replace the contents of the playground with an empty implementation of a bag:

struct Bag<Element: Hashable> {
 
}

Bag is a generic structure that requires an element type that is Hashable. Requiring Hashable elements allows you to compare and only store unique values at O(1) time complexity. This means that no matter the size of its contents, Bag will perform at constant speeds. You used a struct to enforce value semantics as found with Swift’s standard collections.

Next, add the following properties to Bag:

// 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. You’re using 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 dictionary value for each element is its count. It’s marked 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 of 4 oranges and 2 apples would return a uniqueCount of 2.
3. totalCount returns the total number of items in the bag. In the same example as before, totalCount would return 6.

Now you’ll need 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. The mutating keyword is used to let variables be modified in a struct or enum. These methods will not be available if the instance is defined as a constant let rather than var.
2. precondition(_:_:) takes a Boolean value as its first parameter. If false, execution of the program will stop and the String from the second parameter will be output in the Debug area. You’ll use preconditions a number of times in this tutorial to ensure Bag is used in its intended way. You’ll also use it as a sanity check to make sure things work as expected as you add functionality.
3. This checks if the element already exists in the bag. If it does, increment the count, if not, create a new element.

On the flip side, you’ll need a way to remove elements from Bag. 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 {
    preconditionFailure("Removed non-existent elements")
  }
 
  // 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)
  }
}

remove(_:occurrences:) takes the same parameters as add(_:occurrences:) and does the opposite with them. Here’s how it works:

1. First it checks that the element exists and that it has at least the number of occurrences as being removed.
2. Next it makes sure that the number of occurrences to remove is greater than 0.
3. Finally, it checks if the element exists and decrements the count. If the count drops to zero, it removes the element entirely.

At this point Bag doesn’t do much; its contents aren’t even accessible once added. You also lose the useful higher-order methods available on Dictionary.

But all is not lost. You’ve started the process of separating all this wrapper code into it’s own object. That’s one step in the right direction of keeping your code clean!

But wait, there’s more! Swift provides all the tools you need to make Bag into a legitimate collection.

To do this you’ll need to look at what makes an object a collection in Swift.

What’s a Custom Collection?

To understand what a Swift Collection is, you first need to look at its protocol inheritance hierarchy.

The Sequence protocol 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.

There are way too many Pokemon to keep track these days

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.
• reduce(_:_:): Returns a single value by combining each element in the sequence using the provided closure.
• 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 the Sequence docs.

One caveat to Sequence is that it makes no requirement for conforming types to be destructive or not. This means that after iteration, there’s no guarantee that future iterations will start from the beginning.

Thats a huge issue if you plan on iterating over your data more than once. To enforce nondestructive iteration, your object needs to conform to the Collection protocol.

Collection inherits from Sequence and Indexable. 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. If you’d like a sneak peek for yourself, check out the Collection docs.

Before implementing these protocols, there are some easy improvements you can make to Bag.

Textual Representation

Currently, Bag objects expose little information through print(_:) or the results sidebar.

Add the following code to the end of the playground to see for yourself:

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

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

You can fix this using a single protocol provided by the Swift standard library. Add the following just after the closing brace of Bag above shoppingCart:

extension Bag: CustomStringConvertible {
  var description: String {
    return contents.description
  }
}

Conforming to CustomStringConvertible requires that you implement 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 reuse the value of description from contents.

Take a look at the previously useless debug information for shoppingCart:

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

While you’re writing code in a playground, you can verify a result you expect by calling precondition(_:_:). Doing so incrementally will also keep you from accidentally breaking functionality that was previously working. You can use this tool just as you would with unit tests — it would be a great idea to incorporate this into your every day coding!

Add the following just after the last call to add(_:occurrences:) at the end of the playground:

precondition("\(shoppingCart)" == "\(shoppingCart.contents)", "Expected bag description to match its contents description")

This will result in an error if the description of shoppingCart diverges from contents.

Next up for creating powerful collection types that feel native is initialization.

Initialization

It’s pretty annoying that you have to add each element one at a time. One common expectation is to be able to initialize collection types with other collections.

Here’s how you might expect to create a Bag. Add the following code to the end of the playground:

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 won’t compile because you haven’t defined any initializers for these types. 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)
  }
}

Let’s take a look at what you just added:

1. First, you created an empty initialization method. You’re required to add this after defining additional init methods to avoid compiler errors.
2. Next, you added an initialization method that accepts any Sequence of elements. This sequence must have a matching Element type. This, for example, covers both Array and Set objects. You iterate over the passed in sequence and add each element one at a time.
3. The last method works similarly, but for tuple elements 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.

These generic initializers enable a much wider variety of data sources for Bag objects. They do, however, need you to initialize another sequence simply to pass to Bag.

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

Add the following code to the end of your playground to see an example of how this is used:

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

Again, you’ll see compiler errors which you’ll fix next. This is an example of initialization using Array and Dictionary literals rather than objects.

Add the following two extensions just below the other Bag extension:

extension Bag: ExpressibleByArrayLiteral {
  init(arrayLiteral elements: Element...) {
    self.init(elements)
  }
}
 
extension Bag: ExpressibleByDictionaryLiteral {
  init(dictionaryLiteral elements: (Element, Int)...) {
    // The map converts elements to the "named" tuple the initializer expects.
    self.init(elements.map { (key: $0.0, value: $0.1) })
  }
}

Both ExpressibleByArrayLiteral and ExpressibleByDictionaryLiteral require an initializer that handles their matching literal parameter. These were incredibly easy to implement due to the previous initializers you added.

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

Sequence

By far the most common action performed on a collection type is iterating through its elements. To see an example of this, add the following to the end of the playground:

for element in shoppingCart {
  print(element)
}

Super basic stuff here. 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.

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. Let’s look at what you just added:

1. You defined a typealias named Iterator that Sequence defines as conforming to IteratorProtocol. 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. makeIterator() returns an Iterator that can step through each element of the sequence.
3. You create an iterator by calling makeIterator() on contents, which already 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)")
}

Open the Debug area and you’ll see the printout of the elements in the sequence.

Being able to iterate through 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 counts
let itemList = shoppingCart.map { $0.0 }
itemList
precondition(itemList == ["Orange", "Banana"], "Expected itemList contents to be [\"Orange\", \"Banana\"]")
 
// 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 decending 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)]")

These are all useful methods for working with sequences — and they were given to you 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 and dandy 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. Creating your own iterator type is definitely possible, but fortunately not needed.

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()
    }
  }
}

The playground will show a couple errors which you’ll fix soon. This looks close to the previous implementation with the addition of AnyIterator:

1. AnyIterator is a type-erased iterator that forwards its next() method to an underlying iterator. This allows you to hide the actual iterator type used.
2. Just like before, you create a new DictionaryIterator from contents.
3. Finally, you wrap iterator in a new AnyIterator object to forward its next() method.

Now to fix the errors. You’ll see the following two errors:

Before, you were using the DictionaryIterator tuple names 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.

Your preconditions should now pass and work just as they did before. This is why preconditions are awesome at making sure things don’t change unexpectedly.

Now no one will know that you’re just using a dictionary to do everything for you.

Now that you’re feeling better about Bag, it’s time to bring it home. Ok, ok, collect your excitement, it’s Collection time! :]

Collection

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

To add Collection conformance, you’ll need to provide the following details:

• startIndex and endIndex: Define the bounds of a collection and expose starting points for transversal.
• subscript (position:): Lets you to access 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. 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((startIndex ..< endIndex).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:

1. First, you declare the Index type defined in Collection as DictionaryIndex. You’ll pass these indices to contents. Note that the compiler could infer this type based on the rest of your implementation. Explicitly defining it keeps the code clean and maintainable.
2. Next, you return the start and end indices from contents.
3. Here, you use a precondition to enforce valid indices. You return the value from contents at that index as a new tuple.
4. Finally, you echo the value of index(after:) called on contents.

By simply adding these properties and methods, you’ve created a fully functional 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, "Expected first item of shopping cart to be (\"Orange\", 1)")
 
// 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)")

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 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 just after the Collection extension:

// 1
struct BagIndex<Element: Hashable>: Comparable {
  // 2
  fileprivate var index: DictionaryIndex<Element, Int>
 
  // 3
  fileprivate init(_ dictionaryIndex: DictionaryIndex<Element, Int>) {
    self.index = dictionaryIndex
  }
}

There’s nothing too strange here, but let’s walk through what you just added:

1. You’re defining a new generic type, BagIndex, that conforms to Comparable. Collection requires Index to be comparable to allow comparing two indexes to perform operations. Like Bag, this requires a generic type that’s Hashable for use with dictionaries.
2. The underlying data for this index type is a fileprivate DictionaryIndex object. BagIndex is really just a wrapper that hides its true index from the outside world.
3. Last, you create a fileprivate initializer that accepts a DictionaryIndex to store.

Because BagIndex conforms to Comparable, add the following methods just after BagIndex to fix the errors:

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

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

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. First you replaced the Index type from DictionaryIndex to BagIndex.
2. Next, for both startIndex and endIndex, you created a new BagIndex from contents.
3. Then you used the property of BagIndex to access and return an element from contents.
4. Finally, you used a combination of the previous steps. You got the DictionaryIndex value from contents using the property of BagIndex. Then you created 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.

To finish off, you’ll take a look at how a Slice works with collections.

Slice

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 reference to the base collection it was created from. It also keeps reference to the start and end indices to mark the subsequence range. Slices maintain a O(1) complexity because they directly reference their base collection.

Slices share indices with their base collection which make them incredibly useful.

To see how this works in practice, 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() // No pun intended ;]
 
// 3
if let fruitMinIndex = fruitSlice.indices.min(by: { fruitSlice[$0] > fruitSlice[$1] }) {
  // 4
  let minFruitFromSlice = fruitSlice[fruitMinIndex]
  let minFruitFromBasket = fruitBasket[fruitMinIndex]
}

Let’s take a look at what’s going on here and why it’s significant:

1. First you create a fruit basket made up of four different fruits.
2. Next, you remove the first type of fruit to eat. This creates a new Slice view into the fruit basket excluding the first element you removed.
3. With the remaining fruit in the slice, you find the index of the least occurring fruit.
4. Although you calculated the index from a slice, you’re able to use it successfully from both the base collection as well as the slice.

Congratulations — you’re a collection pro now! You can celebrate by making a bag of whatever prizes you want. :]

Where to Go From Here?

You can download the complete playground with all the code in this article here. If you’d like to view or contribute to a more complete Bag implementation, check it out on github.

In this article you learned what makes a data structure a collection in Swift by creating your own. You added conformance to Sequence, Collection, CustomStringConvertible, ExpressibleByArrayLiteral, ExpressibleByDictionaryLiteral as well as creating your own index type.

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 that weren’t covered here, check out the following:

• Array’s protocol hierarchy: Array Hierarchy
• Dictionary’s protocol hierarchy: Dictionary Hierarchy

I hope you enjoyed this tutorial! Building your own custom collection might not be the most common requirement, but hopefully it’s given you a better understanding of Swift’s provided collection types.

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

The post Building a Custom Collection in Swift appeared first on Ray Wenderlich.