Add UniTaskAsyncEnumerable.SkipUntil, TakeUntil. Fix SkipUntilCanceled behaviour.

src/UniTask.NetCoreTests/Linq/TakeInfinityTest.cs

@@ -43,7 +43,7 @@ namespace NetCoreTests.Linq
         }
 
         [Fact]
-        public async Task TakeUntil()
+        public async Task TakeUntilCanceled()
         {
             var cts = new CancellationTokenSource();
 
@@ -72,7 +72,7 @@ namespace NetCoreTests.Linq
         }
 
         [Fact]
-        public async Task SkipUntil()
+        public async Task SkipUntilCanceled()
         {
             var cts = new CancellationTokenSource();
 
@@ -85,7 +85,7 @@ namespace NetCoreTests.Linq
             await c;
             var foo = await xs;
 
-            foo.Should().BeEquivalentTo(new[] { 30, 40 });
+            foo.Should().BeEquivalentTo(new[] { 20, 30, 40 });
 
             async Task CancelAsync()
             {
@@ -102,5 +102,64 @@ namespace NetCoreTests.Linq
             }
         }
 
+        [Fact]
+        public async Task TakeUntil()
+        {
+            var cts = new AsyncReactiveProperty<int>(0);
+
+            var rp = new AsyncReactiveProperty<int>(1);
+
+            var xs = rp.TakeUntil(cts.WaitAsync()).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.Value = 9999;
+                rp.Value = 30;
+                await Task.Yield();
+                rp.Value = 40;
+            }
+        }
+
+        [Fact]
+        public async Task SkipUntil()
+        {
+            var cts = new AsyncReactiveProperty<int>(0);
+
+            var rp = new AsyncReactiveProperty<int>(1);
+
+            var xs = rp.SkipUntil(cts.WaitAsync()).ToArrayAsync();
+
+            var c = CancelAsync();
+
+            await c;
+            var foo = await xs;
+
+            foo.Should().BeEquivalentTo(new[] { 20, 30, 40 });
+
+            async Task CancelAsync()
+            {
+                rp.Value = 10;
+                await Task.Yield();
+                rp.Value = 20;
+                await Task.Yield();
+                cts.Value = 9999;
+                rp.Value = 30;
+                await Task.Yield();
+                rp.Value = 40;
+
+                rp.Dispose(); // complete.
+            }
+        }
     }
 }

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

@@ -0,0 +1,187 @@
+using Cysharp.Threading.Tasks.Internal;
+using System;
+using System.Threading;
+
+namespace Cysharp.Threading.Tasks.Linq
+{
+    public static partial class UniTaskAsyncEnumerable
+    {
+        public static IUniTaskAsyncEnumerable<TSource> SkipUntil<TSource>(this IUniTaskAsyncEnumerable<TSource> source, UniTask other)
+        {
+            Error.ThrowArgumentNullException(source, nameof(source));
+
+            return new SkipUntil<TSource>(source, other, null);
+        }
+
+        public static IUniTaskAsyncEnumerable<TSource> SkipUntil<TSource>(this IUniTaskAsyncEnumerable<TSource> source, Func<CancellationToken, UniTask> other)
+        {
+            Error.ThrowArgumentNullException(source, nameof(source));
+            Error.ThrowArgumentNullException(source, nameof(other));
+
+            return new SkipUntil<TSource>(source, default, other);
+        }
+    }
+
+    internal sealed class SkipUntil<TSource> : IUniTaskAsyncEnumerable<TSource>
+    {
+        readonly IUniTaskAsyncEnumerable<TSource> source;
+        readonly UniTask other;
+        readonly Func<CancellationToken, UniTask> other2;
+
+        public SkipUntil(IUniTaskAsyncEnumerable<TSource> source, UniTask other, Func<CancellationToken, UniTask> other2)
+        {
+            this.source = source;
+            this.other = other;
+            this.other2 = other2;
+        }
+
+        public IUniTaskAsyncEnumerator<TSource> GetAsyncEnumerator(CancellationToken cancellationToken = default)
+        {
+            if (other2 != null)
+            {
+                return new _SkipUntil(source, this.other2(cancellationToken), cancellationToken);
+            }
+            else
+            {
+                return new _SkipUntil(source, this.other, cancellationToken);
+            }
+        }
+
+        sealed class _SkipUntil : MoveNextSource, IUniTaskAsyncEnumerator<TSource>
+        {
+            static readonly Action<object> CancelDelegate1 = OnCanceled1;
+            static readonly Action<object> MoveNextCoreDelegate = MoveNextCore;
+
+            readonly IUniTaskAsyncEnumerable<TSource> source;
+            CancellationToken cancellationToken1;
+
+            bool completed;
+            CancellationTokenRegistration cancellationTokenRegistration1;
+            IUniTaskAsyncEnumerator<TSource> enumerator;
+            UniTask<bool>.Awaiter awaiter;
+            bool continueNext;
+            Exception exception;
+
+            public _SkipUntil(IUniTaskAsyncEnumerable<TSource> source, UniTask other, CancellationToken cancellationToken1)
+            {
+                this.source = source;
+                this.cancellationToken1 = cancellationToken1;
+                if (cancellationToken1.CanBeCanceled)
+                {
+                    this.cancellationTokenRegistration1 = cancellationToken1.RegisterWithoutCaptureExecutionContext(CancelDelegate1, this);
+                }
+
+                TaskTracker.TrackActiveTask(this, 3);
+                RunOther(other).Forget();
+            }
+
+            public TSource Current { get; private set; }
+
+            public UniTask<bool> MoveNextAsync()
+            {
+                if (exception != null)
+                {
+                    return UniTask.FromException<bool>(exception);
+                }
+
+                if (cancellationToken1.IsCancellationRequested)
+                {
+                    return UniTask.FromCanceled<bool>(cancellationToken1);
+                }
+
+                if (enumerator == null)
+                {
+                    enumerator = source.GetAsyncEnumerator(cancellationToken1);
+                }
+                completionSource.Reset();
+
+                if (completed)
+                {
+                    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 = (_SkipUntil)state;
+
+                if (self.TryGetResult(self.awaiter, out var result))
+                {
+                    if (result)
+                    {
+                        self.Current = self.enumerator.Current;
+                        self.completionSource.TrySetResult(true);
+                        if (self.continueNext)
+                        {
+                            self.SourceMoveNext();
+                        }
+                    }
+                    else
+                    {
+                        self.completionSource.TrySetResult(false);
+                    }
+                }
+            }
+
+            async UniTaskVoid RunOther(UniTask other)
+            {
+                try
+                {
+                    await other;
+                    completed = true;
+                    SourceMoveNext();
+                }
+                catch (Exception ex)
+                {
+                    exception = ex;
+                    completionSource.TrySetException(ex);
+                }
+            }
+
+            static void OnCanceled1(object state)
+            {
+                var self = (_SkipUntil)state;
+                self.completionSource.TrySetCanceled(self.cancellationToken1);
+            }
+
+            public UniTask DisposeAsync()
+            {
+                TaskTracker.RemoveTracking(this);
+                cancellationTokenRegistration1.Dispose();
+                if (enumerator != null)
+                {
+                    return enumerator.DisposeAsync();
+                }
+                return default;
+            }
+        }
+    }
+}

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

@@ -32,13 +32,17 @@ namespace Cysharp.Threading.Tasks.Linq
 
         sealed class _SkipUntilCanceled : 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;
+            int isCanceled;
             IUniTaskAsyncEnumerator<TSource> enumerator;
             UniTask<bool>.Awaiter awaiter;
             bool continueNext;
@@ -48,6 +52,14 @@ namespace Cysharp.Threading.Tasks.Linq
                 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);
+                }
                 TaskTracker.TrackActiveTask(this, 3);
             }
 
@@ -55,15 +67,18 @@ namespace Cysharp.Threading.Tasks.Linq
 
             public UniTask<bool> MoveNextAsync()
             {
-                if (cancellationToken1.IsCancellationRequested) isCanceled = true;
-                if (cancellationToken2.IsCancellationRequested) isCanceled = true;
-
                 if (enumerator == null)
                 {
+                    if (cancellationToken1.IsCancellationRequested) isCanceled = 1;
+                    if (cancellationToken2.IsCancellationRequested) isCanceled = 1;
                     enumerator = source.GetAsyncEnumerator(cancellationToken2); // use only AsyncEnumerator provided token.
                 }
                 completionSource.Reset();
-                SourceMoveNext();
+
+                if (isCanceled != 0)
+                {
+                    SourceMoveNext();
+                }
                 return new UniTask<bool>(this, completionSource.Version);
             }
 
