From edf8ee6c123316614efca04edf89406b4f2b902a Mon Sep 17 00:00:00 2001
From: neuecc <ils@neue.cc>
Date: Sat, 6 Jan 2024 15:22:20 +0900
Subject: [PATCH] FromEvent supports CancellationToken

---
 src/R3/Factories/FromEvent.cs                | 70 +++++++++++++++-----
 tests/R3.Tests/FactoryTests/FromEventTest.cs | 34 ++++++++++
 2 files changed, 86 insertions(+), 18 deletions(-)

diff --git a/src/R3/Factories/FromEvent.cs b/src/R3/Factories/FromEvent.cs
index cf4e9bd5..6eb2c749 100644
--- a/src/R3/Factories/FromEvent.cs
+++ b/src/R3/Factories/FromEvent.cs
@@ -2,43 +2,43 @@
 
 public static partial class Observable
 {
-    public static Observable<(object? sender, EventArgs e)> FromEventHandler(Action<EventHandler> addHandler, Action<EventHandler> removeHandler)
+    public static Observable<(object? sender, EventArgs e)> FromEventHandler(Action<EventHandler> addHandler, Action<EventHandler> removeHandler, CancellationToken cancellationToken = default)
     {
-        return new FromEvent<EventHandler, (object? sender, EventArgs e)>(h => (sender, e) => h((sender, e)), addHandler, removeHandler);
+        return new FromEvent<EventHandler, (object? sender, EventArgs e)>(h => (sender, e) => h((sender, e)), addHandler, removeHandler, cancellationToken);
     }
 
-    public static Observable<(object? sender, TEventArgs e)> FromEventHandler<TEventArgs>(Action<EventHandler<TEventArgs>> addHandler, Action<EventHandler<TEventArgs>> removeHandler)
+    public static Observable<(object? sender, TEventArgs e)> FromEventHandler<TEventArgs>(Action<EventHandler<TEventArgs>> addHandler, Action<EventHandler<TEventArgs>> removeHandler, CancellationToken cancellationToken = default)
     {
-        return new FromEvent<EventHandler<TEventArgs>, (object? sender, TEventArgs e)>(h => (sender, e) => h((sender, e)), addHandler, removeHandler);
+        return new FromEvent<EventHandler<TEventArgs>, (object? sender, TEventArgs e)>(h => (sender, e) => h((sender, e)), addHandler, removeHandler, cancellationToken);
     }
 
-    public static Observable<Unit> FromEvent(Action<Action> addHandler, Action<Action> removeHandler)
+    public static Observable<Unit> FromEvent(Action<Action> addHandler, Action<Action> removeHandler, CancellationToken cancellationToken = default)
     {
-        return new FromEvent<Action>(static h => h, addHandler, removeHandler);
+        return new FromEvent<Action>(static h => h, addHandler, removeHandler, cancellationToken);
     }
 
-    public static Observable<T> FromEvent<T>(Action<Action<T>> addHandler, Action<Action<T>> removeHandler)
+    public static Observable<T> FromEvent<T>(Action<Action<T>> addHandler, Action<Action<T>> removeHandler, CancellationToken cancellationToken = default)
     {
-        return new FromEvent<Action<T>, T>(static h => h, addHandler, removeHandler);
+        return new FromEvent<Action<T>, T>(static h => h, addHandler, removeHandler, cancellationToken);
     }
 
-    public static Observable<Unit> FromEvent<TDelegate>(Func<Action, TDelegate> conversion, Action<TDelegate> addHandler, Action<TDelegate> removeHandler)
+    public static Observable<Unit> FromEvent<TDelegate>(Func<Action, TDelegate> conversion, Action<TDelegate> addHandler, Action<TDelegate> removeHandler, CancellationToken cancellationToken = default)
     {
-        return new FromEvent<TDelegate>(conversion, addHandler, removeHandler);
+        return new FromEvent<TDelegate>(conversion, addHandler, removeHandler, cancellationToken);
     }
 
-    public static Observable<T> FromEvent<TDelegate, T>(Func<Action<T>, TDelegate> conversion, Action<TDelegate> addHandler, Action<TDelegate> removeHandler)
+    public static Observable<T> FromEvent<TDelegate, T>(Func<Action<T>, TDelegate> conversion, Action<TDelegate> addHandler, Action<TDelegate> removeHandler, CancellationToken cancellationToken = default)
     {
-        return new FromEvent<TDelegate, T>(conversion, addHandler, removeHandler);
+        return new FromEvent<TDelegate, T>(conversion, addHandler, removeHandler, cancellationToken);
     }
 }
 
