How to create a variadic ZipSequence type with Swift's parameter packs?

Question

According to SE-0398, we can now create types with a variable number of generic type parameters. The document uses ZipSequence as an example:

In the generic parameter list of a generic type, the each keyword declares a generic parameter pack, just like it does in the generic parameter list of a generic function. The types of stored properties can contain pack expansion types, as in let seq and var iter below.

This lets us define the return type of the variadic zip function as follows:

struct ZipSequence<each S: Sequence>: Sequence {
  typealias Element = (repeat (each S).Element)

  let seq: (repeat each S)

  func makeIterator() -> Iterator {
    return Iterator(iter: (repeat (each seq).makeIterator()))
  }

  struct Iterator: IteratorProtocol {
    typealias Element = (repeat (each S).Element)

    var iter: (repeat (each S).Iterator)

    mutating func next() -> Element? {
      return ...
    }
  }
}

It unfortunately does not show how to implement the next method.

I tried implementing it as:

(repeat (each iter).next()!)

This produces the error:

Cannot use mutating member on immutable value of type 'τ_1_0.Iterator'

I know that I did not handle the case of next returning nil. Let's just assume that the sequences are all infinite - I just want to get to something that at least compiles first. Handling next returning nil is the least of my concerns, if I can't even write something that compiles.

How can I implement this next method?

Answer 1

Apparently each iter is an immutable value, so you can only pass it to functions expecting a non-inout IteratorProtocol parameter. That function can do var copy = iter to make it mutable, and return the modified copy, together with the next element:

// inside "next"
func nextHelper<I: IteratorProtocol>( _ iter: I) -> (I, I.Element) {
    var copy = iter
    let next = copy.next()!
    return (copy, next)
}

Then, if we do (repeat nextHelper(each iter)), we get a nested tuple structured like this:

((iter1, element1), (iter2, element2), ... (iterN, elementN))

All we need to do now is to create a tuple of all the .0 of each nested tuple, and assign it to the iter property to update it. Then create a tuple of all the .1 of each nested tuple, and return that as the next element.

// inside "next"
let nestedTuple = (repeat nextHelper(each iter))
iter = (repeat (each nestedTuple).0)
return (repeat (each nestedTuple).1)

To handle the case of next returning nil, I thought of a trick - to give nextHelper some side effect:

// inside "next"
var anyReachedEnd = false
func nextHelper<I: IteratorProtocol>( _ iter: I) -> (I, I.Element?) {
    var copy = iter
    if let next = copy.next() {
        return (copy, next)
    } else {
        anyReachedEnd = true
        return (copy, nil)
    }
}

When we do (repeat nextHelper(each iter)), nextHelper runs for every iterator. If any of the iterators returned nil, anyReachedEnd would be set to true.

// inside "next"
let nestedTuple = (repeat nextHelper(each iter))
if anyReachedEnd { return nil }
iter = (repeat (each nestedTuple).0)
return (repeat (each nestedTuple).1!)

I do think this is a bit "nasty". Perhaps there will be better ways to do this in future versions of Swift.

Answer 2

The fact that each iter is immutable, regardless of its storage seems to be a compiler limitation, which hopefully will be addressed soon.

The immutability can be circumvented, however, as you already did, just wanted to provide another solution that avoids the extra Bool property, by using exceptions:

func nextFor<I: IteratorProtocol>(for iterator: I) throws -> (I, I.Element) {
    var iterator = iterator
    guard let value = iterator.next() else { throw IteratorError.noMoreElements }
    return (iterator, value)
}

With the above in place, the next method can look like this:

mutating func next() -> Element? {
    do {
        let iterVal = try (repeat nextFor(each iter))
        iter = (repeat (each iterVal).0)
        return (repeat (each iterVal).1)
    } catch {
        return nil
    }
}

Combining everything together:

struct ZipSequence<each S: Sequence>: Sequence {
    typealias Element = (repeat (each S).Element)
    
    let seq: (repeat each S)
    
    func makeIterator() -> Iterator {
        return Iterator(iter: (repeat (each seq).makeIterator()))
    }
    
    struct Iterator: IteratorProtocol {
        typealias Element = (repeat (each S).Element)
        
        var iter: (repeat (each S).Iterator)
        
        mutating func next() -> Element? {
            do {
                let iterVal = try (repeat nextFor(each iter))
                iter = (repeat (each iterVal).0)
                return (repeat (each iterVal).1)
            } catch {
                return nil
            }
        }
    }
}

private func nextFor<I: IteratorProtocol>(_ iterator: I) throws -> (I, I.Element) {
    var iterator = iterator
    guard let value = iterator.next() else { throw IteratorError.noMoreElements }
    return (iterator, value)
}

private enum IteratorError: Error {
    case noMoreElements
}

func zip<each S: Sequence>(_ seq: repeat each S) -> ZipSequence<repeat each S> {
    ZipSequence(seq: (repeat each seq))
}

Usage example:

let zipped = Array(zip([1, 2, 3], ["a", "b", "c", "d"], [true, false]))
print(zipped)

Output:

[(1, "a", true), (2, "b", false)]