Fix lifetime of Records' ViewModelStores

backstack/src/androidMain/kotlin/com/slack/circuit/backstack/ViewModelBackStackRecordLocalProvider.kt

@@ -30,6 +30,7 @@ import androidx.lifecycle.ViewModelStoreOwner
 import androidx.lifecycle.viewmodel.CreationExtras
 import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
 import androidx.lifecycle.viewmodel.compose.viewModel
+import kotlinx.collections.immutable.PersistentList
 import kotlinx.collections.immutable.persistentListOf
 
 private fun Context.findActivity(): Activity? {
@@ -56,25 +57,14 @@ internal object ViewModelBackStackRecordLocalProvider :
       )
     val viewModelStore = containerViewModel.viewModelStoreForKey(record.key)
     val activity = LocalContext.current.findActivity()
-    remember(record, viewModelStore) {
-      object : RememberObserver {
-        override fun onAbandoned() {
-          disposeIfNotChangingConfiguration()
-        }
-
-        override fun onForgotten() {
-          disposeIfNotChangingConfiguration()
-        }
-
-        override fun onRemembered() {}
-
-        fun disposeIfNotChangingConfiguration() {
+    val observer =
+      remember(record, viewModelStore) {
+        NestedRememberObserver {
           if (activity?.isChangingConfigurations != true) {
             containerViewModel.removeViewModelStoreOwnerForKey(record.key)?.clear()
           }
         }
       }
-    }
     return remember(viewModelStore) {
       val list =
         persistentListOf<ProvidedValue<*>>(
@@ -85,7 +75,11 @@ internal object ViewModelBackStackRecordLocalProvider :
         )
       @Suppress("ObjectLiteralToLambda")
       object : ProvidedValues {
-        @Composable override fun provideValues() = list
+        @Composable
+        override fun provideValues(): PersistentList<ProvidedValue<*>> {
+          remember { observer.UiRememberObserver() }
+          return list
+        }
       }
     }
   }
@@ -113,3 +107,47 @@ internal class BackStackRecordLocalProviderViewModel : ViewModel() {
     }
   }
 }
+
+private class NestedRememberObserver(
+  private val onCompletelyForgotten: () -> Unit,
+) : RememberObserver {
+  private var isRememberedForStack: Boolean = false
+    set(value) {
+      field = value
+      recomputeState()
+    }
+
+  private var isRememberedForUi: Boolean = false
+    set(value) {
+      field = value
+      recomputeState()
+    }
+
+  private fun recomputeState() {
+    if (!isRememberedForUi && !isRememberedForStack) {
+      onCompletelyForgotten()
+    }
+  }
+
+  inner class UiRememberObserver : RememberObserver {
+    override fun onRemembered() {
+      isRememberedForUi = true
+    }
+
+    override fun onAbandoned() = onForgotten()
+
+    override fun onForgotten() {
+      isRememberedForUi = false
+    }
+  }
+
+  override fun onRemembered() {
+    isRememberedForStack = true
+  }
+
+  override fun onAbandoned() = onForgotten()
+
+  override fun onForgotten() {
+    isRememberedForStack = false
+  }
+}

circuit-foundation/src/androidUnitTest/AndroidManifest.xml

@@ -1,5 +1,6 @@
 <manifest xmlns:android="http://schemas.android.com/apk/res/android">
     <application>
         <activity android:name="com.slack.circuit.foundation.NavigableCircuitRetainedStateTestActivity"/>
+        <activity android:name="com.slack.circuit.foundation.NavigableCircuitViewModelStateTestActivity"/>
     </application>
 </manifest>

circuit-foundation/src/androidUnitTest/kotlin/com/slack/circuit/foundation/NavigableCircuitViewModelStateAndroidTest.kt

