In-lining InvocationExpressions

This is regards to one of my answers: Linq Expression Calling Combines. The question wasn't very well received, but it proved to be an interesting task that I enjoyed working on.

I ended up with the following:

public static class ExpressionHelpers
{
    public static TExpressionType InlineInvokes(this TExpressionType expression)
        where TExpressionType : Expression
    {
        return (TExpressionType)new InvokeInliner().Inline(expression);
    }

    public static Expression InlineInvokes(this InvocationExpression expression)
    {
        return new InvokeInliner().Inline(expression);
    }

    public class InvokeInliner : ExpressionVisitor
    {
        private Stack> _context = new Stack>();
        public Expression Inline(Expression expression)
        {
            return Visit(expression);
        }

        protected override Expression VisitInvocation(InvocationExpression e)
        {
            var callingLambda = e.Expression as LambdaExpression;
            if (callingLambda == null) //Fix as per comment
                return base.VisitInvocation(e);
            var currentMapping = new Dictionary();
            for (var i = 0; i  0)
            {
                var currentMapping = _context.Peek();
                if (currentMapping.ContainsKey(e))
                    return currentMapping[e];
            }
            return e;
        }
    }
}

At its core, it simply in-lines Invoke() calls in the expression tree, replacing them with the called expression's code verbatim.

Examples:

INPUT:
Invoke(i => (i + 1), 3)

OUTPUT:
(3 + 1)  

INPUT: 
i => Invoke((i, j) => (i * j), Invoke(i => (i + 1), i), Invoke(i => (i + 2), i))

OUTPUT:
i => ((i + 1) * (i + 2))

INPUT:  
b => Invoke((d, e) => (d * e), Invoke(b => (50 + Invoke(z => (25 + Invoke(h => (h * 8), z)), b)), b), Invoke(c => (c + 2), b))  

OUTPUT:   
b => ((50 + (25 + (b * 8))) * (b + 2))

Example usage:

```
Expression> f1 = i => i + 1;
Expression> f2 = i => i + 2;

I ended up using this code a few times, and recently was thinking through a scenario where the above will fail:

var iParam = Expression.Parameter(typeof(int));
var lParam = Expression.Parameter(typeof(int));

var original = Expression.Invoke(
    Expression.Lambda(
        Expression.Invoke(
                Expression.Lambda(
                    Expression.Add(iParam, lParam),
                    iParam),
                Expression.Constant(1)
        ),
        lParam),
    Expression.Constant(2)
);
var transformed = original.InlineInvokes();

The original can be visualized as:

Invoke(Param_0 => Invoke(Param_1 => (Param_1 + Param_0), 1), 2)

Here, the inner invocation is using a parameter from its parent. Since we're using a stack to represent the mappings between parameters, we're not taking into consideration Param_0. Here's the current output:

(1 + Param_0)

Here's the fixed code:

protected override Expression VisitInvocation(InvocationExpression e)
{
    var callingLambda = e.Expression as LambdaExpression;
    if (callingLambda == null)
        return base.VisitInvocation(e);
    var currentMapping = new Dictionary();
    for (var i = 0; i  0)
    {
        var existingContext = _context.Peek();
        foreach (var kvp in existingContext)
        {
            if (!currentMapping.ContainsKey(kvp.Key))
                currentMapping[kvp.Key] = kvp.Value;
        }
    }

    _context.Push(currentMapping);
    var result = Visit(callingLambda.Body);
    _context.Pop();
    return result;
}

Note that we only propagate the parameters if they don't already exist. This is required in situations as follows:

var iParam = Expression.Parameter(typeof(int));
var lParam = Expression.Parameter(typeof(int));

var original = Expression.Invoke(
    Expression.Lambda(
        Expression.Invoke(
                Expression.Lambda(
                    Expression.Add(iParam, lParam),
                    iParam, lParam),
                Expression.Constant(1), Expression.Constant(3)
        ),
        lParam),
    Expression.Constant(2)
);

Invoke(Param_0 => Invoke((Param_1, Param_0) => (Param_1 + Param_0), 1, 3), 2)

Here, the inner invoke is providing a value for Param_0 which takes precedence over what the outer invoke is providing. If we compile that expression (without inlining) - we see the result is 4.

Thus, we must ensure that we only add parameters from our parent context if we haven't already been provided them.

StackExchange Code Review Q#116530, answer score: 2