Skip to main content
5 of 5
Commonmark migration

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).

overactor
  • 654
  • 3
  • 17