-internal sealed class FromEvent<TDelegate>(Func<Action, TDelegate> conversion, Action<TDelegate> addHandler, Action<TDelegate> removeHandler)
+internal sealed class FromEvent<TDelegate>(Func<Action, TDelegate> conversion, Action<TDelegate> addHandler, Action<TDelegate> removeHandler, CancellationToken cancellationToken)
     : Observable<Unit>
 {
     protected override IDisposable SubscribeCore(Observer<Unit> observer)
     {
-        return new _FromEventPattern(conversion, addHandler, removeHandler, observer);
+        return new _FromEventPattern(conversion, addHandler, removeHandler, observer, cancellationToken);
     }
 
     sealed class _FromEventPattern : IDisposable
@@ -46,13 +46,23 @@ sealed class _FromEventPattern : IDisposable
         Observer<Unit>? observer;
         Action<TDelegate>? removeHandler;
         TDelegate registeredHandler;
+        CancellationTokenRegistration cancellationTokenRegistration;
 
-        public _FromEventPattern(Func<Action, TDelegate> conversion, Action<TDelegate> addHandler, Action<TDelegate> removeHandler, Observer<Unit> observer)
+        public _FromEventPattern(Func<Action, TDelegate> conversion, Action<TDelegate> addHandler, Action<TDelegate> removeHandler, Observer<Unit> observer, CancellationToken cancellationToken)
         {
             this.observer = observer;
             this.removeHandler = removeHandler;
             this.registeredHandler = conversion(OnNext);
             addHandler(this.registeredHandler);
+
+            if (cancellationToken.CanBeCanceled)
+            {
+                this.cancellationTokenRegistration = cancellationToken.UnsafeRegister(static state =>
+                {
+                    var s = (_FromEventPattern)state!;
+                    s.CompleteDispose();
+                }, this);
+            }
         }
 
         void OnNext()
@@ -60,6 +70,12 @@ void OnNext()
             observer?.OnNext(default);
         }
 
+        void CompleteDispose()
+        {
+            observer?.OnCompleted();
+            Dispose();
+        }
+
         public void Dispose()
         {
             var handler = Interlocked.Exchange(ref removeHandler, null);
@@ -67,18 +83,19 @@ public void Dispose()
             {
                 observer = null;
                 removeHandler = null;
+                cancellationTokenRegistration.Dispose();
                 handler(this.registeredHandler);
             }
         }
     }
 }
 
-internal sealed class FromEvent<TDelegate, T>(Func<Action<T>, TDelegate> conversion, Action<TDelegate> addHandler, Action<TDelegate> removeHandler)
+internal sealed class FromEvent<TDelegate, T>(Func<Action<T>, TDelegate> conversion, Action<TDelegate> addHandler, Action<TDelegate> removeHandler, CancellationToken cancellationToken)
     : Observable<T>
 {
     protected override IDisposable SubscribeCore(Observer<T> observer)
     {
-        return new _FromEventPattern(conversion, addHandler, removeHandler, observer);
+        return new _FromEventPattern(conversion, addHandler, removeHandler, observer, cancellationToken);
     }
 
     sealed class _FromEventPattern : IDisposable
@@ -86,13 +103,23 @@ sealed class _FromEventPattern : IDisposable
         Observer<T>? observer;
         Action<TDelegate>? removeHandler;
         TDelegate registeredHandler;
+        CancellationTokenRegistration cancellationTokenRegistration;
 
-        public _FromEventPattern(Func<Action<T>, TDelegate> conversion, Action<TDelegate> addHandler, Action<TDelegate> removeHandler, Observer<T> observer)
+        public _FromEventPattern(Func<Action<T>, TDelegate> conversion, Action<TDelegate> addHandler, Action<TDelegate> removeHandler, Observer<T> observer, CancellationToken cancellationToken)
         {
             this.observer = observer;
             this.removeHandler = removeHandler;
             this.registeredHandler = conversion(OnNext);
             addHandler(this.registeredHandler);
+
+            if (cancellationToken.CanBeCanceled)
+            {
+                this.cancellationTokenRegistration = cancellationToken.UnsafeRegister(static state =>
+                {
+                    var s = (_FromEventPattern)state!;
+                    s.CompleteDispose();
+                }, this);
+            }
         }
 
         void OnNext(T value)
@@ -100,6 +127,12 @@ void OnNext(T value)
             observer?.OnNext(value);
         }
 
