Peekable Iterator in Swift
Swift has the ability to create an iterator out of a sequence by calling the makeIterator method provided by the Sequence protocol.
let numbers = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
var iterator = numbers.makeIterator()
// Print all values from 0 to 9
while let value = iterator.next() {
print(value)
}
To my knowledge there is no way of getting a peek of the next element in the sequence without advancing the iterator. But since this was exactly what I needed, I had to come up with my own solution. I started by creating a new peekable iterator struct that would take in an already existing iterator. Below is an example for an Array<Int>.Iterator.
struct PeekableIntegerIterator {
private var iterator: Array<Int>.Iterator
private var nextElement: Int?
init(_ iterator: Array<Int>.Iterator) {
self.iterator = iterator
}
mutating func peek() -> Int? {
if nextElement == nil {
nextElement = iterator.next()
}
return nextElement
}
mutating func next() -> Int? {
guard let result = nextElement else {
return iterator.next()
}
nextElement = nil
return result
}
}
This would be sufficient to use, but at some point I had a function that took an IteratorProtocol as input. The internal iterator I was using (Array<Int>.Iterator in the example above) did implement the IteratorProtocol so my first instinct was to make the property public and pass it to the function. But exposing the internals because I needed to pass them to a function felt wrong. A much better way was to make the peekable iterator conform to IteratorProtocol instead.
struct PeekableIntegerIterator: IteratorProtocol {
typealias Element = Array<Int>.Iterator.Element
private var iterator: Array<Int>.Iterator
private var nextElement: Element?
init(_ iterator: Array<Int>.Iterator) {
self.iterator = iterator
}
mutating func peek() -> Element? {
// ...
}
mutating func next() -> Element? {
// ...
}
}
At some point I implemented a second peekable iterator over a different iterator type and realised that the implementation did not change apart from the element and iterator type. This meant I could have one generic peekable iterator and use that for both types instead.
struct PeekableIterator<Iterator: IteratorProtocol>: IteratorProtocol {
typealias Element = Iterator.Element
private var iterator: Iterator
private var nextElement: Element?
init(_ iterator: Iterator) {
self.iterator = iterator
}
mutating func peek() -> Element? {
// ...
}
mutating func next() -> Element? {
// ...
}
}
As a bonus I also added a PeekableProtocol that contained the peek method definition and made the PeekableIterator conform to it.
// The peekable protocol
protocol PeekableProtocol<Element> {
associatedtype Element
mutating func peek() -> Self.Element?
}
// The new peekable iterator definition
struct PeekableIterator<Iterator: IteratorProtocol>: PeekableProtocol, IteratorProtocol {
// ...
}
To make it easier to create a peekable iterator from an existing iterator I added an extension that would allow me to call peekable() on any iterator and return a new PeekableIterator instance.
extension IteratorProtocol {
func peekable() -> PeekableIterator
However, I removed this extension as soon as I realised that PeekableIterator also implements IteratorProtocol, which meant I could call peekable on a PeekableIterator itself.
let numbers = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
var peekable = numbers.makeIterator().peekable().peekable().peekable()
A less confusing way was to provide an extension over a Sequence instead, since that protocol requires a makeIterator method implementation.
extension Sequence {
func makePeekableIterator() -> PeekableIterator<Self.Iterator> {
return .init(makeIterator())
}
}