From 856a049dd084cddc3a326750db3dc466ed3bbe9c Mon Sep 17 00:00:00 2001
From: neuecc <ils@neue.cc>
Date: Fri, 8 May 2020 03:48:46 +0900
Subject: [PATCH] WithCancellation

---
 src/UniTask.NetCore/IAsyncEnumerable.cs | 54 ++++++++++++++++++++++++-
 src/UniTask.NetCore/Linq/Range.cs       |  2 +-
 src/UniTask.NetCore/Linq/Repeat.cs      |  2 +-
 src/UniTask.NetCoreSandbox/Program.cs   | 11 ++++-
 4 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/src/UniTask.NetCore/IAsyncEnumerable.cs b/src/UniTask.NetCore/IAsyncEnumerable.cs
index 9c67692..dfc88e1 100644
--- a/src/UniTask.NetCore/IAsyncEnumerable.cs
+++ b/src/UniTask.NetCore/IAsyncEnumerable.cs
@@ -1,4 +1,5 @@
-using System.Threading;
+using System.Runtime.InteropServices;
+using System.Threading;
 
 namespace Cysharp.Threading.Tasks
 {
@@ -17,4 +18,55 @@ namespace Cysharp.Threading.Tasks
     {
         UniTask DisposeAsync();
     }
+
+    public static class UniTaskAsyncEnumerableExtensions
+    {
+        public static UniTaskCancelableAsyncEnumerable<T> WithCancellation<T>(this IUniTaskAsyncEnumerable<T> source, CancellationToken cancellationToken)
+        {
+            return new UniTaskCancelableAsyncEnumerable<T>(source, cancellationToken);
+        }
+    }
+
+    [StructLayout(LayoutKind.Auto)]
+    public readonly struct UniTaskCancelableAsyncEnumerable<T>
+    {
+        private readonly IUniTaskAsyncEnumerable<T> enumerable;
+        private readonly CancellationToken cancellationToken;
+
+        internal UniTaskCancelableAsyncEnumerable(IUniTaskAsyncEnumerable<T> enumerable, CancellationToken cancellationToken)
+        {
+            this.enumerable = enumerable;
+            this.cancellationToken = cancellationToken;
+        }
+
+        public Enumerator GetAsyncEnumerator()
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+            return new Enumerator(enumerable.GetAsyncEnumerator(cancellationToken));
+        }
+
+        [StructLayout(LayoutKind.Auto)]
+        public readonly struct Enumerator
+        {
+            private readonly IUniTaskAsyncEnumerator<T> enumerator;
+
+            internal Enumerator(IUniTaskAsyncEnumerator<T> enumerator)
+            {
+                this.enumerator = enumerator;
+            }
+
+            public T Current => enumerator.Current;
+
+            public UniTask<bool> MoveNextAsync()
+            {
+                return enumerator.MoveNextAsync();
+            }
+
+
+            public UniTask DisposeAsync()
+            {
+                return enumerator.DisposeAsync();
+            }
+        }
+    }
 }
\ No newline at end of file
diff --git a/src/UniTask.NetCore/Linq/Range.cs b/src/UniTask.NetCore/Linq/Range.cs
index c48e2da..d07e7d6 100644
--- a/src/UniTask.NetCore/Linq/Range.cs
+++ b/src/UniTask.NetCore/Linq/Range.cs
@@ -55,7 +55,7 @@ namespace Cysharp.Threading.Tasks.Linq
 
             public UniTask<bool> MoveNextAsync()
             {
-                if (cancellationToken.IsCancellationRequested) return CompletedTasks.False;
+                cancellationToken.ThrowIfCancellationRequested();
 
                 current++;
 
diff --git a/src/UniTask.NetCore/Linq/Repeat.cs b/src/UniTask.NetCore/Linq/Repeat.cs
index a4f6415..42cdc10 100644
--- a/src/UniTask.NetCore/Linq/Repeat.cs
+++ b/src/UniTask.NetCore/Linq/Repeat.cs
@@ -50,7 +50,7 @@ namespace Cysharp.Threading.Tasks.Linq
 
             public UniTask<bool> MoveNextAsync()
             {
-                if (cancellationToken.IsCancellationRequested) return CompletedTasks.False;
+                cancellationToken.ThrowIfCancellationRequested();
 
                 if (remaining-- != 0)
                 {
diff --git a/src/UniTask.NetCoreSandbox/Program.cs b/src/UniTask.NetCoreSandbox/Program.cs
index 3d797d3..2aa34f9 100644
--- a/src/UniTask.NetCoreSandbox/Program.cs
+++ b/src/UniTask.NetCoreSandbox/Program.cs
@@ -32,17 +32,24 @@ namespace NetCoreSandbox
 
         static async Task Main(string[] args)
         {
+            var cts = new CancellationTokenSource();
 
 
+            await foreach (var item in UniTaskAsyncEnumerable.Range(1, 3).WithCancellation(cts.Token))
+            {
+                Console.WriteLine(item);
+                cts.Cancel();
+            }
 
 
-            await UniTaskAsyncEnumerable.Range(1, 3).ForEachAsync(x =>
+            /*
+            .ForEachAsync(x =>
            {
                if (x == 2) throw new Exception();
 
                Console.WriteLine(x);
            });
-
+           */