@@ -0,0 +1,190 @@
+// Copyright (C) 2023 Slack Technologies, LLC
+// SPDX-License-Identifier: Apache-2.0
+package com.slack.circuit.foundation
+
+import androidx.compose.ui.test.SemanticsMatcher
+import androidx.compose.ui.test.assertAll
+import androidx.compose.ui.test.assertAny
+import androidx.compose.ui.test.assertCountEquals
+import androidx.compose.ui.test.assertTextEquals
+import androidx.compose.ui.test.hasTestTag
+import androidx.compose.ui.test.hasTextExactly
+import androidx.compose.ui.test.junit4.ComposeTestRule
+import androidx.compose.ui.test.junit4.createAndroidComposeRule
+import androidx.compose.ui.test.onAllNodesWithTag
+import androidx.compose.ui.test.onNodeWithTag
+import androidx.compose.ui.test.performClick
+import androidx.test.core.app.ActivityScenario
+import com.slack.circuit.internal.test.TestContentTags.TAG_COUNT
+import com.slack.circuit.internal.test.TestContentTags.TAG_GO_NEXT
+import com.slack.circuit.internal.test.TestContentTags.TAG_INCREASE_COUNT
+import com.slack.circuit.internal.test.TestContentTags.TAG_LABEL
+import com.slack.circuit.internal.test.TestContentTags.TAG_POP
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(ComposeUiTestRunner::class)
+class NavigableCircuitViewModelStateAndroidTest {
+
+  @get:Rule
+  val composeTestRule = createAndroidComposeRule<NavigableCircuitViewModelStateTestActivity>()
+
+  private val scenario: ActivityScenario<NavigableCircuitViewModelStateTestActivity>
+    get() = composeTestRule.activityRule.scenario
+
+  @Test
+  fun retainedStateScopedToBackstackWithRecreations() {
+    composeTestRule.run {
+      mainClock.autoAdvance = false
+
+      // Current: Screen A. Increase count to 1
+      onNodeWithTag(TAG_LABEL).assertTextEquals("A")
+      onNodeWithTag(TAG_COUNT).assertTextEquals("0")
+      onNodeWithTag(TAG_INCREASE_COUNT).performClick()
+      mainClock.advanceTimeByFrame()
+      onNodeWithTag(TAG_COUNT).assertTextEquals("1")
+
+      // Now recreate the Activity and assert that the values were retained
+      scenario.recreate()
+      onNodeWithTag(TAG_LABEL).assertTextEquals("A")
+      onNodeWithTag(TAG_COUNT).assertTextEquals("1")
+
+      // Navigate to Screen B. Increase count to 1
+      onNodeWithTag(TAG_GO_NEXT).performClick()
+      mainClock.advanceTimeBy(1_000)
+      onNodeWithTag(TAG_LABEL).assertTextEquals("B")
+      onNodeWithTag(TAG_COUNT).assertTextEquals("0")
+      onNodeWithTag(TAG_INCREASE_COUNT).performClick()
+      mainClock.advanceTimeByFrame()
+      onNodeWithTag(TAG_COUNT).assertTextEquals("1")
+
+      // Now recreate the Activity and assert that the values were retained
+      scenario.recreate()
+      onNodeWithTag(TAG_LABEL).assertTextEquals("B")
+      onNodeWithTag(TAG_COUNT).assertTextEquals("1")
+
+      // Navigate to Screen C. Increase count to 1
+      onNodeWithTag(TAG_GO_NEXT).performClick()
+      mainClock.advanceTimeBy(1_000)
+      onNodeWithTag(TAG_LABEL).assertTextEquals("C")
+      onNodeWithTag(TAG_COUNT).assertTextEquals("0")
+      onNodeWithTag(TAG_INCREASE_COUNT).performClick()
+      mainClock.advanceTimeByFrame()
+      onNodeWithTag(TAG_COUNT).assertTextEquals("1")
+
+      // Now recreate the Activity and assert that the values were retained
+      scenario.recreate()
+      onNodeWithTag(TAG_LABEL).assertTextEquals("C")
+      onNodeWithTag(TAG_COUNT).assertTextEquals("1")
+
+      // Pop to Screen B. Increase count from 1 to 2.
+      onNodeWithTag(TAG_POP).performClick()
+
+      // Part-way through pop, both screens should be visible
+      onEachFrameWhileMultipleScreens(hasTestTag(TAG_LABEL)) {
+        onAllNodesWithTag(TAG_LABEL)
+          .assertCountEquals(2)
+          .assertAny(hasTextExactly("C"))
+          .assertAny(hasTextExactly("B"))
+        onAllNodesWithTag(TAG_COUNT).assertCountEquals(2).assertAll(hasTextExactly("1"))
+      }
+      onNodeWithTag(TAG_LABEL).assertTextEquals("B")
+      onNodeWithTag(TAG_COUNT).assertTextEquals("1")
+      onNodeWithTag(TAG_INCREASE_COUNT).performClick()
+      mainClock.advanceTimeByFrame()
+      onNodeWithTag(TAG_COUNT).assertTextEquals("2")
+
+      // Navigate to Screen C. Assert that it's state was not retained
+      onNodeWithTag(TAG_GO_NEXT).performClick()
+
+      // Part-way through push, both screens should be visible
+      onEachFrameWhileMultipleScreens(hasTestTag(TAG_LABEL)) {
+        onAllNodesWithTag(TAG_LABEL)
+          .assertCountEquals(2)
+          .assertAny(hasTextExactly("C"))
+          .assertAny(hasTextExactly("B"))
+        onAllNodesWithTag(TAG_COUNT)
+          .assertCountEquals(2)
+          .assertAny(hasTextExactly("0"))
+          .assertAny(hasTextExactly("2"))
+      }
+      onNodeWithTag(TAG_LABEL).assertTextEquals("C")
+      onNodeWithTag(TAG_COUNT).assertTextEquals("0")
+
+      // Now recreate the Activity and assert that the values were retained
+      scenario.recreate()
+      onNodeWithTag(TAG_LABEL).assertTextEquals("C")
+      onNodeWithTag(TAG_COUNT).assertTextEquals("0")
+
+      // Pop to Screen B. Assert that it's state was retained
+      onNodeWithTag(TAG_POP).performClick()
+
+      // Part-way through pop, both screens should be visible
+      onEachFrameWhileMultipleScreens(hasTestTag(TAG_LABEL)) {
+        onAllNodesWithTag(TAG_LABEL)
+          .assertCountEquals(2)
+          .assertAny(hasTextExactly("C"))
+          .assertAny(hasTextExactly("B"))
+        onAllNodesWithTag(TAG_COUNT)
+          .assertCountEquals(2)
+          .assertAny(hasTextExactly("0"))
+          .assertAny(hasTextExactly("2"))
+      }
+      onNodeWithTag(TAG_LABEL).assertTextEquals("B")
+      onNodeWithTag(TAG_COUNT).assertTextEquals("2")
+
+      // Now recreate the Activity and assert that the values were retained
+      scenario.recreate()
+      onNodeWithTag(TAG_LABEL).assertTextEquals("B")
+      onNodeWithTag(TAG_COUNT).assertTextEquals("2")
+
+      // Pop to Screen A. Assert that it's state was retained
+      onNodeWithTag(TAG_POP).performClick()
+
+      // Part-way through pop, both screens should be visible
+      onEachFrameWhileMultipleScreens(hasTestTag(TAG_LABEL)) {
+        onAllNodesWithTag(TAG_LABEL)
+          .assertCountEquals(2)
+          .assertAny(hasTextExactly("B"))
+          .assertAny(hasTextExactly("A"))
+        onAllNodesWithTag(TAG_COUNT)
+          .assertCountEquals(2)
+          .assertAny(hasTextExactly("2"))
+          .assertAny(hasTextExactly("1"))
+      }
+      onNodeWithTag(TAG_LABEL).assertTextEquals("A")
+      onNodeWithTag(TAG_COUNT).assertTextEquals("1")
+
+      // Now recreate the Activity and assert that the values were retained
+      scenario.recreate()
+      onNodeWithTag(TAG_LABEL).assertTextEquals("A")
+      onNodeWithTag(TAG_COUNT).assertTextEquals("1")
+
+      // Navigate to Screen B. Assert that it's state was not retained
+      onNodeWithTag(TAG_GO_NEXT).performClick()
+      mainClock.advanceTimeBy(1_000)
+      onNodeWithTag(TAG_LABEL).assertTextEquals("B")
+      onNodeWithTag(TAG_COUNT).assertTextEquals("0")
+    }
+  }
+
+  private fun ComposeTestRule.onEachFrameWhileMultipleScreens(
+    matcher: SemanticsMatcher,
+    block: ComposeTestRule.() -> Unit
+  ) {
+    var i = 0
+    while (true) {
+      mainClock.advanceTimeByFrame()
+      if (onAllNodes(matcher).fetchSemanticsNodes().size <= 1) {
+        break
+      }
+      try {
+        block()
+      } catch (e: Throwable) {
+        throw AssertionError("Error on frame $i", e)
+      }
+      i++
+    }
+  }
+}

