Safely Zipping Sequences

This week John Sundell tweeted about how he uses zip(sequence1, sequence2) to quickly iterate over two Sequences at once.

Seeing this code in front of me I realized that it’s kind of unsafe to do this unless you’re checking if sequence1.count == sequence2.count. A case where zip(_:_:) comes in handy if you’re assigning text from an array to a Collection of IBOutlets. If the text array is not hardcoded and maybe comes from the internet, the count could easily change without you being notified.

Checking if both sequences have the same count is easy. Writing exceptions if there is a mismatch is doable but may not look very appealing.

This was my angle to write SafeZip2Sequence where the checks are done inside the Sequence so the call site looks much cleaner.

SafeZip2Sequence

First, we have to declare the SafeZip2Sequence which can store two Sequences.

public struct SafeZip2Sequence<Sequence1: Sequence, Sequence2: Sequence> {
	fileprivate let sequence1: Sequence1
	fileprivate let sequence2: Sequence2
}

Then, of course, the new struct has to comply to Sequence. A Sequence has three constraints that have to be fit:

associatedtype Element where Self.Element == Self.Iterator.Element
associatedtype Iterator: IteratorProtocol
func makeIterator() -> Self.Iterator

All other constraints can be inferred.

Conforming to Sequence

So we extend our SafeZip2Sequence. To save namespaces we declare the iterator and element directly inside the extension as types. This saves us writing additional types.

The iterator has to hold the iterators of the base sequences.

extension SafeZip2Sequence: Sequence {
	public struct Iterator {
		fileprivate var iterator1: Sequence1.Iterator
		fileprivate var iterator2: Sequence2.Iterator
	}
}

Then we have to provide a method which creates the iterator:

extension SafeZip2Sequence: Sequence {
	public func makeIterator() -> Iterator {
		return Iterator(
			iterator1: sequence1.makeIterator(),
			iterator2: sequence2.makeIterator()
		)
	}
}

Next, we have to define a value that can store all possible value combination. An enum with associated values seems the right fit for that: It can be a pair or either of the Sequences has a value and the other does not.

extension SafeZip2Sequence: Sequence {
	public enum Element {
		case pair(first: Sequence1.Element, second: Sequence2.Element)
		case first(Sequence1.Element)
		case second(Sequence2.Element)
	}
}

The Iterator

Now we have to make the Iterator conform to the IteratorProtocol. To do this it only needs a typealias for the Element and mutating func next() -> Element?. If you define the next method to return an actual type you can omit the typealias as the compiler infers it.

The method next() has to call next() on both stored iterators and create a SafeZip2Sequence.Element or nil if there is nothing more to iterate over.

extension SafeZip2Sequence.Iterator: IteratorProtocol {
	public mutating func next() -> SafeZip2Sequence.Element? {
		let value1 = iterator1.next()
		let value2 = iterator2.next()
		return Element(first: value1, second: value2)
	}
}

To keep the next() clean we implement the logic of the Elements creation in a failing initializer.

fileprivate extension SafeZip2Sequence.Element {
	init?(first: Sequence1.Element?, second: Sequence2.Element?) {
		switch (first, second) {
		case let (first?, second?):
			self = .pair(first: first, second: second)
		case let (first?, nil):
			self = .first(first)
		case let (nil, second?):
			self = .second(second)
		default:
			return nil
		}
	}
}

Here we use a simple switch statement over both values and match patterns.

Adding sugar

Now we have everything in place to use our new sequence. To create one we declare a free function similar to the Swift Standard Libraries zip(_:_:).

public func safeZip<Sequence1: Sequence, Sequence2: Sequence>(
	_ sequence1: Sequence1,
	_ sequence2: Sequence2
) -> SafeZip2Sequence<Sequence1, Sequence2> {
	return SafeZip2Sequence(sequence1: sequence1, sequence2: sequence2)
}

As most users would just want to iterate over the sequence, we could overload the function to accept a method which will be executed on every Element.

public func safeZip<Sequence1, Sequence2>(
	_ sequence1: Sequence1,
	_ sequence2: Sequence2,
	executing: (SafeZip2Sequence<Sequence1, Sequence2>.Element) -> Void
) {
	for value in safeZip(sequence1, sequence2) {
		executing(value)
	}
}

So you can do both:

for element in safeZip(labels, texts) {
	print(element)
}

safeZip(labels, texts) { print($0) }

safeZip(labels, texts, executing: doSomething)

To match the whole implementation of zip(_:_:) we have to add an extension to Sequence as well:

public extension Sequence {
	func mySafeZip<S: Sequence>(_ sequence: S) -> SafeZip2Sequence<Self, S> {
		return safeZip(self, sequence)
	}

	func mySafeZip<S: Sequence>(
		_ sequence: S,
		executing: (SafeZip2Sequence<Self, S>.Element) -> Void
	) {
		return safeZip(self, sequence, executing: executing)
	}
}

This wraps it up. You can find the whole code on my gist on GitHub.