I thought I finally nailed SelectRecursive--my Linq-friendly extension method for navigating homogenous hierarchies. In my infinite cleverness, I figured I could avoid paying the common call stack cost of recursive algorithms by creating my own stack.

The solution I posted Wednesday was an okay start of the philosophy, but I really messed up the implementation.

I thought I got it right yesterday, but after being bitten once, I was cautious--notice I said "Smart(er)".

Well, the implementation yesterday was certainly better, but it was missing a critical component: correct clean-up.

IEnumerator<T> implements IDisposable.

Why? Well, suppose you're enumerating over a database table. When one is done walking through the items, you might want to close your connection.

We rarely see the use of IDisposable in this context, because it's hidden behind the 99.999% use of IEnumerable<T>--foreach.

Don't forget that this:

foreach (object item in sourceEnum)
{
    // do something
}

Looks like this to the runtime:

using (IEnumerator<object> source = sourceEnum.GetEnumerator())
{
    while (source.MoveNext())
    {
        // do something
    }
}

Well, actually it looks like this:

IEnumerator<object> source = sourceEnum.GetEnumerator();
try
{
    while (source.MoveNext())
    {
        // do something
    }
}
finally
{
    source.Dispose();
}

foreach does the work of calling Dispose() on the enumerator--both when things go well and if there is an exception thrown in the while loop.

Since the "simple" version of SelectEnumerator I posted yesterday uses foreach, everything works great. Even if an exception is thrown deep in the recursion, the unrolling of the stack will go through all of the nested finally blocks.

The "smart(er)" version of SelectEnumerator will have to be smart to handle disposing every enumerator it creates--both when the stack is popped normally and in the case of an exception.

public static IEnumerable<TSource> SelectRecursive<TSource>(
    this IEnumerable<TSource> source,
    Func<TSource, IEnumerable<TSource>> recursiveSelector)
{
    Stack<IEnumerator<TSource>> stack = new Stack<IEnumerator<TSource>>();
    stack.Push(source.GetEnumerator());

    try
    {
        while (stack.Count > 0)
        {
            if (stack.Peek().MoveNext())
            {
                TSource current = stack.Peek().Current;

                yield return current;

                stack.Push(recursiveSelector(current).GetEnumerator());
            }
            else
            {
                stack.Pop().Dispose();
            }
        }
    }
    finally
    {
        while (stack.Count > 0)
        {
            stack.Pop().Dispose();
        }
    }
}

Take a look at the code from yesterday to see the diff.

At this point, I'm pretty nervous about saying I've gotten this "right". In all honesty, I did the stack flattening to see if I could do it and do it well.

As I've mentioned before, this pattern is only useful if you plan on having prohibitively deep hierarchies where the cost of the churning the runtime stack or the chance of stack overflow is high.

If neither is the case, keep things simple.

Remember what I said about clever.

 

Happy hacking.