Zip 3

In the Swift Standard Library, you’ll find a convenient free function called zip which will interlace any two sequences you give it.

let seq1 = [1, 2, 3, 4, 5]
let seq2 = ["a", "b", "c", "d", "e"]

let result = zip(seq1, seq2)
result.forEach { print($0) }

This print out:

(1, "a")
(2, "b")
(3, "c")
(4, "d")
(5, "e")

Sometimes though, you need to interlace three sequences. Behold! Zip3, an implementation closely patterned on the Standard Library’s implementation of zip: formatting, naming conventions, etc.

The function itself is extremely short. It just takes three sequences of possibly three separate types and returns some type that is generic over those same three sequence types. We will overload the Standard Library’s zip function and the third parameter will differentiate this implementation so the compiler will know which to use.

func zip<Sequence1, Sequence2, Sequence3>(
    _ sequence1: Sequence1, _ sequence2: Sequence2, _ sequence3: Sequence3
    ) -> Zip3Sequence<Sequence1, Sequence2, Sequence3> {
    return Zip3Sequence(sequence1, sequence2, sequence3)
}

One admirable trait of the Standard Library’s implementation is the way it breaks its task into most dirt simple parts. The return type of the zip function is a Zip3Sequence struct. It is initialized with the three sequences which are stored as properties. Consequently the types of these three sequences are associated with the aliases Sequence1, Sequence2, and Sequence3.

struct Zip3Sequence<Sequence1: Sequence, Sequence2: Sequence, Sequence3: Sequence> {
    let _sequence1: Sequence1
    let _sequence2: Sequence2
    let _sequence3: Sequence3
    init(_ sequence1: Sequence1, _ sequence2: Sequence2, _ sequence3: Sequence3) {
        (_sequence1, _sequence2, _sequence3) = (sequence1, sequence2, sequence3)
    }
}

Eventually, we want Zip3Sequence to conform to Sequence. As a prerequisite, it needs to define a type conforming to IteratorProtocol that it will use as the return type of the makeIterator function. We’ll nest this definition within the Zip3Sequence itself since it is specific to that type.

extension Zip3Sequence {
    struct Iterator {
        var _baseStream1: Sequence1.Iterator
        var _baseStream2: Sequence2.Iterator
        var _baseStream3: Sequence3.Iterator
        var _reachedEnd: Bool = false
        init(
            _ iterator1: Sequence1.Iterator,
            _ iterator2: Sequence2.Iterator,
            _ iterator3: Sequence3.Iterator
        ) {
            (_baseStream1, _baseStream2, _baseStream3) = (iterator1, iterator2, iterator3)
        }
    }
}

This Iterator struct holds the three iterators created by the three sequences of the Zip3Sequence respectively and a flag indicating whether the iterator should at last return nil.

Of course Zip3Sequence.Iterator does not yet conform to IteratorProtocol. It must implement the next() -> Element? function. This can be done within an extension on Zip3Sequence.Iterator.

extension Zip3Sequence.Iterator: IteratorProtocol {
    public typealias Element = (Sequence1.Element, Sequence2.Element, Sequence3.Element)
    public mutating func next() -> Element? {
        if _reachedEnd {
            return nil
        }
        guard let element1 = _baseStream1.next(),
              let element2 = _baseStream2.next(),
              let element3 = _baseStream3.next() else {
            _reachedEnd = true
            return nil
        }
        return (element1, element2, element3)
    }
}

The return type of next() is an optional Element which is a tuple of the Element types of the three underlying sequence iterators. This the type that each successive iteration of the sequence will yield, a three-part tuple of the three corresponding elements from the three underlying sequences. As long as all three return elements, the sequence continues. Once any one of them returns nil, the whole sequence ends. That’s all that is needed for Zip3Sequence.Iterator to conform to IteratorProtocol. Now that it does, Zip3Sequence can finally conform to Sequence.

extension Zip3Sequence: Sequence {
    typealias Element = (Sequence1.Element, Sequence2.Element, Sequence3.Element)
    func makeIterator() -> Iterator {
        return Iterator(
            _sequence1.makeIterator(),
            _sequence2.makeIterator(),
            _sequence3.makeIterator()
        )
    }
}

All Zip3Sequence needs to do is implement makeIterator() which will create an iterator of the type it just defined and return it. The iterator is initialized with freshly created iterators from the three underlying sequences.

let seq1 = [1, 2, 3, 4, 5]
let seq2 = [2, 4, 6, 8, 10]
let seq3 = [10, 20, 30, 40, 50]

let result = zip(seq1, seq2, seq3)
result.forEach { print($0) }

/*
 (1, 2, 10)
 (2, 4, 20)
 (3, 6, 30)
 (4, 8, 40)
 (5, 10, 50)
 */

The code for zip3 is available a gist. Enjoy!