Implementing Find and Replace for Expression trees

UPDATE September 2011

The code described below in now available from NuGet that installs a single file (ExpressionExtensions.cs) into your project.

In addition to search and replace, methods for expression composition (AND/OR) has been added.

Get it from NuGet…….PM> Install-Package ExpressionTools

The source is now hosted at GitHub

Introduction

Imagine that you are working with an Expression tree and you need to check
if the tree contains a certain Expression with certain property values.

What we basically would need to do is to visit every node(Expression) in the tree
and check to see if the Expression matches what we are searching for.

First of all we need a class that visits all the nodes in the tree.

The ExpressionVisitor

The ExpressionVisitor class that can be found in the MSDN documentation is sort of an
all purpose class when dealing with expressions.

It was originally written by Matt Warren (the father of Linq to Sql) and was first published
as part of his "Building an IQueryable Provider" series

The expression visitor was considered so essential for everyone working with expression trees,
so they finally included the source in the MSDN documentation.
The implementation is really quite simple. It takes care of visiting all the nodes in the tree as well as it takes care of rewriting the tree if we make any changes to it. Great stuff!!

So what we are aiming for here is some kind of method that returns a list of expressions
that matches any predicate defined for the expression (sub)tree.

Implementing Find

Extension methods are a great way of adding new features to existing types and in this case
we want to extend the Expression class.

/// <summary>

/// Returns a list of <typeparamref name="TExpression"/> instances

/// that matches the <paramref name="predicate"/>.

/// </summary>

/// <typeparam name="TExpression">The type of <see cref="Expression"/>

/// to search for.</typeparam>

/// <param name="expression">The <see cref="Expression"/> that represents the sub tree for which to start searching.</param>

/// <param name="predicate">The <see cref="Func{T,TResult}"/> used to filter the result</param>

/// <returns>A list of <see cref="Expression"/> instances that matches the given predicate.</returns>

public static IEnumerable<TExpression> Find<TExpression>(this Expression expression, Func<TExpression, bool> predicate) where TExpression : Expression

{

    var finder = new ExpressionFinder<TExpression>();

    return finder.Find(expression, predicate);

}

Given the following expression:

Expression<Func<string, bool>> expression = (s) => s.Length == 5;

we can do a search in the expression tree by doing:

var binaryExpressions = expression.Find<BinaryExpression>(b => true);

or vi could search for ConstantExpression like this:

var constantExpressions = expression.Find<ConstantExpression>(c => (int)c.Value == 5);

A top-down search is performed in the sub tree that the target expression represents,
meaning that we can choose to investigate the whole tree or just individual sub trees.

The ExpressionFinder

This is a very simple class that override the Visit method that is called for every node in the expression tree.

If the expression being visited matches the given predicate, we add it to the list of matching expressions.

/// <summary>
/// A class used to search for <see cref="Expression"/> instances. 
/// </summary>
/// <typeparam name="TExpression">The type of <see cref="Expression"/> to search for.</typeparam>
public class ExpressionFinder<TExpression> : ExpressionVisitor where TExpression : Expression
{
    private readonly IList<TExpression> _result = new List<TExpression>();
    private Func<TExpression, bool> _predicate;

    /// <summary>
    /// Returns a list of <typeparamref name="TExpression"/> instances that matches the <paramref name="predicate"/>.
    /// </summary>
    /// <param name="expression">The <see cref="Expression"/> that represents the sub tree for which to start searching.</param>
    /// <param name="predicate">The <see cref="Func{T,TResult}"/> used to filter the result</param>
    /// <returns>A list of <see cref="Expression"/> instances that matches the given predicate.</returns>
    public IEnumerable<TExpression> Find(Expression expression, Func<TExpression, bool> predicate)
    {
        _result.Clear();
        _predicate = predicate;
        Visit(expression);
        return _result;
    }

