Generate

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

@@ -5,6 +5,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace System.Linq
@@ -20,44 +21,69 @@ namespace System.Linq
             if (resultSelector == null)
                 throw new ArgumentNullException(nameof(resultSelector));
 
-            return CreateEnumerable(
-                () =>
-                {
-                    var i = initialState;
-                    var started = false;
-                    var current = default(TResult);
+            return new GenerateAsyncIterator<TState, TResult>(initialState, condition, iterate, resultSelector);
+        }
 
-                    return CreateEnumerator(
-                        ct =>
-                        {
-                            var b = false;
-                            try
-                            {
-                                if (started)
-                                    i = iterate(i);
+        private sealed class GenerateAsyncIterator<TState, TResult> : AsyncIterator<TResult>
+        {
+            private readonly Func<TState, bool> condition;
+            private readonly TState initialState;
+            private readonly Func<TState, TState> iterate;
+            private readonly Func<TState, TResult> resultSelector;
+            private TState currentState;
 
-                                b = condition(i);
+            private bool started;
 
-                                if (b)
-                                    current = resultSelector(i);
-                            }
-                            catch (Exception ex)
-                            {
-                                return TaskExt.Throw<bool>(ex);
-                            }
+            public GenerateAsyncIterator(TState initialState, Func<TState, bool> condition, Func<TState, TState> iterate, Func<TState, TResult> resultSelector)
+            {
+                this.initialState = initialState;
+                this.condition = condition;
+                this.iterate = iterate;
+                this.resultSelector = resultSelector;
+            }
 
-                            if (!b)
-                                return TaskExt.False;
+            public override AsyncIterator<TResult> Clone()
+            {
+                return new GenerateAsyncIterator<TState, TResult>(initialState, condition, iterate, resultSelector);
+            }
 
-                            if (!started)
-                                started = true;
+            public override void Dispose()
+            {
+                currentState = default(TState);
+
+                base.Dispose();
+            }
+
+            protected override Task<bool> MoveNextCore(CancellationToken cancellationToken)
+            {
+                switch (state)
+                {
+                    case AsyncIteratorState.Allocated:
+                        started = false;
+                        currentState = initialState;
 
+                        state = AsyncIteratorState.Iterating;
+                        goto case AsyncIteratorState.Iterating;
+
+                    case AsyncIteratorState.Iterating:
+                        if (started)
+                        {
+                            currentState = iterate(currentState);
+                        }
+
+                        started = true;
+
+                        if (condition(currentState))
+                        {
+                            current = resultSelector(currentState);
                             return TaskExt.True;
-                        },
-                        () => current,
-                        () => { }
-                    );
-                });
+                        }
+                        break;
+                }
+
+                Dispose();
+                return TaskExt.False;
+            }
         }
     }
 }