3
\$\begingroup\$

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?

\$\endgroup\$
1
\$\begingroup\$

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;
    }
}
  1. 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.
\$\endgroup\$
3
  • \$\begingroup\$ It is important to me that the operation be lazy, as I intend to work with some fairly large expressions and would like to take advantage of time-saving functions like Enumerable.Any. \$\endgroup\$ – Rhaokiel Nov 20 '19 at 15:46
  • \$\begingroup\$ I've also tried route #2 and found that there was no easy way to store a simple stack of Expressions along with a detailed value of the process index of that Expression's children which is necessary to continue when moving back up the stack. Because each Expression differs so much, the storage would require the whole set of children to be added to the stack, which requires almost as much storage as your Flatten example. Only to add the characteristic of being lazy. Whereas my Task example uses the storage method already present in each Expression as it enumerates. \$\endgroup\$ – Rhaokiel Nov 20 '19 at 16:11
  • \$\begingroup\$ In that case, put your code in the visitor. \$\endgroup\$ – Adam Brown Nov 21 '19 at 16:39

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.