Weekly Swift articles, podcasts and tips by John Sundell.

Wrapping sequences in Swift

Published on 12 May 2019

One major benefit of Swift’s protocol-oriented design is that it enables us to write generic code that’s compatible with a wide range of types, rather than being specifically implemented for just a single one. Especially if such generic code targets one of the protocols found within the standard library — which’ll make it possible to use it both with built-in types, as well as user-defined ones.

An example of such a protocol is Sequence, which is adopted by all standard library types that can be iterated over, such as Array, Dictionary, Set, and many more. This week, let’s take a look at how we can wrap Sequence in generic containers, that’ll let us encapsulate various algorithms behind easy-to-use APIs.

The art of being lazy

It’s quite easy to fall into the trap of thinking that all sequences are like Array, in that all of the elements are instantly loaded into memory when the sequence is created. However, since the Sequence protocol’s only requirement is that adopters should be capable of being iterated over, we can’t really make any assumptions about how an unknown sequence’s elements are loaded or stored.

For example, like we took a look at in “Swift sequences: The art of being lazy”, sequences can sometimes load their elements lazily — either for performance reasons, or because the entire sequence isn’t guaranteed to fit into memory. Here are some examples of sequences that might do that:

// A sequence of database records, in which pages of records are
// loaded as an iteration goes on, to avoid loading all search
// results into memory in order to iterate over them.
let records = database.records(matching: searchQuery)

// A sequence of subfolders within a folder on disk, in which each
// folder is only opened once the iteration has reached it.
let folders = folder.subfolders

// A sequence of nodes within a graph, which is lazy so that we
// don't need to evaluate the entire graph at the beginning of
// each iteration.
let nodes = node.children

Since all of the above sequences are lazy for a reason, we wouldn’t want to force them into an array, for example by calling Array(folder.subfolders). But we still might want to modify and work with such sequences in different ways — so let’s take a look at how we could do just that by creating a sequence wrapper type.

Building a base

Let’s start by building a base type that we’ll be able to use to create all sorts of convenience APIs on top of any sequence. We’ll name it WrappedSequence, an it’ll be a generic containing both the type of the sequence that we’re wrapping, as well as the type of Element that we want our new sequence to produce.

The core feature of our wrapper will be its IteratorFunction, which’ll enable us to take control over the underlying sequence’s iterations — by mutating the Iterator used for each iteration:

struct WrappedSequence<Wrapped: Sequence, Element>: Sequence {
    typealias IteratorFunction = (inout Wrapped.Iterator) -> Element?

    private let wrapped: Wrapped
    private let iterator: IteratorFunction

    init(wrapping wrapped: Wrapped,
         iterator: @escaping IteratorFunction) {
        self.wrapped = wrapped
        self.iterator = iterator
    }

    func makeIterator() -> AnyIterator<Element> {
        var wrappedIterator = wrapped.makeIterator()
        return AnyIterator { self.iterator(&wrappedIterator) }
    }
}

As you can see above, Sequence uses the factory pattern, to have each sequence create a new iterator instance for each iteration — through the makeIterator() method.

Above we use the standard library’s AnyIterator type, which is a type-erased iterator that can use any underlying IteratorProtocol implementation to produce Element values. In our case, we’ll produce an element by calling our IteratorFunction, passing the wrapped sequence’s own iterator as an argument — and since that argument as marked as inout, we’re able to mutate the underlying iterator in-place within our function.

Since WrappedSequence is also a sequence just like any other, we can use all of the standard library’s Sequence-bound functionality with it — such as iterating over it, or transforming its values using map:

let folderNames = WrappedSequence(wrapping: folders) { iterator in
    return iterator.next()?.name
}

for name in folderNames {
    ...
}

let uppercasedNames = folderNames.map { $0.uppercased() }

Now, let’s start taking our new WrappedSequence for a spin!

Prefixing and suffixing

When working with sequences, it’s very common to want to insert either a prefix or a suffix into the sequence that we’re working with — but wouldn’t it be great if we’d be able to do that without mutating the underlying sequence? That could both lead to better performance, and would enable us to add prefixes and suffixes to any sequence — not just to common types, like Array.

Using WrappedSequence, we can do just that quite easily. All we have to do is to extend Sequence with a method that produces a wrapped sequence, given an array of elements to insert as a prefix. Then, when we’re performing an iteration, we’ll start by iterating over all of the prefixed elements — before proceeding with the underlying sequence — like this:

extension Sequence {
    func prefixed(
        with prefixElements: Element...
    ) -> WrappedSequence<Self, Element> {
        var prefixIndex = 0

        return WrappedSequence(wrapping: self) { iterator in
            // If we still have prefixed elements left to serve,
            // then return the next one by incrementing our index:
            guard prefixIndex >= prefixElements.count else {
                let element = prefixElements[prefixIndex]
                prefixIndex += 1
                return element
            }

            // Otherwise, return an element from our underlying
            // sequence's own iterator:
            return iterator.next()
        }
    }
}

Above we use a variadic parameter (by adding ... to its type), to enable either a single element — or multiple ones — to be passed to the same method.

Similarly, we can also create a method that adds a given set of suffixes to the end of a sequence — by first completing the underlying sequence’s own iteration, and then iterating over the suffixed elements:

extension Sequence {
    func suffixed(
        with suffixElements: Element...
    ) -> WrappedSequence<Self, Element> {
        var suffixIndex = 0

        return WrappedSequence(wrapping: self) { iterator in
            guard let next = iterator.next() else {
                // This is our exit condition, in which we return
                // nil after both the underlying iteration, and
                // the suffixed one, have been completed:
                guard suffixIndex < suffixElements.count else {
                    return nil
                }

                let element = suffixElements[suffixIndex]
                suffixIndex += 1
                return element
            }

            return next
        }
    }
}

With the above two methods in place, we can now add prefixes and suffixes to any sequence that we want. Here are some example of how our new APIs may be used:

// Including the parent folder in a subfolder iteration:
let allFolders = rootFolder.subfolders.prefixed(with: rootFolder)

// Append a draft as a suffix to a message sequence:
let messages = inbox.messages.suffixed(with: composer.message)

// Enclosing a string in brackets before iterating over it,
// without having to create any new string instances:
let characters = code.prefixed(with: "{").suffixed(with: "}")

While all of the above examples could be implemented using concrete types (like Array and String), the benefit of using our WrappedSequence type is that everything can be done lazily — we don’t perform any mutations or evaluate any sequences in order to add our prefixes and suffixes — which can be really beneficial in performance-critical situations, or when working with large datasets.

Segmentation

Next, let’s take a look at how we can wrap sequences in order to create segmented versions of them. In certain iterations, it’s not enough to know what the current element is — we might also need to be aware of the next and previous ones.

When working with indexed sequences, we can often accomplish that by using the enumerated() API — which also uses sequence wrapping to give us access to both the current element, as well as its index:

for (index, current) in list.items.enumerated() {
    let previous = (index > 0) ? list.items[index - 1] : nil
    let next = (index < list.items.count - 1) ? list.items[index + 1] : nil
    ...
}

However, the above technique isn’t only quite verbose at the call site, it also once again relies on using arrays — or at least some form of sequence that gives us random access to its elements — which many sequences, especially lazy ones, often don’t.

Instead, let’s use our WrappedSequence once more — to create a wrapping sequence that lazily provides segmented views into its underlying sequence, by keeping track of the previous and current elements, and updating those as the iteration continues:

extension Sequence {
    typealias Segment = (
        previous: Element?,
        current: Element,
        next: Element?
    )

    var segmented: WrappedSequence<Self, Segment> {
        var previous: Element?
        var current: Element?
        var endReached = false

        return WrappedSequence(wrapping: self) { iterator in
            // Here our exit condition is either that we've
            // reached the end of the underlying sequence, or
            // that a first current element couldn't be created,
            // because the sequence was empty.
            guard !endReached,
                  let element = current ?? iterator.next() else {
                return nil
            }

            let next = iterator.next()
            let segment = (previous, element, next)

            // Before we return the new segment, we update our
            // iteration state to be ready for the next element:
            previous = element
            current = next
            endReached = (next == nil)

            return segment
        }
    }
}

We can now use the above API to create a segmented version of any sequence whenever we need to either look ahead or behind when performing an iteration. For example, here’s how we could use segmentation to be able to easily determine when we’ve reached the end of a list:

for segment in list.items.segmented {
    addTopBorder()
    addView(for: segment.current)

    if segment.next == nil {
        // We've reached the end, time to add a bottom border
        addBottomBorder()
    }
}

