Change rememberRetain not to retain the value of removed node

circuit-foundation/src/commonMain/kotlin/com/slack/circuit/foundation/NavigableCircuitContent.kt

@@ -39,8 +39,10 @@ import com.slack.circuit.backstack.providedValuesForBackStack
 import com.slack.circuit.retained.CanRetainChecker
 import com.slack.circuit.retained.LocalCanRetainChecker
 import com.slack.circuit.retained.LocalRetainedStateRegistry
+import com.slack.circuit.retained.RetainedStateHolder
 import com.slack.circuit.retained.RetainedStateRegistry
 import com.slack.circuit.retained.rememberRetained
+import com.slack.circuit.retained.rememberRetainedStateHolder
 import com.slack.circuit.runtime.InternalCircuitApi
 import com.slack.circuit.runtime.Navigator
 import com.slack.circuit.runtime.screen.Screen
@@ -103,21 +105,23 @@ public fun <R : Record> NavigableCircuitContent(
    */
   val outerKey = "_navigable_registry_${currentCompositeKeyHash.toString(MaxSupportedRadix)}"
   val outerRegistry = rememberRetained(key = outerKey) { RetainedStateRegistry() }
-
   val saveableStateHolder = rememberSaveableStateHolder()
 
   CompositionLocalProvider(LocalRetainedStateRegistry provides outerRegistry) {
-    decoration.DecoratedContent(activeContentProviders, backStack.size, modifier) { provider ->
-      val record = provider.record
-
-      saveableStateHolder.SaveableStateProvider(record.key) {
-        // Remember the `providedValues` lookup because this composition can live longer than
-        // the record is present in the backstack, if the decoration is animated for example.
-        val values = remember(record) { providedValues[record] }?.provideValues()
-        val providedLocals = remember(values) { values?.toTypedArray() ?: emptyArray() }
-
-        CompositionLocalProvider(LocalBackStack provides backStack, *providedLocals) {
-          provider.content(record)
+    val retainedStateHolder = rememberRetainedStateHolder()
+    CompositionLocalProvider(LocalRetainedStateHolder provides retainedStateHolder) {
+      decoration.DecoratedContent(activeContentProviders, backStack.size, modifier) { provider ->
+        val record = provider.record
+
+        saveableStateHolder.SaveableStateProvider(record.key) {
+          // Remember the `providedValues` lookup because this composition can live longer than
+          // the record is present in the backstack, if the decoration is animated for example.
+          val values = remember(record) { providedValues[record] }?.provideValues()
+          val providedLocals = remember(values) { values?.toTypedArray() ?: emptyArray() }
+
+          CompositionLocalProvider(LocalBackStack provides backStack, *providedLocals) {
+            provider.content(record)
+          }
         }
       }
     }
@@ -174,26 +178,22 @@ private fun <R : Record> buildCircuitContentProviders(
 
       val lifecycle =
         remember { MutableRecordLifecycle() }.apply { isActive = lastBackStack.topRecord == record }
+      val retainedStateHolder = LocalRetainedStateHolder.current
 
       CompositionLocalProvider(LocalCanRetainChecker provides recordInBackStackRetainChecker) {
         // Now provide a new registry to the content for it to store any retained state in,
         // along with a retain checker which is always true (as upstream registries will
         // maintain the lifetime), and the other provided values
-        val recordRetainedStateRegistry =
-          rememberRetained(key = record.registryKey) { RetainedStateRegistry() }
-
-        CompositionLocalProvider(
-          LocalRetainedStateRegistry provides recordRetainedStateRegistry,
-          LocalCanRetainChecker provides CanRetainChecker.Always,
-          LocalRecordLifecycle provides lifecycle,
-        ) {
-          CircuitContent(
-            screen = record.screen,
-            navigator = lastNavigator,
-            circuit = lastCircuit,
-            unavailableContent = lastUnavailableRoute,
-            key = record.key,
-          )
+        retainedStateHolder.RetainedStateProvider(record.registryKey) {
+          CompositionLocalProvider(LocalRecordLifecycle provides lifecycle) {
+            CircuitContent(
+              screen = record.screen,
+              navigator = lastNavigator,
+              circuit = lastCircuit,
+              unavailableContent = lastUnavailableRoute,
+              key = record.key,
+            )
+          }
         }
       }
     }
@@ -366,3 +366,6 @@ public object NavigatorDefaults {
 public val LocalBackStack: ProvidableCompositionLocal<BackStack<out Record>?> = compositionLocalOf {
   null
 }
+
+private val LocalRetainedStateHolder =
+  compositionLocalOf<RetainedStateHolder> { error("No RetainedStateHolder provided") }

circuit-foundation/src/commonMain/kotlin/com/slack/circuit/foundation/PausableState.kt

@@ -7,10 +7,6 @@ package com.slack.circuit.foundation
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.Stable
 import androidx.compose.runtime.remember
-import com.slack.circuit.foundation.internal.withCompositionLocalProvider
-import com.slack.circuit.retained.LocalRetainedStateRegistry
-import com.slack.circuit.retained.RetainedStateRegistry
-import com.slack.circuit.retained.rememberRetained
 import com.slack.circuit.runtime.CircuitUiState
 import com.slack.circuit.runtime.presenter.Presenter
 
@@ -60,14 +56,13 @@ public fun <T> pausableState(
   val state = remember(key) { MutableRef<T>(null) }
 
   val saveableStateHolder = rememberSaveableStateHolderWithReturn()
+  val retainedStateHolder = rememberRetainedStateHolderWithReturn()
 
   return if (isActive || state.value == null) {
-    val retainedStateRegistry = rememberRetained(key = key) { RetainedStateRegistry() }
-    withCompositionLocalProvider(LocalRetainedStateRegistry provides retainedStateRegistry) {
-        saveableStateHolder.SaveableStateProvider(
-          key = key ?: "pausable_state",
-          content = produceState,
-        )
+    val finalKey = key ?: "pausable_state"
+    saveableStateHolder
+      .SaveableStateProvider(finalKey) {
+        retainedStateHolder.RetainedStateProvider(key = finalKey, content = produceState)
       }
       .also {
         // Store the last emitted state

circuit-foundation/src/commonMain/kotlin/com/slack/circuit/foundation/RetainedStateHolder.kt

@@ -0,0 +1,77 @@
+// Copyright (C) 2024 Slack Technologies, LLC
+// SPDX-License-Identifier: Apache-2.0
+package com.slack.circuit.foundation
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.DisposableEffect
+import com.slack.circuit.foundation.internal.withCompositionLocalProvider
+import com.slack.circuit.retained.CanRetainChecker
+import com.slack.circuit.retained.LocalCanRetainChecker
+import com.slack.circuit.retained.LocalRetainedStateRegistry
+import com.slack.circuit.retained.RetainedStateRegistry
+import com.slack.circuit.retained.RetainedValueProvider
+import com.slack.circuit.retained.rememberRetained
+
+/** Copy of [RetainedStateHolder] to return content value */
+internal interface RetainedStateHolder {
+
+  @Composable fun <T> RetainedStateProvider(key: String, content: @Composable () -> T): T
+}
+
+/** Creates and remembers the instance of [RetainedStateHolder]. */
+@Composable
+internal fun rememberRetainedStateHolderWithReturn(): RetainedStateHolder {
+  return rememberRetained { RetainedStateHolderImpl() }
+}
+
+private class RetainedStateHolderImpl : RetainedStateHolder, RetainedStateRegistry {
+
+  private val registry = RetainedStateRegistry()
+
+  @Composable
+  override fun <T> RetainedStateProvider(key: String, content: @Composable (() -> T)): T {
+    return withCompositionLocalProvider(LocalRetainedStateRegistry provides this) {
+      val canRetainChecker = LocalCanRetainChecker.current ?: CanRetainChecker.Always
+      val childRegistry = rememberRetained(key = key) { RetainedStateRegistry() }
+      withCompositionLocalProvider(
+          LocalRetainedStateRegistry provides childRegistry,
+          LocalCanRetainChecker provides CanRetainChecker.Always,
+        ) {
+          content()
+        }
+        .also {
+          DisposableEffect(key, childRegistry) {
+            onDispose {
+              childRegistry.saveAll()
+              if (canRetainChecker.canRetain(this@RetainedStateHolderImpl)) {
+                saveValue(key)
+              }
+            }
+          }
+        }
+    }
+  }
+
+  override fun consumeValue(key: String): Any? {
+    return registry.consumeValue(key)
+  }
+
+  override fun registerValue(
+    key: String,
+    valueProvider: RetainedValueProvider,
+  ): RetainedStateRegistry.Entry {
+    return registry.registerValue(key, valueProvider)
+  }
+
+  override fun saveAll() {
+    registry.saveAll()
+  }
+
+  override fun saveValue(key: String) {
+    registry.saveValue(key)
+  }
+
+  override fun forgetUnclaimedValues() {
+    registry.forgetUnclaimedValues()
+  }
+}