Iterate through Expression Tree

Question

I'm trying to build an iterator (enumerator) that can select specific elements in an Expression tree by traversing the tree and deferring further iteration until needed.

My following implementation intends to keep it simple by leveraging the ExpressionVisitor. But since a visitor-pattern uses recursive stack calls to keep its place, I am using a Task thread in a synchronous manner to store the position of the IEnumerator between calls to MoveNext.

/// <summary>
/// Iterates through an expression tree collecting only expression elements that match the predicate.
/// </summary>
internal sealed class ExpressionCollector : IEnumerable<Expression>
{
    private ExpressionCollector(Expression root, Func<Expression, bool> predicate)
    {
        this.Root = root;
        this.Predicate = predicate;
    }

    private Expression Root { get; }
    private Func<Expression, bool> Predicate { get; }

    public static IEnumerable<Expression> Collect(Expression root, Func<Expression, bool> predicate)
        => new ExpressionCollector(root, predicate);

    public IEnumerator<Expression> GetEnumerator() => new Iterator(this);
    IEnumerator IEnumerable.GetEnumerator() => this.GetEnumerator();

    private class Iterator : ExpressionVisitor, IEnumerator<Expression>
    {
        private Task itr;
        private CancellationTokenSource cancel;
        private TaskCompletionSource<Expression> found;
        private TaskCompletionSource<bool> isReady;

        public Iterator(ExpressionCollector owner) { this.Owner = owner; }

        public ExpressionCollector Owner { get; }

        public Expression Current { get; private set; }
        object IEnumerator.Current => this.Current;

        public void Dispose() => this.Reset();
        public bool MoveNext()
        {
            if (this.found is null || this.found.Task.IsCompleted)
                this.found = new TaskCompletionSource<Expression>();

            if (this.itr is null)
                this.itr = Task.Run(Iterate, (cancel = new CancellationTokenSource()).Token);

            this.isReady?.SetResult(true);
            switch (Task.WaitAny(itr, found.Task))
            {
                case 0:
                    this.Current = null;
                    this.found = null;
                    return false;
                case 1:
                    this.Current = found.Task.Result;
                    this.found = null;
                    return true;
                default: throw new NotSupportedException();
            }
        }
        public void Reset()
        {
            if (this.itr != null)
            {
                this.cancel.Cancel();
                this.itr.Wait();

                this.itr.Dispose();
                this.itr = null;

                this.cancel.Dispose();
                this.cancel = null;

                this.Current = null;
            }
        }

        private void Iterate() => this.Visit(this.Owner.Root);
        public override Expression Visit(Expression node)
        {
            if (cancel.IsCancellationRequested)
                cancel.Token.ThrowIfCancellationRequested();

            if (this.Owner.Predicate(node))
            {
                this.isReady = new TaskCompletionSource<bool>();
                this.found.SetResult(node);
                this.isReady.Task.Wait();
            }

            return base.Visit(node);
        }
    }
}

My concern stems from whether or not I am using the Task properly. Any thoughts?

Answer 1

Just some thoughts:

1. Not sure why you don't just flatten the expression tree with a visitor and then return the collection of expressions. This leads to simpler code:

class Program
{
    static void Main(string[] args)
    {
        Expression<Func<int, int, int>> expr = (a, b) => ((a + b) * (a - b));

        //Flatten the entire tree.
        foreach (var elem in expr.Flatten())
        {
            Console.WriteLine(elem.NodeType);
        }

        //Use Linq to select what you want.
        foreach (var elem in expr.Flatten().Where(t => t.NodeType == ExpressionType.Parameter))
        {
            Console.WriteLine(elem.NodeType);
        }
        Console.ReadKey();
    }
}

public static class ExpressionExtensions
{
    public static IEnumerable<Expression> Flatten(this Expression expr)
    {
        return Visitor.Flatten(expr);
    }
}

public sealed class Visitor : ExpressionVisitor
{
    private readonly Action<Expression> nodeAction;

    private Visitor(Action<Expression> nodeAction)
    {
        this.nodeAction = nodeAction;
    }

    public override Expression Visit(Expression node)
    {
        nodeAction(node);
        return base.Visit(node);
    }

    public static IEnumerable<Expression> Flatten(Expression expr)
    {
        var ret = new List<Expression>();
        var visitor = new Visitor(t => ret.Add(t));
        visitor.Visit(expr);
        return ret;
    }
}

2. Even if, for some reason, you absolutely need to have lazy execution, you're better off making your own expression iterator, rather than trying to pause the execution of a class meant to eagerly visit everything.