Segmented sequences are also great for visualizing maps, paths, or graphs. Here we’re using segments to be able to compute, and render, the enter and exit directions for each node within a path:

for segment in path.nodes.segmented {
    let directions = (
        enter: segment.previous?.direction(to: segment.current),
        exit: segment.next.map(segment.current.direction)
    )

    let nodeView = NodeView(directions: directions)
    nodeView.center = segment.current.position.cgPoint
    view.addSubview(nodeView)
}

Now we’re starting to see the true power of wrapping sequences — in that they let us hide increasingly complex algorithms behind a really simple API. All the caller needs to do to segment a sequence is to access the segmented property on any Sequence, and our underlying implementation takes care of the rest.

Recursion

Finally, let’s take a look at how even recursive iterations can be modeled by wrapping sequences. Let’s say that we wanted to provide an easy way to recursively iterate over a hierarchy of values — in which each element in the hierarchy contains a sequence of child elements. This is something that can be quite tricky to get right, so it’d be great if we could use a single implementation to perform all such iterations within our code base.

Using WrappedSequence, we can make that happen by extending Sequence with a method that uses a same type generic constraint to ensure that each element is able to provide a nested sequence that has the same type of iterator as our original one. To be able to dynamically access each nested sequence, we’ll also ask the caller to supply the KeyPath for the property that should be used for recursion — giving us an implementation that looks like this:

extension Sequence {
    func recursive<S: Sequence>(
        for keyPath: KeyPath<Element, S>
    ) -> WrappedSequence<Self, Element> where S.Iterator == Iterator {
        var parentIterators = [Iterator]()

        func moveUp() -> (iterator: Iterator, element: Element)? {
            guard !parentIterators.isEmpty else {
                return nil
            }

            var iterator = parentIterators.removeLast()

            guard let element = iterator.next() else {
                // We'll keep moving up our chain of parents
                // until we find one that can be advanced to
                // its next element:
                return moveUp()
            }

            return (iterator, element)
        }

        return WrappedSequence(wrapping: self) { iterator in
            // We either use the current iterator's next element,
            // or we move up the chain of parent iterators in
            // order to obtain the next element in the sequence:
            let element = iterator.next() ?? {
                return moveUp().map {
                    iterator = $0
                    return $1
                }
            }()

            // Our recursion is performed depth-first, meaning
            // that we'll dive as deep as possible within the
            // sequence before advancing to the next element on
            // the level above.
            if let nested = element?[keyPath: keyPath].makeIterator() {
                let parent = iterator
                parentIterators.append(parent)
                iterator = nested
            }

            return element
        }
    }
}

Using the above, we’re now able to recursively iterate over any sequence, no matter how it’s constructed internally — and without having to load the entire hierarchy up-front. For example, here’s how we could use this new API to recursively iterate over a folder hierarchy:

let allFolders = folder.subfolders.recursive(for: \.subfolders)

for folder in allFolders {
    try loadContent(from: folder)
}

We could also use it to iterate over all nodes within a tree, or to recursively traverse a set of database records — for example to enumerate all user groups within an organization:

let allNodes = tree.recursive(for: \.children)
let allGroups = database.groups.recusive(for: \.subgroups)

The one thing we have to be careful of when it comes to recursive iterations is to not end up with circular references — when a certain path leads us back to an element we already encountered — which will cause us to keep iterating through that loop over and over.

One way to fix that would be to keep track of all encountered elements (but that could be problematic in terms of memory), to ensure that no circular references exists within our dataset, or to handle such cases at each call site (by using break, continue, or return to end any circular iterations).

Conclusion

Sequence is one of the simplest protocols in the standard library — it only has a single method requirement — but it’s still one of the most powerful ones, especially when it comes to how much functionality we can create on top of it. Just like how the standard library contains wrapping sequences for things like enumerations, we can create our own wrappers as well — that let us hide advanced functionality behind really simple APIs.

While abstractions are never free, and it’s important to consider when (and perhaps more importantly — when not) to introduce one, if we can build our abstractions right on top of what the standard library provides — using the same conventions — then those abstractions usually have a higher chance of standing the test of time.

What do you think? Have you created your own wrapping sequence for some kind of algorithm, or is it something you’ll try out? Let me know — along with your questions, comments, and feedback — on Twitter @johnsundell.

Thanks for reading! 🚀