Splitting a Swift Sequence into head and tail

A common pattern in functional programming languages is to split a list into its head (the first element) and tail (all remaining elements). For example, in Haskell, the pattern x:xs matches against a non-empty list and binds the list’s head to the variable x and the tail to xs.

Swift is not a functional language. It neither has a built-in List type, nor a special pattern matching syntax for collections.1

Collections

Nonetheless, splitting a Sequence or Collection into head and tail can occasionally be useful. For collections, this is straightforward:

extension Collection {
    var headAndTail: (head: Element, tail: SubSequence)? {
        guard let head = first else { return nil }
        return (head, dropFirst())
    }
}

if let (firstLetter, remainder) = "Hello".headAndTail {
    // firstLetter: Character == "H"
    // remainder: Substring == "ello"
}

Sequences

For sequences it’s trickier because they are allowed to be single-pass: some sequences can only be iterated over once as iteration consumes their elements. Think of a network stream as an example. Once you read a byte from the buffer, the operating system throws it away. You can’t tell it to start over.

One possible solution is to create an iterator to read in the first element, and then wrap the current iterator state in a new AnySequence instance:

extension Sequence {
    var headAndTail: (head: Element, tail: AnySequence<Element>)? {
        var iterator = makeIterator()
        guard let head = iterator.next() else { return nil }
        let tail = AnySequence { iterator }
        return (head, tail)
    }
}

This code works, but it’s not a nice generic solution, especially for types that also conform to Collection. Wrapping the tail in an AnySequence is a big performance killer, and you can’t use the affordances of a collection’s proper SubSequence type.

You’d be better off writing two extensions, one for Collection and one for Sequence, in order to preserve the SubSequence type for collections. (And as we’ll see, this is also the preferred solution from Swift 5 onward, but I’ll get to that.)

Preserving the SubSequence type

I couldn’t figure out a generic solution for Sequence that kept the SubSequence type for the tail intact and worked correctly with single-pass sequences. I thank Dennis Vennink for coming up with a solution and sharing it with me. Here’s his code (slightly modified by me for style):

extension Sequence {
    var headAndTail: (head: Element, tail: SubSequence)? {
        var first: Element? = nil
        let tail = drop(while: { element in
            if first == nil {
                first = element
                return true
            } else {
                return false
            }
        })
        guard let head = first else {
            return nil
        }
        return (head, tail)
    }
}

Dennis’s trick is to call Sequence.drop(while:), which preserves the SubSequence type for the tail, and then “catch” the first element inside the drop(while:) predicate closure using a captured local variable. Nicely done!

Swift 5

The code above targets Swift 4.2. It will break in Swift 5 because sequences will no longer have an associated SubSequence type, only collections (Swift Evolution proposal SE-0234).2

This change has many advantages, but it means we’re no longer able to write generic algorithms that make use of SubSequence and work on Sequence and Collection at the same time.

Instead, we can add the straightforward solution to Collection:

extension Collection {
    var headAndTail: (head: Element, tail: SubSequence)? {
        guard let head = first else { return nil }
        return (head, dropFirst())
    }
}

If we find that we need the same functionality for Sequence as well, we should add a separate extension that uses the new DropWhileSequence type as the return type for the tail:

extension Sequence {
    var headAndTail: (head: Element, tail: DropWhileSequence<Self>)? {
        var first: Element? = nil
        let tail = drop(while: { element in
            if first == nil {
                first = element
                return true
            } else {
                return false
            }
        })
        guard let head = first else {
            return nil
        }
        return (head, tail)
    }
}

(The implementation is the same as above. Only the return type has changed.)

  1. Introducing a pattern matching construct for collections has been mentioned as a possible feature on the forums several times. With it, you’d be able to write something like this to destructure an ordered collection into head and tail:

    let numbers = 1...10
    let [head, tail...] = numbers
    // head == 1
    // tail == 2...10
    

    This would be very handy in switch statements. ↩︎

  2. It’s unfortunate that we’re now stuck with the misleading name SubSequence. Renaming Collection.SubSequence to the more appropriate Slice was considered too source-breaking↩︎