Add SkipUntilCanceled, TakeUntilCanceled

src/UniTask.NetCoreTests/Linq/TakeInfinityTest.cs

@@ -4,6 +4,7 @@ using FluentAssertions;
 using System;
 using System.Collections.Generic;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Xunit;
 
@@ -40,5 +41,66 @@ namespace NetCoreTests.Linq
 
             (await xs).Should().BeEquivalentTo(1, 2, 3, 4);
         }
+
+        [Fact]
+        public async Task TakeUntil()
+        {
+            var cts = new CancellationTokenSource();
+
+            var rp = new AsyncReactiveProperty<int>(1);
+
+            var xs = rp.TakeUntilCanceled(cts.Token).ToArrayAsync();
+
+            var c = CancelAsync();
+
+            await c;
+            var foo = await xs;
+
+            foo.Should().BeEquivalentTo(new[] { 1, 10, 20 });
+
+            async Task CancelAsync()
+            {
+                rp.Value = 10;
+                await Task.Yield();
+                rp.Value = 20;
+                await Task.Yield();
+                cts.Cancel();
+                rp.Value = 30;
+                await Task.Yield();
+                rp.Value = 40;
+            }
+        }
+
+        [Fact]
+        public async Task SkipUntil()
+        {
+            var cts = new CancellationTokenSource();
+
+            var rp = new AsyncReactiveProperty<int>(1);
+
+            var xs = rp.SkipUntilCanceled(cts.Token).ToArrayAsync();
+
+            var c = CancelAsync();
+
+            await c;
+            var foo = await xs;
+
+            foo.Should().BeEquivalentTo(new[] { 30, 40 });
+
+            async Task CancelAsync()
+            {
+                rp.Value = 10;
+                await Task.Yield();
+                rp.Value = 20;
+                await Task.Yield();
+                cts.Cancel();
+                rp.Value = 30;
+                await Task.Yield();
+                rp.Value = 40;
+
+                rp.Dispose(); // complete.
+            }
+        }
+
     }
 }

src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/SkipUntilCanceled.cs

@@ -0,0 +1,142 @@
+using Cysharp.Threading.Tasks.Internal;
+using System;
+using System.Threading;
+
+namespace Cysharp.Threading.Tasks.Linq
+{
+    public static partial class UniTaskAsyncEnumerable
+    {
+        public static IUniTaskAsyncEnumerable<TSource> SkipUntilCanceled<TSource>(this IUniTaskAsyncEnumerable<TSource> source, CancellationToken cancellationToken)
+        {
+            Error.ThrowArgumentNullException(source, nameof(source));
+
+            return new SkipUntilCanceled<TSource>(source, cancellationToken);
+        }
+    }
+
+    internal sealed class SkipUntilCanceled<TSource> : IUniTaskAsyncEnumerable<TSource>
+    {
+        readonly IUniTaskAsyncEnumerable<TSource> source;
+        readonly CancellationToken cancellationToken;
+
+        public SkipUntilCanceled(IUniTaskAsyncEnumerable<TSource> source, CancellationToken cancellationToken)
+        {
+            this.source = source;
+            this.cancellationToken = cancellationToken;
+        }
+
+        public IUniTaskAsyncEnumerator<TSource> GetAsyncEnumerator(CancellationToken cancellationToken = default)
+        {
+            return new _SkipUntilCanceled(source, this.cancellationToken, cancellationToken);
+        }
+
+        sealed class _SkipUntilCanceled : MoveNextSource, IUniTaskAsyncEnumerator<TSource>
+        {
+            static readonly Action<object> MoveNextCoreDelegate = MoveNextCore;
+
+            readonly IUniTaskAsyncEnumerable<TSource> source;
+            CancellationToken cancellationToken1;
+            CancellationToken cancellationToken2;
+
+            bool isCanceled;
+            IUniTaskAsyncEnumerator<TSource> enumerator;
+            UniTask<bool>.Awaiter awaiter;
+            bool continueNext;
+
+            public _SkipUntilCanceled(IUniTaskAsyncEnumerable<TSource> source, CancellationToken cancellationToken1, CancellationToken cancellationToken2)
+            {
+                this.source = source;
+                this.cancellationToken1 = cancellationToken1;
+                this.cancellationToken2 = cancellationToken2;
+            }
+
+            public TSource Current { get; private set; }
+
+            public UniTask<bool> MoveNextAsync()
+            {
+                if (cancellationToken1.IsCancellationRequested) isCanceled = true;
+                if (cancellationToken2.IsCancellationRequested) isCanceled = true;
+
+                if (enumerator == null)
+                {
+                    enumerator = source.GetAsyncEnumerator(cancellationToken2); // use only AsyncEnumerator provided token.
+                }
+                completionSource.Reset();
+                SourceMoveNext();
+                return new UniTask<bool>(this, completionSource.Version);
+            }
+
+            void SourceMoveNext()
+            {
+                try
+                {
+                    LOOP:
+                    awaiter = enumerator.MoveNextAsync().GetAwaiter();
+                    if (awaiter.IsCompleted)
+                    {
+                        continueNext = true;
+                        MoveNextCore(this);
+                        if (continueNext)
+                        {
+                            continueNext = false;
+                            goto LOOP;
+                        }
+                    }
+                    else
+                    {
+                        awaiter.SourceOnCompleted(MoveNextCoreDelegate, this);
+                    }
+                }
+                catch (Exception ex)
+                {
+                    completionSource.TrySetException(ex);
+                }
+            }
+
+            static void MoveNextCore(object state)
+            {
+                var self = (_SkipUntilCanceled)state;
+
+                if (self.TryGetResult(self.awaiter, out var result))
+                {
+                    if (result)
+                    {
+                        AGAIN:
+
+                        if (self.isCanceled)
+                        {
+                            self.continueNext = false;
+                            self.Current = self.enumerator.Current;
+                            self.completionSource.TrySetResult(true);
+                        }
+                        else
+                        {
+                            if (self.cancellationToken1.IsCancellationRequested) self.isCanceled = true;
+                            if (self.cancellationToken2.IsCancellationRequested) self.isCanceled = true;
+
+                            if (self.isCanceled) goto AGAIN;
+
+                            if (!self.continueNext)
+                            {
+                                self.SourceMoveNext();
+                            }
+                        }
+                    }
+                    else
+                    {
+                        self.completionSource.TrySetResult(false);
+                    }
+                }
+            }
+
+            public UniTask DisposeAsync()
+            {
+                if (enumerator != null)
+                {
+                    return enumerator.DisposeAsync();
+                }
+                return default;
+            }
+        }
+    }
+}

