From 72694d4a44800140dd5cd90b69c7e68d2b28f0da Mon Sep 17 00:00:00 2001
From: Shimoda Takao <c3.hoge.fuga.piyo@gmail.com>
Date: Tue, 15 Apr 2025 11:29:14 +0900
Subject: [PATCH] Add SubscribeAwait<T, TState> overload

---
 src/ZeroMessenger/AnonymousMessageHandler.cs     | 10 +++++++++-
 src/ZeroMessenger/MessageSubscriberExtensions.cs |  8 +++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/ZeroMessenger/AnonymousMessageHandler.cs b/src/ZeroMessenger/AnonymousMessageHandler.cs
index addca0b..df5169e 100644
--- a/src/ZeroMessenger/AnonymousMessageHandler.cs
+++ b/src/ZeroMessenger/AnonymousMessageHandler.cs
@@ -22,4 +22,12 @@ protected override ValueTask HandleAsyncCore(T message, CancellationToken cancel
     {
         return handler(message, cancellationToken);
     }
-}
\ No newline at end of file
+}
+
+internal sealed class AnonymousAsyncMessageHandler<T, TState>(TState state, Func<T, TState, CancellationToken, ValueTask> handler) : AsyncMessageHandler<T>
+{
+    protected override ValueTask HandleAsyncCore(T message, CancellationToken cancellationToken = default)
+    {
+        return handler(message, state, cancellationToken);
+    }
+}
diff --git a/src/ZeroMessenger/MessageSubscriberExtensions.cs b/src/ZeroMessenger/MessageSubscriberExtensions.cs
index 200a47d..329f329 100644
--- a/src/ZeroMessenger/MessageSubscriberExtensions.cs
+++ b/src/ZeroMessenger/MessageSubscriberExtensions.cs
@@ -21,4 +21,10 @@ public static IDisposable SubscribeAwait<T>(this IMessageSubscriber<T> subscribe
     {
         return subscriber.SubscribeAwait(new AnonymousAsyncMessageHandler<T>(handler), subscribeStrategy);
     }
-}
\ No newline at end of file
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static IDisposable SubscribeAwait<T, TState>(this IMessageSubscriber<T> subscriber, TState state, Func<T, TState, CancellationToken, ValueTask> handler, AsyncSubscribeStrategy subscribeStrategy = AsyncSubscribeStrategy.Sequential)
+    {
+        return subscriber.SubscribeAwait(new AnonymousAsyncMessageHandler<T, TState>(state, handler), subscribeStrategy);
+    }
+}