How to iterate through a tree structure using a generator?

Question

I'm trying to figure out how to implement a function in a tree node which returns all its descendant leaves (whether direct or indirect). However, I don't want to pass a container in which the leaf nodes will be put recursively (the tree could be huge), instead I'd like to use a generator to iterate through the tree. I've tried a few approaches but none of them worked so far. This one is the closest I've come to a possible solution:

    public interface ITreeNode
    {
        IEnumerable<ITreeNode> EnumerateLeaves();            
    }

    class Leaf : ITreeNode
    {
        public IEnumerable<ITreeNode> EnumerateLeaves()
        {
            throw new NotImplementedException();
        }
    }

    class Branch : ITreeNode
    {
        private List<ITreeNode> m_treeNodes = new List<ITreeNode>();

        public IEnumerable<ITreeNode> EnumerateLeaves()
        {
            foreach( var node in m_treeNodes )
            {
                if( node is Leaf )
                    yield return node;
                else
                    node.EnumerateLeaves();
            }
        }
    }

But this is not working either. What am I doing wrong? Seems like calling .EnumerateLeaves recursively won't work if there's a yield statement in the same function.

Any help would be very much appreciated. Thanks in advance.

Edit: I forgot to mention that a branch can have either leaves or branches as children, hence the recursion.

Answer 1

Here's how you should implement Branch.EnumerateLeaves:

public IEnumerable<ITreeNode> EnumerateLeaves()
{
    foreach( var node in m_treeNodes )
    {
        if( node is Leaf )
            yield return node;
        else
        {
            foreach (ITreeNode childNode in node.EnumerateLeaves())
                yield return childNode;
        }
    }
}

Answer 2

lassevk is right - one potential issue with that method, however, is that recursively calling into enumerables can yield O(n^2) performance. If this is an issue, then you should instead factor the recursion out and use an internal stack.

public IEnumerable<ITreeNode> EnumerateLeaves()
{
    Stack<ITreeNode> workStack = new Stack<ITreeNode>(m_treeNodes);

    while(workStack.Count > 0) {
        var current = workStack.Pop();
        if(current is Leaf)
            yield return current;
        else {
            for(n = 0; n < current.m_treeNodes.Count; n++) {
                workStack.Push(current.m_treeNodes[n]);
            }
        }
    }
}

This should perform the same function, without the recursion.

Edit: Daniel Plaisted posted in a comment about a bigger problem with recursively calling Enumerators, summed up in a post on the MSDN Blogs regarding iterators. Thanks Daniel. =)

Another Edit: Always remember that code simplicity can be more important than performance, especially if you expect others to maintain your code. If you don't expect your tree to grow very large, use the recursion method lassevk outlined in his answer.