circuit-foundation/src/androidUnitTest/kotlin/com/slack/circuit/foundation/NavigableCircuitViewModelStateTestActivity.kt

@@ -0,0 +1,32 @@
+// Copyright (C) 2023 Slack Technologies, LLC
+// SPDX-License-Identifier: Apache-2.0
+package com.slack.circuit.foundation
+
+import android.os.Bundle
+import androidx.activity.ComponentActivity
+import androidx.activity.compose.setContent
+import com.slack.circuit.backstack.rememberSaveableBackStack
+import com.slack.circuit.internal.test.TestCountPresenter
+import com.slack.circuit.internal.test.TestScreen
+import com.slack.circuit.internal.test.createTestCircuit
+
+class NavigableCircuitViewModelStateTestActivity : ComponentActivity() {
+
+  override fun onCreate(savedInstanceState: Bundle?) {
+    super.onCreate(savedInstanceState)
+
+    val circuit = createTestCircuit(rememberType = TestCountPresenter.RememberType.ViewModel)
+
+    setContent {
+      CircuitCompositionLocals(circuit) {
+        val backstack = rememberSaveableBackStack { push(TestScreen.ScreenA) }
+        val navigator =
+          rememberCircuitNavigator(
+            backstack = backstack,
+            onRootPop = {} // no-op for tests
+          )
+        NavigableCircuitContent(navigator = navigator, backstack = backstack)
+      }
+    }
+  }
+}