From c592f833b7ebef4e1d40d2b30e79104ad285f0f2 Mon Sep 17 00:00:00 2001
From: mcatta <marcocatta@gmail.com>
Date: Mon, 10 Jul 2023 21:23:27 +0200
Subject: [PATCH] fix(reducer): ReducerFactory improve definition searching

---
 .../dev/mcatta/polpetta/ReducerFactory.kt     | 13 ++++-----
 .../polpetta/core/ReducerFactoryBuilder.kt    |  2 ++
 .../dev/mcatta/polpetta/StateStoreTest.kt     | 29 +++++++++++++++++++
 3 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/src/commonMain/kotlin/dev/mcatta/polpetta/ReducerFactory.kt b/src/commonMain/kotlin/dev/mcatta/polpetta/ReducerFactory.kt
index 2dba669..088748e 100644
--- a/src/commonMain/kotlin/dev/mcatta/polpetta/ReducerFactory.kt
+++ b/src/commonMain/kotlin/dev/mcatta/polpetta/ReducerFactory.kt
@@ -16,8 +16,7 @@ public abstract class ReducerFactory<A : Action, S : State, E : SideEffect>(
     sideEffectFactory: SideEffectFactory<E>
 ) {
 
-    private val _reducerDefinition: MutableMap<KClass<out State>, MutableList<ReducerFactoryBuilder<A, S>>> =
-        mutableMapOf()
+    private val _reducerDefinition: MutableList<ReducerFactoryBuilder<A, S>> = mutableListOf()
     private val _sideEffectFactory = sideEffectFactory
 
     /**
@@ -30,10 +29,9 @@ public abstract class ReducerFactory<A : Action, S : State, E : SideEffect>(
     internal fun <FromState : S> getReducer(
         action: A,
         fromState: FromState
-    ): Reducer<S>? =
-        _reducerDefinition[fromState::class]
-            ?.firstOrNull { item -> item.kClassAction == action::class }
-            ?.build(action)
+    ): Reducer<S>? = _reducerDefinition
+        .firstOrNull { item -> item.kClassAction == action::class && item.kClassFromState.isInstance(fromState) }
+        ?.build(action)
 
     /**
      * Define a [Reducer]'s body for the defined action [A]
@@ -58,9 +56,10 @@ public abstract class ReducerFactory<A : Action, S : State, E : SideEffect>(
         kClassFromState: KClass<FromState>,
         block: suspend SideEffectFactory<E>.(RA, StateModifier<FromState>) -> S
     ) {
-        _reducerDefinition.getOrPut(kClassFromState) { mutableListOf() }.add(
+        _reducerDefinition.add(
             ReducerFactoryBuilder(
                 kClassAction = kClassAction,
+                kClassFromState = kClassFromState,
                 handler = { action ->
                     @Suppress("UNCHECKED_CAST")
                     (reducer { state -> block(_sideEffectFactory, action as RA, state as StateModifier<FromState>) })
diff --git a/src/commonMain/kotlin/dev/mcatta/polpetta/core/ReducerFactoryBuilder.kt b/src/commonMain/kotlin/dev/mcatta/polpetta/core/ReducerFactoryBuilder.kt
index 780acb2..e4bd394 100644
--- a/src/commonMain/kotlin/dev/mcatta/polpetta/core/ReducerFactoryBuilder.kt
+++ b/src/commonMain/kotlin/dev/mcatta/polpetta/core/ReducerFactoryBuilder.kt
@@ -8,10 +8,12 @@ import kotlin.reflect.KClass
  * Reducer Factory Builder used to wrap the Reduce function based on an action
  *
  * @param kClassAction action type
+ * @param kClassFromState input state type
  * @param handler reducer
  */
 internal class ReducerFactoryBuilder<A : Action, S : State>(
     val kClassAction: KClass<out A>,
+    val kClassFromState: KClass<out S>,
     val handler: (A) -> Reducer<S>
 ) {
     fun build(action: A): Reducer<S> = handler.invoke(action)
diff --git a/src/commonTest/kotlin/dev/mcatta/polpetta/StateStoreTest.kt b/src/commonTest/kotlin/dev/mcatta/polpetta/StateStoreTest.kt
index eb781d0..af7bd4e 100644
--- a/src/commonTest/kotlin/dev/mcatta/polpetta/StateStoreTest.kt
+++ b/src/commonTest/kotlin/dev/mcatta/polpetta/StateStoreTest.kt
@@ -99,6 +99,35 @@ internal class StateStoreTest {
         }
     }
 
+    @Test
+    fun `Test action based on a Generic State`() = runTest(context = testScope.coroutineContext) {
+        // Given
+        val testStore = object : StateStore<TestAction, TestState, Nothing>(
+            coroutineScope = testScope,
+            initialState = TestState.Count(0),
+            reducerFactory = {
+                on<TestAction.Increase, TestState.Count> { _, stateModifier ->
+                    stateModifier.transform { copy(counter = counter + 1) }
+                }
+
+                on<TestAction.ToString, TestState> { _, stateModifier ->
+                    stateModifier.transform { TestState.Result(this.toString()) }
+                }
+            }
+        ) {}
+
+        testStore.stateFlow.test {
+            // When
+            testStore.dispatchAction(TestAction.Increase)
+            testStore.dispatchAction(TestAction.ToString)
+
+            // Then
+            assertIs<TestState.Count>(awaitItem())
+            assertIs<TestState.Count>(awaitItem())
+            assertIs<TestState.Result>(awaitItem())
+        }
+    }
+
     @Test
     fun `Test Intents execution and side effect`() = runTest(context = testScope.coroutineContext) {
         // Given