objc.io

00:06 Today we'll do another episode about collections, the third in the series. Before, we looked at implementing custom extensions on the collection types and protocols, and today, we'll implement a simple custom collection.

00:33 Let's look at our example. We have a Calendar struct where we want to store some events. The invariant we'll use is that the events are sorted by start date. We could add a comment to the property to document this:

struct Event {
    var startDate: Date
}

struct Calendar {
    // INVARIANT: sorted by start date
    var events: [Event] = []
}

01:07 To insert an event, we can add a method. The simplest way is to append it to events and then sort the entire array. This isn't very efficient, but it does the job for the example:

mutating func insert(_ event: Event) {
    events.append(event)
    events.sort { e1, e2 in
        e1.startDate < e2.startDate
    }
}

A Custom Struct

02:00 It would be much nicer if we didn't have to pay attention to the sorting in every mutating method of our calendar. This way, if a method changes something in the events, it doesn't have to worry about keeping the events sorted. We could pull the sorting logic out of the Calendar struct into its own data type and make it into a custom collection called SortedArray. It's good to separate the sorted array from the calendar because this way, we can add more logic to our Calendar struct without having to worry about the invariant:

struct Calendar {
    var events: SortedArray<Event> = []

    // ...
}

02:40 We'd like to have a struct, SortedArray. It'll be generic over Element, just like Array. Inside, we can have a property, elements. Our struct is simply a wrapper around Array; we won't implement much functionality ourselves. In the insert method, we do something similar to what happened in the Calendar: we append and sort. We need a function that can compare two elements — areInIncreasingOrder — which we'll store as a property in the struct:

struct SortedArray<Element> {
    var elements: [Element]
    let areInIncreasingOrder: (Element, Element) -> Bool

    mutating func insert(_ element: Element) {
        elements.append(element)
        elements.sort(by: areInIncreasingOrder)
    }
}

04:20 Now we have to see that we can initialize our struct with an existing sequence of elements and a function that compares two elements. We already get a memberwise initializer, but it'll only work correctly if the elements are already sorted. Instead, we'll write a custom initializer. We can easily make our initializer work with a sequence by constraining the sequence to make sure the Iterator.Element is the same as our generic type, Element:

init<S: Sequence>(unsorted: S, areInIncreasingOrder: @escaping (Element, Element) -> Bool) where S.Iterator.Element == Element {
    elements = unsorted.sorted(by: areInIncreasingOrder)
    self.areInIncreasingOrder = areInIncreasingOrder
}

06:42 By putting the custom initializer inside the struct definition, the memberwise initializer is gone. Now we can create a sorted array and insert something:

var x = SortedArray(unsorted: [3,1,2], areInIncreasingOrder: <)
x.insert(-1)
x.elements // [-1, 1, 2, 3]

Conforming to Collection

07:46 Once we finish our struct, we want to hide elements as a private implementation detail. Next, we make SortedArray conform to the Collection protocol so that we can iterate over the elements and use methods like count, map, and filter. We start by writing an extension on SortedArray. In order to conform to Collection, we need four things: a startIndex, an endIndex, a subscript, and a method called index(after:_), which advances an index.

08:43 We can implement these four methods by forwarding the call to the backing array. That will make our collection simple:

extension SortedArray: Collection {
    var startIndex: Int {
        return elements.startIndex
    }

    var endIndex: Int {
        return elements.endIndex
    }

    subscript(index: Int) -> Element {
        return elements[index]
    }

    func index(after i: Int) -> Int {
        return elements.index(after: i)
    }
}

10:05 That was straightforward, and we now have access to all the functionality that Collection and Sequence provide:

var x = SortedArray(unsorted: [3,1,2])
x.insert(-1)
x.elements

for y in x {
    print(y)
}
x.count // 4

10:36 There are many methods we get by conforming to Collection. Most of these methods are on Sequence, a parent protocol. We didn't have to conform to Sequence explicitly; we got the conformance for free by conforming to Collection. If we say x.makeIterator(), we can see that we get an IndexingIterator<SortedArray<Int>> as the type. Once you implement the four methods above and conform to the Collection protocol, the Collection protocol will provide you with things — such as a default iterator — for free. You could still write your own, but you don't have to. At first, it's surprising we have to conform to Sequence, because the things we've implemented are so different. But out of these four methods, everything for Sequence can be provided for us.

Overriding Default Methods

11:32 Another interesting method we get from the Collection conformance is min(). This will return the minimum element within the collection, and we can override this for SortedArray. The default implementation doesn't know the elements are sorted, so it'll loop over all the elements and return the minimum value. Instead, we could be more efficient by returning the first element in our sorted array:

extension SortedArray {
    func min() -> Element? {
        return elements.first
    }

}

A Custom Initializer

12:36 There are some additional improvements we could make. For example, if a type is already Comparable, we could add an initializer so that we don't have to provide a comparison function. Again, we'll pass in an unsorted sequence:

extension SortedArray where Element: Comparable {
    init<S: Sequence>(unsorted: S) where S.Iterator.Element == Element {
        self.init(unsorted: unsorted, areInIncreasingOrder: <)
    }
}

14:13 Now we can write the following code, because Ints already conform to Comparable:

var x = SortedArray(unsorted: [3,1,2])

14:24 If we have something that doesn't conform, we can still use the original initializer with an explicit comparison function. Or if we want our Ints sorted in the reverse direction, we could pass in the > operator as our comparison function.

14:41 Of course, our implementation is naive and inefficient. In a real implementation, you'd insert the elements immediately at the correct place instead of appending and sorting. However, this example demonstrates it can be useful to create a custom collection: we can encapsulate the fact that the array is sorted.

15:06 There are some other protocols we could also conform to: we could conform to RangeReplaceCollection, but it wouldn't make much sense, because then we could replace a range with random integers. A simpler example is MutableCollection, which allows us to assign new values to existing indices. But then we'd immediately need to invalidate the index, because the new value will probably be stored at a different index. However, we can conform to BidirectionalCollection so that we get a cheap reversed(), and we can conform to RandomAccessCollection as well. There might also be some more methods, such as min() and max(), that we could override.

16:22 In future episodes, we'll have a look at creating a completely custom collection instead of just a simple wrapper.