Recursive flattening of Swift sequences

Question

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<S : SequenceType>
    (seq : S, children : (S.Generator.Element) -> S) -> SequenceOf<S.Generator.Element>
{
    return SequenceOf {
        () -> GeneratorOf<S.Generator.Element> 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 } ).map( { $0 as! UILabel } )

creates a sequence of all views in the view hierarchy of someView, and then extracts all labels from that sequence. (Note that no intermediate array of all views is created.)

Example 2:

let seq = sequentialFlatten([0], { n in n < 50 ? [ 2*n+1, 2*n+2] : []})
for x in seq { println(x) }

is an unconventional method to print the numbers 0 ... 100.

I wrote this function mainly for my own educational purpose, to get more familiar with Swift sequences and generators. All feedback is welcome, for example:

• Naming (function and variables).
• Swifty-ness: things which should be done different in Swift.
• Implementation: can the same be achieved easier, perhaps using available functions from the Swift standard library?
• Can sequentialFlatten() be implemented in a way that it calls itself recursively?

Answer 1

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<S.Generator.Element>? = AnyGenerator(seq.generate())
// Queue of generators for sequences which still have to be enumerated:
var queue: [AnyGenerator<S.Generator.Element>] = []

---

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<S : SequenceType>(seq : S, children : (S.Generator.Element) -> S) -> AnySequence<S.Generator.Element> {
    return AnySequence {
        () -> AnyGenerator<S.Generator.Element> in
        
        // Current generator, or `nil` if all sequences are exhausted:
        var generator: AnyGenerator<S.Generator.Element>? = AnyGenerator(seq.generate())
        // Queue of generators for sequences which still have to be enumerated:
        var queue: [AnyGenerator<S.Generator.Element>] = []
        
        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<S.Generator.Element>?

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<S : SequenceType>(seq : S, children : (S.Generator.Element) -> S) -> AnySequence<S.Generator.Element> {
    return AnySequence {
        () -> AnyGenerator<S.Generator.Element> in
        
        var mainGen = seq.generate()
        // Current generator, or `nil` if all sequences are exhausted:
        var generator: AnyGenerator<S.Generator.Element>?
        
        return AnyGenerator {
            
            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
        }
    }
}

which is 5 lines of code shorter and requires O(log n) storage instead of O(n).

Answer 2

Here is a Swift 3 version as an extension:

Code was been restructured to make it easier to read.

extension Sequence
{
    public func recursiveFlatten <O: Sequence> (children: @escaping (Self.Iterator.Element) -> O?) -> AnySequence<Self.Iterator.Element>

        where O.Iterator.Element == Self.Iterator.Element

    {
        return AnySequence
        {
            () -> AnyIterator<Self.Iterator.Element> in

            var iterator = self.makeIterator()
            var childIterator: AnyIterator<Self.Iterator.Element>?

            return AnyIterator
            {
                if let childIterator = childIterator, let nextChild = childIterator.next()
                {
                    return nextChild
                }
                else
                {
                    if let next = iterator.next()
                    {
                        childIterator = children(next)?.recursiveFlatten(children: children).makeIterator()
                        return next
                    }
                    else
                    {
                        return nil
                    }
                }
            }
        }
    }
}