src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/SkipUntilCanceled.cs.meta

@@ -0,0 +1,11 @@
+fileFormatVersion: 2
+guid: 4b1a778aef7150d47b93a49aa1bc34ae
+MonoImporter:
+  externalObjects: {}
+  serializedVersion: 2
+  defaultReferences: []
+  executionOrder: 0
+  icon: {instanceID: 0}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 

src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/TakeUntilCanceled.cs

@@ -0,0 +1,162 @@
+using Cysharp.Threading.Tasks.Internal;
+using System;
+using System.Threading;
+
+namespace Cysharp.Threading.Tasks.Linq
+{
+    public static partial class UniTaskAsyncEnumerable
+    {
+        public static IUniTaskAsyncEnumerable<TSource> TakeUntilCanceled<TSource>(this IUniTaskAsyncEnumerable<TSource> source, CancellationToken cancellationToken)
+        {
+            Error.ThrowArgumentNullException(source, nameof(source));
+
+            return new TakeUntilCanceled<TSource>(source, cancellationToken);
+        }
+    }
+
+    internal sealed class TakeUntilCanceled<TSource> : IUniTaskAsyncEnumerable<TSource>
+    {
+        readonly IUniTaskAsyncEnumerable<TSource> source;
+        readonly CancellationToken cancellationToken;
+
+        public TakeUntilCanceled(IUniTaskAsyncEnumerable<TSource> source, CancellationToken cancellationToken)
+        {
+            this.source = source;
+            this.cancellationToken = cancellationToken;
+        }
+
+        public IUniTaskAsyncEnumerator<TSource> GetAsyncEnumerator(CancellationToken cancellationToken = default)
+        {
+            return new _TakeUntilCanceled(source, this.cancellationToken, cancellationToken);
+        }
+
+        sealed class _TakeUntilCanceled : MoveNextSource, IUniTaskAsyncEnumerator<TSource>
+        {
+            static readonly Action<object> CancelDelegate1 = OnCanceled1;
+            static readonly Action<object> CancelDelegate2 = OnCanceled2;
+            static readonly Action<object> MoveNextCoreDelegate = MoveNextCore;
+
+            readonly IUniTaskAsyncEnumerable<TSource> source;
+            CancellationToken cancellationToken1;
+            CancellationToken cancellationToken2;
+            CancellationTokenRegistration cancellationTokenRegistration1;
+            CancellationTokenRegistration cancellationTokenRegistration2;
+
+            bool isCanceled;
+            IUniTaskAsyncEnumerator<TSource> enumerator;
+            UniTask<bool>.Awaiter awaiter;
+
+            public _TakeUntilCanceled(IUniTaskAsyncEnumerable<TSource> source, CancellationToken cancellationToken1, CancellationToken cancellationToken2)
+            {
+                this.source = source;
+                this.cancellationToken1 = cancellationToken1;
+                this.cancellationToken2 = cancellationToken2;
+
+                if (cancellationToken1.CanBeCanceled)
+                {
+                    this.cancellationTokenRegistration1 = cancellationToken1.RegisterWithoutCaptureExecutionContext(CancelDelegate1, this);
+                }
+
+                if (cancellationToken1 != cancellationToken2 && cancellationToken2.CanBeCanceled)
+                {
+                    this.cancellationTokenRegistration2 = cancellationToken2.RegisterWithoutCaptureExecutionContext(CancelDelegate2, this);
+                }
+            }
+
+            public TSource Current { get; private set; }
+
+            public UniTask<bool> MoveNextAsync()
+            {
+                if (cancellationToken1.IsCancellationRequested) isCanceled = true;
+                if (cancellationToken2.IsCancellationRequested) isCanceled = true;
+
+                if (enumerator == null)
+                {
+                    enumerator = source.GetAsyncEnumerator(cancellationToken2); // use only AsyncEnumerator provided token.
+                }
+
+                if (isCanceled) return CompletedTasks.False;
+
+                completionSource.Reset();
+                SourceMoveNext();
+                return new UniTask<bool>(this, completionSource.Version);
+            }
+
+            void SourceMoveNext()
+            {
+                try
+                {
+                    awaiter = enumerator.MoveNextAsync().GetAwaiter();
+                    if (awaiter.IsCompleted)
+                    {
+                        MoveNextCore(this);
+                    }
+                    else
+                    {
+                        awaiter.SourceOnCompleted(MoveNextCoreDelegate, this);
+                    }
+                }
+                catch (Exception ex)
+                {
+                    completionSource.TrySetException(ex);
+                }
+            }
+
+            static void MoveNextCore(object state)
+            {
+                var self = (_TakeUntilCanceled)state;
+
+                if (self.TryGetResult(self.awaiter, out var result))
+                {
+                    if (result)
+                    {
+                        if (self.isCanceled)
+                        {
+                            self.completionSource.TrySetResult(false);
+                        }
+                        else
+                        {
+                            self.Current = self.enumerator.Current;
+                            self.completionSource.TrySetResult(true);
+                        }
+                    }
+                    else
+                    {
+                        self.completionSource.TrySetResult(false);
+                    }
+                }
+            }
+
+            static void OnCanceled1(object state)
+            {
+                var self = (_TakeUntilCanceled)state;
+                if (!self.isCanceled)
+                {
+                    self.cancellationTokenRegistration2.Dispose();
+                    self.completionSource.TrySetResult(false);
+                }
+            }
+
+            static void OnCanceled2(object state)
+            {
+                var self = (_TakeUntilCanceled)state;
+                if (!self.isCanceled)
+                {
+                    self.cancellationTokenRegistration1.Dispose();
+                    self.completionSource.TrySetResult(false);
+                }
+            }
+
+            public UniTask DisposeAsync()
+            {
+                cancellationTokenRegistration1.Dispose();
+                cancellationTokenRegistration2.Dispose();
+                if (enumerator != null)
+                {
+                    return enumerator.DisposeAsync();
+                }
+                return default;
+            }
+        }
+    }
+}

src/UniTask/Assets/Plugins/UniTask/Runtime/Linq/TakeUntilCanceled.cs.meta

@@ -0,0 +1,11 @@
+fileFormatVersion: 2
+guid: e82f498cf3a1df04cbf646773fc11319
+MonoImporter:
+  externalObjects: {}
+  serializedVersion: 2
+  defaultReferences: []
+  executionOrder: 0
+  icon: {instanceID: 0}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: