Take optimizations

Ix.NET/Source/System.Interactive.Async/Take.cs

@@ -4,6 +4,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -16,168 +17,324 @@ namespace System.Linq
         {
             if (source == null)
                 throw new ArgumentNullException(nameof(source));
-            if (count < 0)
-                throw new ArgumentOutOfRangeException(nameof(count));
 
-            return CreateEnumerable(
-                () =>
-                {
-                    var e = source.GetEnumerator();
-                    var n = count;
-
-                    var cts = new CancellationTokenDisposable();
-                    var d = Disposable.Create(cts, e);
-                    var current = default(TSource);
+            if (count <= 0)
+            {
+                return Empty<TSource>();
+            }
 
-                    return CreateEnumerator(
-                        async ct =>
-                        {
-                            if (n == 0)
-                                return false;
-
-                            var result = await e.MoveNext(cts.Token)
-                                                .ConfigureAwait(false);
-
-                            --n;
-                            if (result)
-                                current = e.Current;
-
-                            if (n == 0)
-                                e.Dispose();
-
-                            return result;
-                        },
-                        () => current,
-                        d.Dispose,
-                        e
-                    );
-                });
+            return new TakeAsyncIterator<TSource>(source, count);
         }
 
         public static IAsyncEnumerable<TSource> TakeLast<TSource>(this IAsyncEnumerable<TSource> source, int count)
         {
             if (source == null)
                 throw new ArgumentNullException(nameof(source));
-            if (count < 0)
-                throw new ArgumentOutOfRangeException(nameof(count));
 
-            return CreateEnumerable(
-                () =>
+            if (count <= 0)
+            {
+                return Empty<TSource>();
+            }
+
+            return new TakeLastAsyncIterator<TSource>(source, count);
+        }
+
+        public static IAsyncEnumerable<TSource> TakeWhile<TSource>(this IAsyncEnumerable<TSource> source, Func<TSource, bool> predicate)
+        {
+            if (source == null)
+                throw new ArgumentNullException(nameof(source));
+            if (predicate == null)
+                throw new ArgumentNullException(nameof(predicate));
+
+            return new TakeWhileAsyncIterator<TSource>(source, predicate);
+        }
+
+        public static IAsyncEnumerable<TSource> TakeWhile<TSource>(this IAsyncEnumerable<TSource> source, Func<TSource, int, bool> predicate)
+        {
+            if (source == null)
+                throw new ArgumentNullException(nameof(source));
+            if (predicate == null)
+                throw new ArgumentNullException(nameof(predicate));
+
+            return new TakeWhileWithIndexAsyncIterator<TSource>(source, predicate);
+        }
+
+        private sealed class TakeAsyncIterator<TSource> : AsyncIterator<TSource>
+        {
+            private readonly int count;
+            private readonly IAsyncEnumerable<TSource> source;
+            private int currentCount;
+            private IAsyncEnumerator<TSource> enumerator;
+
+            public TakeAsyncIterator(IAsyncEnumerable<TSource> source, int count)
+            {
+                this.source = source;
+                this.count = count;
+                currentCount = count;
+            }
+
+            public override AsyncIterator<TSource> Clone()
+            {
+                return new TakeAsyncIterator<TSource>(source, count);
+            }
+
+            public override void Dispose()
+            {
+                if (enumerator != null)
+                {
+                    enumerator.Dispose();
+                    enumerator = null;
+                }
+
+                base.Dispose();
+            }
+
+
+            protected override async Task<bool> MoveNextCore(CancellationToken cancellationToken)
+            {
+                switch (state)
+                {
+                    case AsyncIteratorState.Allocated:
+                        enumerator = source.GetEnumerator();
+
+                        state = AsyncIteratorState.Iterating;
+                        goto case AsyncIteratorState.Iterating;
+
+                    case AsyncIteratorState.Iterating:
+                        if (currentCount > 0 && await enumerator.MoveNext(cancellationToken)
+                                                                .ConfigureAwait(false))
+                        {
+                            current = enumerator.Current;
+                            currentCount--;
+                            return true;
+                        }
+
+                        break;
+                }
+
+                Dispose();
+                return false;
+            }
+        }
+
+        private sealed class TakeLastAsyncIterator<TSource> : AsyncIterator<TSource>
+        {
+            private readonly int count;
+            private readonly IAsyncEnumerable<TSource> source;
+            private IAsyncEnumerator<TSource> enumerator;
+            private bool isDone;
+            private Queue<TSource> queue;
+
+            public TakeLastAsyncIterator(IAsyncEnumerable<TSource> source, int count)
+            {
+                this.source = source;
+                this.count = count;
+            }
+
+            public override AsyncIterator<TSource> Clone()
+            {
+                return new TakeLastAsyncIterator<TSource>(source, count);
+            }
+
+            public override void Dispose()
+            {
+                if (enumerator != null)
                 {
-                    var e = source.GetEnumerator();
+                    enumerator.Dispose();
+                    enumerator = null;
+                }
+                queue = null; // release the memory
 
-                    var cts = new CancellationTokenDisposable();
-                    var d = Disposable.Create(cts, e);
+                base.Dispose();
+            }
 
-                    var q = new Queue<TSource>(count);
-                    var done = false;
-                    var current = default(TSource);
 
-                    var f = default(Func<CancellationToken, Task<bool>>);
-                    f = async ct =>
+            protected override async Task<bool> MoveNextCore(CancellationToken cancellationToken)
+            {
+                switch (state)
+                {
+                    case AsyncIteratorState.Allocated:
+                        enumerator = source.GetEnumerator();
+                        queue = new Queue<TSource>();
+                        isDone = false;
+
+                        state = AsyncIteratorState.Iterating;
+                        goto case AsyncIteratorState.Iterating;
+
+
+                    case AsyncIteratorState.Iterating:
+                        if (!isDone)
                         {
-                            if (!done)
+                            if (await enumerator.MoveNext(cancellationToken)
+                                                .ConfigureAwait(false))
                             {
-                                if (await e.MoveNext(ct)
-                                           .ConfigureAwait(false))
+                                if (count > 0)
                                 {
-                                    if (count > 0)
+                                    var item = enumerator.Current;
+                                    if (queue.Count >= count)
                                     {
-                                        var item = e.Current;
-                                        if (q.Count >= count)
-                                            q.Dequeue();
-                                        q.Enqueue(item);
+                                        queue.Dequeue();
                                     }
+                                    queue.Enqueue(item);
                                 }
-                                else
-                                {
-                                    done = true;
-                                    e.Dispose();
-                                }
-
-                                return await f(ct)
-                                           .ConfigureAwait(false);
                             }
-                            if (q.Count > 0)
+                            else
                             {
-                                current = q.Dequeue();
-                                return true;
+                                isDone = true;
+                                // Dispose early here as we can
+                                enumerator.Dispose();
+                                enumerator = null;
                             }
-                            return false;
-                        };
-
-                    return CreateEnumerator(
-                        f,
-                        () => current,
-                        d.Dispose,
-                        e
-                    );
-                });
+
+                            goto case AsyncIteratorState.Iterating; // loop until queue is drained
+                        }
+
+                        if (queue.Count > 0)
+                        {
+                            current = queue.Dequeue();
+                            return true;
+                        }
+
+                        break;
+                }
+
+                Dispose();
+                return false;
+            }
         }
 