@@ -102,25 +117,11 @@ namespace Cysharp.Threading.Tasks.Linq
                 {
                     if (result)
                     {
-                        AGAIN:
+                        self.Current = self.enumerator.Current;
-
+                        self.completionSource.TrySetResult(true);
-                        if (self.isCanceled)
+                        if (self.continueNext)
                         {
-                            self.continueNext = false;
+                            self.SourceMoveNext();
-                            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
@@ -130,9 +131,37 @@ namespace Cysharp.Threading.Tasks.Linq
                 }
             }
 
+            static void OnCanceled1(object state)
+            {
+                var self = (_SkipUntilCanceled)state;
+                if (self.isCanceled == 0)
+                {
+                    if (Interlocked.Increment(ref self.isCanceled) == 1)
+                    {
+                        self.cancellationTokenRegistration2.Dispose();
+                        self.SourceMoveNext();
+                    }
+                }
+            }
+
+            static void OnCanceled2(object state)
+            {
+                var self = (_SkipUntilCanceled)state;
+                if (self.isCanceled == 0)
+                {
+                    if (Interlocked.Increment(ref self.isCanceled) == 1)
+                    {
+                        self.cancellationTokenRegistration2.Dispose();
+                        self.SourceMoveNext();
+                    }
+                }
+            }
+
             public UniTask DisposeAsync()
             {
                 TaskTracker.RemoveTracking(this);
+                cancellationTokenRegistration1.Dispose();
+                cancellationTokenRegistration2.Dispose();
                 if (enumerator != null)
                 {
                     return enumerator.DisposeAsync();

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

@@ -0,0 +1,190 @@
+using Cysharp.Threading.Tasks.Internal;
+using System;
+using System.Threading;
+
+namespace Cysharp.Threading.Tasks.Linq
+{
+    public static partial class UniTaskAsyncEnumerable
+    {
+        public static IUniTaskAsyncEnumerable<TSource> TakeUntil<TSource>(this IUniTaskAsyncEnumerable<TSource> source, UniTask other)
+        {
+            Error.ThrowArgumentNullException(source, nameof(source));
+
+            return new TakeUntil<TSource>(source, other, null);
+        }
+
+        public static IUniTaskAsyncEnumerable<TSource> TakeUntil<TSource>(this IUniTaskAsyncEnumerable<TSource> source, Func<CancellationToken, UniTask> other)
+        {
+            Error.ThrowArgumentNullException(source, nameof(source));
+            Error.ThrowArgumentNullException(source, nameof(other));
+
+            return new TakeUntil<TSource>(source, default, other);
+        }
+    }
+
+    internal sealed class TakeUntil<TSource> : IUniTaskAsyncEnumerable<TSource>
+    {
+        readonly IUniTaskAsyncEnumerable<TSource> source;
+        readonly UniTask other;
+        readonly Func<CancellationToken, UniTask> other2;
+
+        public TakeUntil(IUniTaskAsyncEnumerable<TSource> source, UniTask other, Func<CancellationToken, UniTask> other2)
+        {
+            this.source = source;
+            this.other = other;
+            this.other2 = other2;
+        }
+
+        public IUniTaskAsyncEnumerator<TSource> GetAsyncEnumerator(CancellationToken cancellationToken = default)
+        {
+            if (other2 != null)
+            {
+                return new _TakeUntil(source, this.other2(cancellationToken), cancellationToken);
+            }
+            else
+            {
+                return new _TakeUntil(source, this.other, cancellationToken);
+            }
+        }
+
+        sealed class _TakeUntil : MoveNextSource, IUniTaskAsyncEnumerator<TSource>
+        {
+            static readonly Action<object> CancelDelegate1 = OnCanceled1;
+            static readonly Action<object> MoveNextCoreDelegate = MoveNextCore;
+
+            readonly IUniTaskAsyncEnumerable<TSource> source;
+            CancellationToken cancellationToken1;
+            CancellationTokenRegistration cancellationTokenRegistration1;
+
+            bool completed;
+            Exception exception;
+            IUniTaskAsyncEnumerator<TSource> enumerator;
+            UniTask<bool>.Awaiter awaiter;
+
+            public _TakeUntil(IUniTaskAsyncEnumerable<TSource> source, UniTask other, CancellationToken cancellationToken1)
+            {
+                this.source = source;
+                this.cancellationToken1 = cancellationToken1;
+
+                if (cancellationToken1.CanBeCanceled)
+                {
+                    this.cancellationTokenRegistration1 = cancellationToken1.RegisterWithoutCaptureExecutionContext(CancelDelegate1, this);
+                }
+
+                TaskTracker.TrackActiveTask(this, 3);
+
+                RunOther(other).Forget();
+            }
+
+            public TSource Current { get; private set; }
+
+            public UniTask<bool> MoveNextAsync()
+            {
+                if (completed)
+                {
+                    return CompletedTasks.False;
+                }
+
+                if (exception != null)
+                {
+                    return UniTask.FromException<bool>(exception);
+                }
+
+                if (cancellationToken1.IsCancellationRequested)
+                {
+                    return UniTask.FromCanceled<bool>(cancellationToken1);
+                }
+
+                if (enumerator == null)
+                {
+                    enumerator = source.GetAsyncEnumerator(cancellationToken1);
+                }
+
+                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 = (_TakeUntil)state;
+
+                if (self.TryGetResult(self.awaiter, out var result))
+                {
+                    if (result)
+                    {
+                        if (self.exception != null)
+                        {
+                            self.completionSource.TrySetException(self.exception);
+                        }
+                        else if (self.cancellationToken1.IsCancellationRequested)
+                        {
+                            self.completionSource.TrySetCanceled(self.cancellationToken1);
+                        }
+                        else
+                        {
+                            self.Current = self.enumerator.Current;
+                            self.completionSource.TrySetResult(true);
+                        }
+                    }
+                    else
+                    {
+                        self.completionSource.TrySetResult(false);
+                    }
+                }
+            }
+
+            async UniTaskVoid RunOther(UniTask other)
+            {
+                try
+                {
+                    await other;
+                    completed = true;
+                    completionSource.TrySetResult(false);
+                }
+                catch (Exception ex)
+                {
+                    exception = ex;
+                    completionSource.TrySetException(ex);
+                }
+            }
+
+            static void OnCanceled1(object state)
+            {
+                var self = (_TakeUntil)state;
+                self.completionSource.TrySetCanceled(self.cancellationToken1);
+            }
+
+            public UniTask DisposeAsync()
+            {
+                TaskTracker.RemoveTracking(this);
+                cancellationTokenRegistration1.Dispose();
+                if (enumerator != null)
+                {
+                    return enumerator.DisposeAsync();
+                }
+                return default;
+            }
+        }
+    }
+}

src/UniTask/Assets/Scenes/SandboxMain.cs

@@ -119,7 +119,7 @@ public class AsyncMessageBroker<T> : IDisposable
 
 public class SandboxMain : MonoBehaviour
 {
-    public Camera camera;
+    public Camera mycamera;
 
     public Button okButton;
     public Button cancelButton;
@@ -998,11 +998,11 @@ public class SandboxMain : MonoBehaviour
         var height = 240;
         var depth = 24;
 
-        camera.targetTexture = new RenderTexture(width, height, depth, RenderTextureFormat.ARGB32, RenderTextureReadWrite.Default)
+        mycamera.targetTexture = new RenderTexture(width, height, depth, RenderTextureFormat.ARGB32, RenderTextureReadWrite.Default)
         {
             antiAliasing = 8
         };
-        camera.enabled = true;
+        mycamera.enabled = true;
 
         //myRenderTexture = new RenderTexture(width, height, depth, RenderTextureFormat.ARGB32, RenderTextureReadWrite.Default)
         //{
@@ -1014,7 +1014,7 @@ public class SandboxMain : MonoBehaviour
 
     async UniTask ShootAsync()
     {
-        var rt = camera.targetTexture;
+        var rt = mycamera.targetTexture;