I sometimes find myself writing code like:
for x in xs {
for y in ys {
for z in zs {
...
}
}
}
wishing it could be written as:
for (x, y, z) in xs * ys * zs {
...
}
where the result of xs * ys * zs
would be an instance of some kind of sequence type with Element == (XS.Element, YS.Element, ZS.Element)
.
Here's a working implementation of such a sequence type.
struct Product<BaseA, BaseB, Element>: Sequence, IteratorProtocol
where BaseA: Sequence, BaseB: Collection
{
typealias Transform = (BaseA.Element, BaseB.Element) -> Element
let baseA: BaseA
let baseB: BaseB
var iteratorA: BaseA.Iterator
var iteratorB: BaseB.Iterator
var elementA: BaseA.Element?
let transform: Transform
init(_ baseA: BaseA, _ baseB: BaseB, _ transform: @escaping Transform) {
self.baseA = baseA
self.baseB = baseB
self.iteratorA = baseA.makeIterator()
self.iteratorB = baseB.makeIterator()
self.transform = transform
self.elementA = self.iteratorA.next()
}
mutating func next() -> Element? {
guard let elementA = elementA else { return nil }
if let elementB = iteratorB.next() { return transform(elementA, elementB) }
guard let elementA = self.iteratorA.next() else { return nil }
self.elementA = elementA
iteratorB = baseB.makeIterator()
if let elementB = iteratorB.next() { return transform(elementA, elementB) }
return nil
}
}
typealias P2<A: Sequence, B: Collection> = Product<A, B, (A.Element, B.Element)>
typealias P3<A: Sequence, B: Collection, C: Collection> = Product<P2<A, B>, C, (A.Element, B.Element, C.Element)>
typealias P4<A: Sequence, B: Collection, C: Collection, D: Collection> = Product<P3<A, B, C>, D, (A.Element, B.Element, C.Element, D.Element)>
func *<A, B>(lhs: A, rhs: B) -> Product<A, B, (A.Element, B.Element)>
where A: Sequence, B: Collection {
Product(lhs, rhs, { ($0, $1) })
}
func *<A, B, C>(lhs: P2<A, B>, rhs: C) -> P3<A, B, C>
where A: Sequence, B: Sequence, C: Collection {
Product(lhs, rhs, { ($0.0, $0.1, $1) })
}
func *<A, B, C, D>(lhs: P3<A, B, C>, rhs: D) -> P4<A, B, C, D>
where A: Sequence, B: Sequence, C: Sequence, D: Collection {
Product(lhs, rhs, { ($0.0, $0.1, $0.2, $1) })
}
func test() {
let xs = ["a", "b"]
let ys = [10, 20]
let zs = [1.2, 2.3]
let ws = [true, false]
for (x, y, z, w) in xs * ys * zs * ws {
debugPrint(x, y, z, w)
}
}
test()
It supports up to 4 base sequences of different types, but it feels like a nicer implementation should be possible, perhaps supporting any number of base sequences with no or less boiler plate / special casing (this code needs a complicated *
operator overload for every extra number of base sequences we wish to support).
And ideally, it would be as performant as the corresponding nested for loop implementation (I haven't checked the performance of the example program).
Any ideas for improvements?