+        void CompleteDispose()
+        {
+            observer?.OnCompleted();
+            Dispose();
+        }
+
         public void Dispose()
         {
             var handler = Interlocked.Exchange(ref removeHandler, null);
@@ -107,6 +140,7 @@ public void Dispose()
             {
                 observer = null;
                 removeHandler = null;
+                cancellationTokenRegistration.Dispose();
                 handler(this.registeredHandler);
             }
         }
diff --git a/tests/R3.Tests/FactoryTests/FromEventTest.cs b/tests/R3.Tests/FactoryTests/FromEventTest.cs
index 062df3df..3ca9391f 100644
--- a/tests/R3.Tests/FactoryTests/FromEventTest.cs
+++ b/tests/R3.Tests/FactoryTests/FromEventTest.cs
@@ -41,6 +41,40 @@ public void Event()
         ev.InvocationListCount().Should().Be((0, 0, 0, 0, 0, 0, 0));
     }
 
+    [Fact]
+    public void Cancel()
+    {
+        var cts = new CancellationTokenSource();
+
+        var ev = new EventPattern();
+
+        var l1 = Observable.FromEventHandler(h => ev.E1 += h, h => ev.E1 -= h, cts.Token).ToLiveList();
+        var l2 = Observable.FromEventHandler<int>(h => ev.E2 += h, h => ev.E2 -= h, cts.Token).ToLiveList();
+        var l3 = Observable.FromEvent(h => ev.A1 += h, h => ev.A1 -= h, cts.Token).ToLiveList();
+        var l4 = Observable.FromEvent<int>(h => ev.A2 += h, h => ev.A2 -= h, cts.Token).ToLiveList();
+        var l5 = Observable.FromEvent<MyDelegate1>(h => new MyDelegate1(h), h => ev.M1 += h, h => ev.M1 -= h, cts.Token).ToLiveList();
+        var l6 = Observable.FromEvent<MyDelegate2, int>(h => new MyDelegate2(h), h => ev.M2 += h, h => ev.M2 -= h, cts.Token).ToLiveList();
+        var l7 = Observable.FromEvent<MyDelegate3, (int x, int y)>(h => (x, y) => h((x, y)), h => ev.M3 += h, h => ev.M3 -= h, cts.Token).ToLiveList();
+
+        ev.Raise(10, 20);
+        ev.Raise(100, 200);
+
+        l1.Should().HaveCount(2);
+        l3.Should().HaveCount(2);
+        l5.Should().HaveCount(2);
+
+        l2.Select(x => x.e).Should().Equal([10, 100]);
+        l4.AssertEqual([10, 100]);
+        l6.AssertEqual([10, 100]);
+        l7.AssertEqual([(10, 20), (100, 200)]);
+
+        ev.InvocationListCount().Should().Be((1, 1, 1, 1, 1, 1, 1));
+
+        cts.Cancel();
+
+        ev.InvocationListCount().Should().Be((0, 0, 0, 0, 0, 0, 0));
+    }
+
 
     class EventPattern
     {