Implementing recursive generator for simple tree structure in Swift

Question

I have a simple tree structure in memory based on an XML document and I am trying to write a recursive generator to support SequenceType, but I am stuck on how

Answer 1

I don't know if a generator itself can be recursive. Will M proved me wrong!

Here is a possible implementation for a pre-order traversal, using a stack for the child nodes which still have to be enumerated:

extension XMLNode : SequenceType {
    public func generate() -> AnyGenerator<XMLNode> {
        var stack : [XMLNode] = [self]
        return anyGenerator {
            if let next = stack.first {
                stack.removeAtIndex(0)
                stack.insertContentsOf(next.childNodes, at: 0)
                return next
            }
            return nil
        }
    }
}

For a level-order traversal, replace

stack.insertContentsOf(next.childNodes, at: 0)

by

stack.appendContentsOf(next.childNodes)

Answer 2

While Martin's answer is certainly more concise, it has the downside of making a lot of using a lot of array/insert operations and is not particularly usable in lazy sequence operations. This alternative should work in those environments, I've used something similar for UIView hierarchies.

public typealias Generator = AnyGenerator<XMLNode>

public func generate() -> AnyGenerator<XMLNode> {
    var childGenerator = childNodes.generate()
    var subGenerator : AnyGenerator<XMLNode>?
    var returnedSelf = false

    return anyGenerator {
        if !returnedSelf {
            returnedSelf = true
            return self
        }

        if let subGenerator = subGenerator,
            let next = subGenerator.next() {
                return next
        }

        if let child = childGenerator.next() {
            subGenerator = child.generate()
            return subGenerator!.next()
        }

        return nil
    }
}

Note that this is preorder iteration, you can move the if !returnedSelf block around for post order.

Answer 3

Here is a recursive post-order generator. Can't say I'd recommend actually using it though. @MartinR's answer seems a bit more practical

 public func generate() -> AnyGenerator<XMLNode> {
    var childGenerator:AnyGenerator<XMLNode>?
    var childArrayGenerator:IndexingGenerator<[XMLNode]>? = self.childNodes.generate()
    var returnedSelf = false

    return anyGenerator {
        if let next = childGenerator?.next() {
            return next
        }
        if let child = childArrayGenerator?.next() {
            childGenerator = child.generate()
            return childGenerator?.next()

        } else if !returnedSelf {

            returnedSelf = true
            return self

        } else {

            return nil
        }
    }
}