LINQ Expression Visitor

For a recent LINQ-based project I needed to merge branches of multiple expression trees into one.  Normally this doesn’t cause any issues, but I wanted to dynamically merge multiple ‘Where’ predicates together with arbitrary logical operators (hope that’s the right term; it sounds about right).  I found a great base class called ExpressionVisitor on Matt Warren’s blog.  Always keen to reduce the amount of typing (and thinking) I need to do, I couldn’t resist making a couple of helpers.  The first allows you to use syntax like this (don’t worry, it gets better…):

    MethodCallExpression call = ...;

    call = ExpressionVisitor<ParameterExpression>.Visit(
        call,
        p => p.Type == typeof( T ) ? elementParam : p ) as MethodCallExpression;

Basically it lets you specify a delegate to visit (and change) a single expression type in the tree.  Here’s the code (you’ll also need Matt’s ExpressionVisitor source):

    public class ExpressionVisitor<T> : ExpressionVisitor where T : Expression
    {
        Func<T, Expression> visitor;

        public ExpressionVisitor( Func<T, Expression> visitor )
        {
            this.visitor = visitor;
        }

        public static Expression Visit(
            Expression exp,
            Func<T, Expression> visitor )
        {
            return new ExpressionVisitor<T>( visitor ).Visit( exp );
        }

        public static Expression<TDelegate> Visit<TDelegate>(
            Expression<TDelegate> exp,
            Func<T, Expression> visitor )
        {
            return (Expression<TDelegate>)new ExpressionVisitor<T>( visitor ).Visit( exp );
        }

        protected override Expression Visit( Expression exp )
        {
            if ( exp is T && visitor != null ) exp = visitor( (T)exp );

            return base.Visit( exp );
        }
    }

As a refinement I added a couple of extension methods for more compact syntax.  For non-generic types you can do this:

    MethodCallExpression call = ...;

    call = call.Visit<ParameterExpression, MethodCallExpression>(
        p => p.Type == typeof( T ) ? elementParam : p );

For generic types you can do this:

    Expression<Func<Profile, bool>> call = ...;

    call = call.Visit<ParameterExpression, Func<Profile, bool>>(
        p => p.Type == typeof( T ) ? elementParam : p );

(However I’d rather the Func<Profile, bool> wasn’t needed in the call to Visit<>() here…  I’d hoped it could infer the TDelegate type and let me just specify the ‘T’ type, but it seems it can’t do both together .  Maybe I’m missing something obvious, or perhaps it’s a beta 2 bug?  Feedback welcome!).

Here’s the code for the extension methods:

    public static class ExpressionExtensions
    {
        public static Expression Visit<T>(
            this Expression exp,
            Func<T, Expression> visitor ) where T : Expression
        {
            return ExpressionVisitor<T>.Visit( exp, visitor );
        }

        public static TExp Visit<T, TExp>(
            this TExp exp,
            Func<T, Expression> visitor )
            where T : Expression
            where TExp : Expression
        {
            return (TExp)ExpressionVisitor<T>.Visit( exp, visitor );
        }

        public static Expression<TDelegate> Visit<T, TDelegate>(
            this Expression<TDelegate> exp,
            Func<T, Expression> visitor ) where T : Expression
        {
            return ExpressionVisitor<T>.Visit<TDelegate>( exp, visitor );
        }
    }

I’m sure there’s scope for other enhancements too…