Change rememberRetain not to retain the value of removed node

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

@@ -39,6 +39,7 @@ 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.RetainedStateProvider
 import com.slack.circuit.retained.RetainedStateRegistry
 import com.slack.circuit.retained.rememberRetained
 import com.slack.circuit.runtime.InternalCircuitApi
@@ -179,21 +180,16 @@ private fun <R : Record> buildCircuitContentProviders(
         // 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,
-          )
+        RetainedStateProvider(record.registryKey) {
+          CompositionLocalProvider(LocalRecordLifecycle provides lifecycle) {
+            CircuitContent(
+              screen = record.screen,
+              navigator = lastNavigator,
+              circuit = lastCircuit,
+              unavailableContent = lastUnavailableRoute,
+              key = record.key,
+            )
+          }
         }
       }
     }

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

@@ -7,10 +7,7 @@ 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.foundation.internal.withRetainedStateProvider
 import com.slack.circuit.runtime.CircuitUiState
 import com.slack.circuit.runtime.presenter.Presenter
 
@@ -62,12 +59,9 @@ public fun <T> pausableState(
   val saveableStateHolder = rememberSaveableStateHolderWithReturn()
 
   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"
+    withRetainedStateProvider(finalKey) {
+        saveableStateHolder.SaveableStateProvider(key = finalKey, content = produceState)
       }
       .also {
         // Store the last emitted state

circuit-foundation/src/commonMain/kotlin/com/slack/circuit/foundation/internal/WithRetainedStateProvider.kt

@@ -0,0 +1,36 @@
+// Copyright (C) 2024 Slack Technologies, LLC
+// SPDX-License-Identifier: Apache-2.0
+package com.slack.circuit.foundation.internal
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.DisposableEffect
+import com.slack.circuit.retained.CanRetainChecker
+import com.slack.circuit.retained.LocalCanRetainChecker
+import com.slack.circuit.retained.LocalRetainedStateRegistry
+import com.slack.circuit.retained.RetainedStateProvider
+import com.slack.circuit.retained.RetainedStateRegistry
+import com.slack.circuit.retained.rememberRetained
+
+/** Copy of [RetainedStateProvider] to return content value */
+@Composable
+internal fun <T> withRetainedStateProvider(key: String, content: @Composable () -> T): T {
+  val canRetainChecker = LocalCanRetainChecker.current ?: CanRetainChecker.Always
+  val parentRegistry = LocalRetainedStateRegistry.current
+  val registry = rememberRetained(key = key) { RetainedStateRegistry() }
+  return withCompositionLocalProvider(
+      LocalRetainedStateRegistry provides registry,
+      LocalCanRetainChecker provides CanRetainChecker.Always,
+    ) {
+      content()
+    }
+    .also {
+      DisposableEffect(key, registry) {
+        onDispose {
+          registry.saveAll()
+          if (canRetainChecker.canRetain(registry)) {
+            parentRegistry.saveValue(key)
+          }
+        }
+      }
+    }
+}