A Sequence type for replacing N nested for loops

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?

2 Likes

Swift Algorithms have something like that already ... called exactly the same Product swift-algorithms/Product.md at main · apple/swift-algorithms · GitHub. Right now it only support product of 2, but you can compose in this case for higher arities like 3 in your example =]

3 Likes

Thanks, wasn't aware of that! But composing them would result in a nested tuple element:

    let xs = ["a", "b"]
    let ys = [10, 20]
    let zs = [1.2, 2.3]
    let ws = [true, false]
    for (((x, y), z), w) in product(product(product(xs, ys), zs), ws) {
        debugPrint(x, y, z, w)
    }

My implementation solves this by using the (escaping) transform closure, at the cost of some performance loss I guess.

1 Like

I think if you’re nesting more than 4 for loops, you may want to refactor your code base.

The boilerplate is unfortunate, but hard to avoid unless you want to make this a language feature. Regarding the operator, you can instead use named functions:

enum Products {
static func cartesian<S1: Sequence, S2: Sequence>
   (_ s1: S1, _ s2: S2) -> Cartesian2<S1, S2> {…}
// other overloads for Cartesian3, 4, 5,…
}

I remember that I essentially made those cartesians just type aliases of some crazy lazy sequence types that you get when implementing the above like this:

s1.lazy.flatMap {e1 in
   s2.lazy.map {(e1, $0)} // and so on
}

Technically speaking, s2 needs to conform to collection for the above code to be valid.

1 Like

I think it would be a good extension of the language if you could even write for x in xs, y in ys, z in zs — or for paragragh in text.paragraphs, sentence in paragraph.sentences, word in sentence.words.
This (maybe with some additional let) would feel quite natural, as if and while already allow this.
It could also solve a well known annoyance by enabling us to write for let sequence = optionalSequence, element in sequence {...

6 Likes

Good idea, that would be a natural and useful extension of the language. Will you propose it?

One issue might be that the nestedness might not be obvious, ie:

let xs = [1, 2]
let ys = ["a", "b"]
for x in xs, y in ys {
  print(x, y)
}

could perhaps be easily (mis)read as if it would print:

1 a
2 b

rather than

1 a
1 b
2 a
2 b

Don't know if this makes it more clear (essentially allowing the brackets to be skipped):

for x in xs, for y in ys {
  print(x, y)
}

But that's almost as verbose as the currently working:

for x in xs { for y in ys {
  print(x, y)
} }
1 Like

That's definitely too big for me :-(however, I would be happy if anyone would "steal" the idea).
Your Swift-only solution has a big advantage (it's more than an idea ;-) and reduces nesting as well, so if you can get it into the stdlib (I really do not understand how that is decided :man_shrugging:), that would be great.

Is the rule now you have to have an implementation to pitch something at all or just to get an official SE-####?

Implementations are required only for proposals.

So, @Tino, you could pitch this without having to write an implementation and we as a community can brainstorm the answers to the questions like, “What is the return type of the iterator for this combo loop?” And if we can figure it out, maybe we can get to the point of Proposal and find help to make an implementation?

Sorry, too little time — and actually too disillusioned about the whole process.
But I don't claim any authorship (wouldn't work anyways :smiley:), so feel free to try your luck.

(and sorry for sidetracking; I think #swift-users should actually protect authors from #evolution :joy:)