Reset ViewStore state when store changes

koma-compose/src/commonMain/kotlin/io/github/komakt/koma/compose/ViewStore.kt

@@ -5,6 +5,7 @@ import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.Stable
 import androidx.compose.runtime.collectAsState
+import androidx.compose.runtime.key
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberUpdatedState
@@ -141,7 +142,9 @@ fun <S : State, A : Action, E : Event> rememberViewStore(key: Any? = null, autoC
     val rememberedStore = holder.value ?: store().also { holder.value = it }
     val closeStoreOnDispose = remember(rememberedStore) { autoClose }
 
-    val state = rememberedStore.state.collectAsState()
+    val state = key(rememberedStore) {
+        rememberedStore.state.collectAsState()
+    }
 
     DisposableEffect(rememberedStore) {
         onDispose {

koma-compose/src/jvmTest/kotlin/io/github/komakt/koma/compose/ViewStoreJvmTest.kt

@@ -257,6 +257,51 @@ class ViewStoreJvmTest {
         }
     }
 
+    @Test
+    fun rememberViewStore_keyChangeUsesNewStoreStateOnFirstComposition() = runTest(testDispatcher) {
+        val firstStore = TestStore(UiState.Ready(1))
+        val secondStore = TestStore(UiState.Ready(2))
+        val renderedStates = mutableListOf<UiState>()
+        val frameClock = BroadcastFrameClock()
+        var frameTimeNanos = 0L
+        var key = "first"
+
+        suspend fun pumpFrame() {
+            testScheduler.runCurrent()
+            frameTimeNanos += 16_000_000L
+            frameClock.sendFrame(frameTimeNanos)
+            testScheduler.runCurrent()
+        }
+
+        withContext(frameClock) {
+            withRunningRecomposer { recomposer ->
+                val composition = Composition(NoOpApplier(), recomposer)
+                try {
+                    composition.setContent {
+                        renderedStates += rememberViewStore(key = key) {
+                            if (key == "first") firstStore else secondStore
+                        }.state
+                    }
+                    repeat(2) { pumpFrame() }
+
+                    renderedStates.clear()
+                    key = "second"
+                    composition.setContent {
+                        renderedStates += rememberViewStore(key = key) {
+                            if (key == "first") firstStore else secondStore
+                        }.state
+                    }
+                    repeat(2) { pumpFrame() }
+
+                    assertEquals(UiState.Ready(2), renderedStates.first())
+                } finally {
+                    composition.dispose()
+                    pumpFrame()
+                }
+            }
+        }
+    }
+
     @Test
     fun rememberViewStore_capturedCallbackReadsLatestState() = runTest(testDispatcher) {
         val store = TestStore(UiState.Loading)