    /// <summary>
    /// Visits each node of the <see cref="Expression"/> tree checks 
    /// if the current expression matches the predicate.         
    /// </summary>
    /// <param name="expression">The <see cref="Expression"/> currently being visited.</param>
    /// <returns><see cref="Expression"/></returns>
    protected override Expression Visit(Expression expression)
    {                        
        if (expression != null && expression is TExpression)
        {
            if (_predicate((TExpression)expression))
                _result.Add((TExpression)expression);
        }

        return base.Visit(expression);
    }
}

Implementing Replace

In addition to being able to search for expressions it would also be nice if we could do a simple replace using a similar syntax.

Another extension method is called for:

/// <summary>

/// Searches for expressions using the given <paramref name="predicate"/> and replaces matches with

/// the result from the <paramref name="replaceWith"/> delegate.

/// </summary>

/// <typeparam name="TExpression">The type of <see cref="Expression"/> to search for.</typeparam>

/// <param name="expression">The <see cref="Expression"/> that represents the sub tree

/// for which to start searching.</param>

/// <param name="predicate">The <see cref="Func{T,TResult}"/> used to filter the result</param>

/// <param name="replaceWith">The <see cref="Func{T,TResult}"/> used to specify the replacement expression.</param>

/// <returns>The modified <see cref="Expression"/> tree.</returns>

public static Expression Replace<TExpression>(this Expression expression, Func<TExpression, bool> predicate, Func<TExpression,Expression> replaceWith) where TExpression : Expression

{

    var replacer = new ExpressionReplacer<TExpression>();

    return replacer.Replace(expression, predicate, replaceWith);

}

We might think of the Replace method as an ad-hoc alternative to modifying the expression tree using an ExpressionVisitor derivate.

Given the following expression:

Expression<Func<string, bool>> expression = (s) => s.Length == 5;

we want to modify the expression tree by replacing the BinaryExpression(Equal) with a BinaryExpression(NotEqual).

This can now be achieved very easily by doing:

var modifiedExpression = expression.Replace<BinaryExpression>(b => true, 
    (b) => Expression.NotEqual(b.Left, b.Right));

The ExpressionReplacer

Similar to the ExpressionFinder it filter outs the target expressions, but instead of returning the matching expressions, it replaces the
matching expression with whatever the Func<TExpression,Expression> delegate returns.

One thing to notice here is that matching expression is passed to the replaceWith delegate.

/// <summary>
/// A class that is capable of doing a find and replace in an <see cref="Expression"/> tree.
/// </summary>
/// <typeparam name="TExpression">The type of <see cref="Expression"/> to find and replace.</typeparam>
public class ExpressionReplacer<TExpression> : ExpressionVisitor where TExpression : Expression
{
    
    private Func<TExpression, Expression> _replaceWith;
    private Func<TExpression, bool> _predicate;
    
    
    /// <summary>
    /// Searches for expressions using the given <paramref name="predicate"/> and 
    /// replaces matches with the result from the <paramref name="replaceWith"/> delegate.          
    /// </summary>
    /// <param name="expression">The <see cref="Expression"/> that 
    /// represents the sub tree for which to start searching.</param>
    /// <param name="predicate">The <see cref="Func{T,TResult}"/> used to filter the result</param>
    /// <param name="replaceWith">The <see cref="Func{T,TResult}"/> 
    /// used to specify the replacement expression.</param>
    /// <returns>The modified <see cref="Expression"/> tree.</returns>
    public Expression Replace(Expression expression, Func<TExpression, bool> predicate, 
        Func<TExpression, Expression> replaceWith)
    {
        _replaceWith = replaceWith;
        _predicate = predicate;
        return Visit(expression);
    }

    /// <summary>
    /// Visits each node of the <see cref="Expression"/> tree checks 
    /// if the current expression matches the predicate. If a match is found 
    /// the expression will be replaced.        
    /// </summary>
    /// <param name="expression">The <see cref="Expression"/> currently being visited.</param>
    /// <returns><see cref="Expression"/></returns>        
    protected override Expression Visit(Expression expression)
    {
        if (expression != null && expression is TExpression)
        {
            if (_predicate((TExpression)expression))
                return _replaceWith((TExpression)expression);
        }
        return base.Visit(expression);
    }
}

Enjoy!!!