-        public static IAsyncEnumerable<TSource> TakeWhile<TSource>(this IAsyncEnumerable<TSource> source, Func<TSource, bool> predicate)
+        private sealed class TakeWhileAsyncIterator<TSource> : AsyncIterator<TSource>
         {
-            if (source == null)
-                throw new ArgumentNullException(nameof(source));
-            if (predicate == null)
-                throw new ArgumentNullException(nameof(predicate));
+            private readonly Func<TSource, bool> predicate;
+            private readonly IAsyncEnumerable<TSource> source;
+
+            private IAsyncEnumerator<TSource> enumerator;
+
+            public TakeWhileAsyncIterator(IAsyncEnumerable<TSource> source, Func<TSource, bool> predicate)
+            {
+                Debug.Assert(predicate != null);
+                Debug.Assert(source != null);
+
+                this.source = source;
+                this.predicate = predicate;
+            }
+
+            public override AsyncIterator<TSource> Clone()
+            {
+                return new TakeWhileAsyncIterator<TSource>(source, predicate);
+            }
+
+            public override void Dispose()
+            {
+                if (enumerator != null)
+                {
+                    enumerator.Dispose();
+                    enumerator = null;
+                }
+
+                base.Dispose();
+            }
 
-            return CreateEnumerable(
-                () =>
+            protected override async Task<bool> MoveNextCore(CancellationToken cancellationToken)
+            {
+                switch (state)
                 {
-                    var e = source.GetEnumerator();
+                    case AsyncIteratorState.Allocated:
+                        enumerator = source.GetEnumerator();
+
+                        state = AsyncIteratorState.Iterating;
+                        goto case AsyncIteratorState.Iterating;
 
-                    var cts = new CancellationTokenDisposable();
-                    var d = Disposable.Create(cts, e);
 
-                    return CreateEnumerator(
-                        async ct =>
+                    case AsyncIteratorState.Iterating:
+                        if (await enumerator.MoveNext(cancellationToken)
+                                            .ConfigureAwait(false))
                         {
-                            if (await e.MoveNext(cts.Token)
-                                       .ConfigureAwait(false))
+                            var item = enumerator.Current;
+                            if (!predicate(item))
                             {
-                                return predicate(e.Current);
+                                break;
                             }
-                            return false;
-                        },
-                        () => e.Current,
-                        d.Dispose,
-                        e
-                    );
-                });
+                            current = item;
+                            return true;
+                        }
+
+                        break;
+                }
+
+                Dispose();
+                return false;
+            }
         }
 
