Lazy tree

Some stupid task about Expression trees.

Write down a simplest expression tree visitor (inherited from System.Linq.Expressions.ExpressionVisitor) that will have a static method:

public static IEnumerable<Expression> All(Expression node)

This method should return lazy enumeration of all expression in a an expression tree represented by node. By "lazy" I mean that a tree should not be entirely flattered into a list like this:

class NonLazyVisitor : ExpressionVisitor

{

       readonly List<Expression> _nodes = new List<Expression>();

       public static IEnumerable<Expression> All(Expression node)

       {

              var visitor = new NonLazyVisitor();

              visitor.Visit(node);

              return visitor._nodes;

       }

       public override Expression Visit(Expression node)

       {

              _nodes.Add(node);

              return base.Visit(node);

       }

}


This visitor will traverse entire expression tree - the Visit() method will be called on every node regardless of how much nodes I really want to visit. The task is to minimize the footprint and visit as less nodes as possible.

The lazy visitor should return expressions in the same order as NonLazyVisitor above.

The following solution is not fully lazy - it enumerates all child nodes (but not all descendants) when visits a particular node:

internal sealed class LazyVisitor : ExpressionVisitor

{

       private readonly List<Expression> _children = new List<Expression>();

       public override Expression Visit(Expression node)

       {

              _children.Add(node);

              return node;

       }

       public static IEnumerable<Expression> All(Expression node)

       {

              var visitor = new LazyVisitor();

              return visitor.AllInternal(node);

       }

       internal IEnumerable<Expression> AllInternal(Expression node)

       {

              yield return node;

              base.Visit(node);

              var nestedVisitor = new LazyVisitor();

              foreach (var descendant in from child in _children

                                     from descendants in nestedVisitor.AllInternal(child)

                                                              select descendants)

              {

                     yield return descendant;

              }

       }

}

The main idea is to collect all children to a list and then enumerates all grand-children of children in lazy manner - SelectMany() enumerator will be called only when a child is actually enumerated.