Recursive flattening of Swift sequences

In Flatten to get all child controls of certain type in a UIView, methods were discussed to
recursively flatten a tree-like structure in Swift, resulting in an array
of all elements.

Motivated by that thread, I have written a flatten function which creates
a sequence instead. This can be an advantage if the resulting elements are to be post-processed (e.g. filtered), because they can be created
"lazily" or "on-demand", instead of creating an array with all
resulting elements first.

func sequentialFlatten
    (seq : S, children : (S.Generator.Element) -> S) -> SequenceOf
{
    return SequenceOf {
        () -> GeneratorOf in

        // Current generator, or `nil` if all sequences are exhausted:
        var gen : S.Generator? = seq.generate()
        // Queue of generators for sequences which still have to be enumerated:
        var queue : [S.Generator] = []

        return GeneratorOf {

            while gen != nil {
                if let next = gen!.next() {
                    queue.append(children(next).generate())
                    return next
                }
                // Current generator is empty, try to deque next one:
                gen = queue.first // `nil` if queue is empty
                if (gen != nil) {
                    queue.removeAtIndex(0)
                }
            }
            return nil
        }
    }
}

The function arguments are an initial sequence, and a closure which
transforms any sequence element into a new sequence (which may be
empty). One can view seq as a set of root nodes in a tree
or forest, and children as the mapping from each node to its children.

(Swift has a related flapMap() function which takes a sequence and a
transformation as argument, but that does not work recursively, and
returns an array and not a sequence.)

Example 1:

```
let someView : UIView = ...
let views = sequentialFlatten([someView], { $0.subviews as! [UIView] })
let labels = lazy(views).filter( { $0 is UILabel } ).ma

This question is more than a year old, so a few things have changed in Swift since. Most notably, SequenceOf and GeneratorOf have been renamed to AnySequence and AnyGenerator respectively.

I'd rename gen to generator so that in the while loop, you can do:

while let gen = generator {
    ...
}

That way, you don't have to force unwrap gen.

You need to declare generator and the elements of queueas AnyGenerators for this to work. That way they are ensured to be reference types and gen and generator refer to the same instance.

// Current generator, or `nil` if all sequences are exhausted:
var generator: AnyGenerator? = AnyGenerator(seq.generate())
// Queue of generators for sequences which still have to be enumerated:
var queue: [AnyGenerator] = []

The four lines where you try to dequeue the first generator could be written in one line if Arrays had a popFirst function. (Which it should but weirdly doesn't.) So we can define that ourselves in an extension. (though this really is a matter of preference)

extension Array {
    mutating func popFirst() -> Generator.Element? {
         if let elem = first {
            removeAtIndex(0)
            return elem
        }
        return nil
    }
}

Now we can write:

generator = queue.popFirst()

This one might be going a bit far because we're extending Array just to get rid of 3 lines, but I really think it improves readability. In the end it's your call.

Put together, that gives us:

func sequentialFlatten(seq : S, children : (S.Generator.Element) -> S) -> AnySequence {
    return AnySequence {
        () -> AnyGenerator in
        
        // Current generator, or `nil` if all sequences are exhausted:
        var generator: AnyGenerator? = AnyGenerator(seq.generate())
        // Queue of generators for sequences which still have to be enumerated:
        var queue: [AnyGenerator] = []
        
        return AnyGenerator {
            
            while let gen = generator {
                if let next = gen.next() {
                    queue.append(AnyGenerator(children(next).generate()))
                    return next
                }
                // Current generator is empty, try to deque next one:
                generator = queue.popFirst() // `nil` if queue is empty
            }
            return nil
        }
    }
}

This version works exactly As your old version, which is breadth-first. This less than ideal, since a breadth first search takes more memory. Additionally, you asked:

• Can sequentialFlatten() be implemented in a way that it calls itself recursively?

Let's see if we can't kill two birds with one stone.
SequentialFlatten(), a recursive approach

Here's my thinking, if we want to go depth first, we must explore the subtree attached to an element before visiting its siblings. We can solve this problem recursively, by maintaining two generators at each level. one for the sequence that was passed in, and another one for the children of the top level element we last generated:

var mainGen = seq.generate()
// Current generator, or `nil` if all sequences are exhausted:
var generator: AnyGenerator?

Initially, generator is nil, so we'll definitely use a guard to unwrap it and return the first element of the main generator if it's nil:

guard let gen = generator else {
    if let elem = mainGen.next() {
        generator = // get the generator for elem's children
        return elem
    }
    return nil
}

This is where recursion comes in:

generator = sequentialFlatten(children(elem), children: children).generate()

It might have been better to define a local function that return a Generator and make that recursive instead of wrapping the Generator up in a Sequence only to then get the Generator again, but I think either way works.

Now all that's left to do is return the right element if generator is not nil and create a new generator when our current one is empty. This can all be solved in one go:

guard let gen = generator, let elem = gen.next() else {
    if let elem = mainGen.next() {
        generator = sequentialFlatten(children(elem), children: children).generate()
            return elem
        }
        return nil
    }
    return elem

We now have the next element of the generator and the generator being nil or empty can be treated the same.

Putting this all together gives us:

``
func sequentialFlatten(seq : S, children : (S.Generator.Element) -> S) -> AnySequence {
return AnySequence {
() -> AnyGenerator in

var mainGen = seq.generate()
// Current generator, or nil` if all sequences are exhausted:
var generator: AnyGenerator?

return AnyGenerator {

guard let gen = generator, let elem = gen.next() else {
if let elem = mainGen.next() {
generator = sequentialFlatten(children(elem), children: children).gen

StackExchange Code Review Q#87004, answer score: 3