-        public static IAsyncEnumerable<TSource> TakeWhile<TSource>(this IAsyncEnumerable<TSource> source, Func<TSource, int, bool> predicate)
+        private sealed class TakeWhileWithIndexAsyncIterator<TSource> : AsyncIterator<TSource>
         {
-            if (source == null)
-                throw new ArgumentNullException(nameof(source));
-            if (predicate == null)
-                throw new ArgumentNullException(nameof(predicate));
+            private readonly Func<TSource, int, bool> predicate;
+            private readonly IAsyncEnumerable<TSource> source;
+
+            private IAsyncEnumerator<TSource> enumerator;
+            private int index;
 
-            return CreateEnumerable(
-                () =>
+            public TakeWhileWithIndexAsyncIterator(IAsyncEnumerable<TSource> source, Func<TSource, int, bool> predicate)
+            {
+                Debug.Assert(predicate != null);
+                Debug.Assert(source != null);
+
+                this.source = source;
+                this.predicate = predicate;
+            }
+
+            public override AsyncIterator<TSource> Clone()
+            {
+                return new TakeWhileWithIndexAsyncIterator<TSource>(source, predicate);
+            }
+
+            public override void Dispose()
+            {
+                if (enumerator != null)
+                {
+                    enumerator.Dispose();
+                    enumerator = null;
+                }
+
+                base.Dispose();
+            }
+
+            protected override async Task<bool> MoveNextCore(CancellationToken cancellationToken)
+            {
+                switch (state)
                 {
-                    var e = source.GetEnumerator();
-                    var index = 0;
+                    case AsyncIteratorState.Allocated:
+                        enumerator = source.GetEnumerator();
+                        index = -1;
+                        state = AsyncIteratorState.Iterating;
+                        goto case AsyncIteratorState.Iterating;
 
-                    var cts = new CancellationTokenDisposable();
-                    var d = Disposable.Create(cts, e);
 
-                    return CreateEnumerator(
-                        async ct =>
+                    case AsyncIteratorState.Iterating:
+                        if (await enumerator.MoveNext(cancellationToken)
+                                            .ConfigureAwait(false))
                         {
-                            if (await e.MoveNext(cts.Token)
-                                       .ConfigureAwait(false))
+                            var item = enumerator.Current;
+                            checked
                             {
-                                return predicate(e.Current, checked(index++));
+                                index++;
                             }
-                            return false;
-                        },
-                        () => e.Current,
-                        d.Dispose,
-                        e
-                    );
-                });
+
+                            if (!predicate(item, index))
+                            {
+                                break;
+                            }
+                            current = item;
+                            return true;
+                        }
+
+                        break;
+                }
+
+                Dispose();
+                return false;
+            }
         }
     }
 }

Ix.NET/Source/Tests/AsyncTests.Single.cs

@@ -711,7 +711,16 @@ namespace Tests
         public void Take_Null()
         {
             AssertThrows<ArgumentNullException>(() => AsyncEnumerable.Take<int>(null, 5));
-            AssertThrows<ArgumentOutOfRangeException>(() => AsyncEnumerable.Take<int>(AsyncEnumerable.Return(42), -1));
+        }
+
+        [Fact]
+        public void Take0()
+        {
+            var xs = new[] { 1, 2, 3, 4 }.ToAsyncEnumerable();
+            var ys = xs.Take(-2);
+
+            var e = ys.GetEnumerator();
+            NoNext(e);
         }
 
         [Fact]
@@ -882,6 +891,16 @@ namespace Tests
             AssertThrows<ArgumentNullException>(() => AsyncEnumerable.Skip<int>(null, 5));
         }
 
+        //[Fact]
+        //public void Skip0()
+        //{
+        //    var xs = new[] { 1, 2, 3, 4 }.ToAsyncEnumerable();
+        //    var ys = xs.Skip(-2);
+
+        //    var e = ys.GetEnumerator();
+        //    NoNext(e);
+        //}
+
         [Fact]
         public void Skip1()
         {
@@ -2414,7 +2433,15 @@ namespace Tests
         public void TakeLast_Null()
         {
             AssertThrows<ArgumentNullException>(() => AsyncEnumerable.TakeLast(default(IAsyncEnumerable<int>), 5));
-            AssertThrows<ArgumentOutOfRangeException>(() => AsyncEnumerable.TakeLast(AsyncEnumerable.Return(42), -1));
+        }
+
+        [Fact]
+        public void TakeLast0()
+        {
+            var xs = new[] { 1, 2, 3, 4 }.ToAsyncEnumerable().TakeLast(-2);
+
+            var e = xs.GetEnumerator();
+            NoNext(e);
         }
 
         [Fact]