IEnumerable<T>.Tee implementation

Question

I am trying to write A C# version of a method in the Java Stream API called Collectors.teeing() (It's like Aggregate but with multiple accumulator functions), just for educational purposes.

I want to make sure that I:

1. don't enumerate the source more than once (because that is trivial)
2. don't store the elements from the source (that's also trivial) meaning each element has to be processed before moving on to the next.
3. don't have to reimplement anything. In other words, if for example I want to find the min and max values, I should just be able to call Min() and Max() and not have to write a loop.

This is what I came up with:

public static class Extensions
{
    public static (T1, T2) Tee<TSource, T1, T2>(
        this IEnumerable<TSource> source,
        Func<IEnumerable<TSource>, T1> accumulator1,
        Func<IEnumerable<TSource>, T2> accumulator2)
    {
        var tee = new Teeing<TSource>(source, 2);
        Task<T1> t1 = Task.Run(() => accumulator1(tee));
        Task<T2> t2 = Task.Run(() => accumulator2(tee));
        Task.WaitAll(t1, t2);
        return (t1.Result, t2.Result);
    }

    //overloads with more parameters. T3, T4, etc

    private class Teeing<T> : IEnumerable<T>, IEnumerator<T>
    {
        private IEnumerator<T> sourceEnumerator;
        private Barrier barrier;
        private bool hasNext;

        public Teeing(IEnumerable<T> source, int count)
        {
            sourceEnumerator = source.GetEnumerator();
            barrier = new Barrier(count, b => {
                hasNext = sourceEnumerator.MoveNext();
            });
        }

        public T Current => sourceEnumerator.Current;
        
        public bool MoveNext()
        {
            barrier.SignalAndWait();
            return hasNext;
        }
        
        public void Dispose()
        {
            barrier.RemoveParticipant();
            if (barrier.ParticipantCount == 0) {
                barrier.Dispose();
                sourceEnumerator.Dispose();
            }
        }
        
        public IEnumerator<T> GetEnumerator() => this;
        
        IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
        object IEnumerator.Current => Current!;
        void IEnumerator.Reset() => throw new NotImplementedException();
    }

}

then call:

var tuple = sequence.Tee(
    s => s.Min(),
    s => s.Max()
);

Any exceptions thrown by one of the accumulator functions will be collected in an AggregateException. I guess that works. Or is there a better way?

Answer 1

I don't know much about Java but it seems that stream operations are sequential unless explicitly parallelized. Your code does the parallel bit already which might be slower for small collections and unexpected. It also blocks on the Tasks Result which is not ideal. (Operations should be "async all the way up")(see comment)

With regards to your use case: Passing a function that needs an IEnumerable makes the multi-threading necessary. Because passing an IEnumerable to the function, one has no control any more how fast it is enumerated without multi-threading.

A simpler alternative would be to provide seeds and change the aggregators so they only apply to one element at a time. MoreLinq has this already. For small sets, this works in 1/10 the time as the multi-thread solution.

So your code works, the signature makes this a bit inefficient though. You could consider an implementation for IAsyncEnumerable as you have more possibilities to control the enumeration.

Just to confirm, I ran a small test with 10 million ints on my computer. Tee took about 2 seconds Min() + Max() took 22 milliseconds