Change rememberRetain not to retain the value of removed node

CHANGELOG.md

@@ -3,6 +3,40 @@ Changelog
 
 **Unreleased**
 --------------
+### Behaviour Change: rememberRetained
+Previously, `rememberRetained` could sometimes restore values when a composable was re-added, depending on whether its parent `RetainedStateRegistry` had been saved (#1783).
+Now, `rememberRetained` aligns with `remember` and `rememberSaveable`: if a composable is removed and later re-added, its value will not be restored unless it is explicitly saved and then restored via the registry.
+
+### Behaviour Change: RetainedStateRegistry
+- `saveAll` now returns the saved values.
+- `RetainedStateRegistry.Entry.unregister` now returns whether the unsaved valueProvider was actually removed.
+- `saveAll` and `saveValue` now skip storing child values when `CanRetainChecker` returns `false`.
+
+### New: RetainedStateHolder
+Similar to `SaveableStateHolder`, `RetainedStateHolder` provides a mechanism to maintain separate `RetainedStateRegistry` entries for specific keys. This allows saving the state defined with `rememberRetained` for a subtree before it is disposed, so that the subtree can later be recomposed with its state restored.
+
+```kotlin
+val retainedStateHolder = rememberRetainedStateHolder()
+var currentTab by remember { mutableStateOf(TabA) }
+
+retainedStateHolder.RetainedStateProvider(key = currentTab.name) {
+  // rememberRetained values in tab content are preserved across tab switches
+  when(currentTab) {
+    TabA -> {
+      TabAContent()
+    }
+    TabB -> {
+      TabBContent()
+    }
+    TabC -> {
+      TabCContent()
+    }
+  }
+}
+```
+
+### Internal Change: NavigableCircuitContent
+The approach of managing a separate `RetainedStateRegistry` for each record has been changed to use `RetainedStateHolder` instead.
 
 0.25.0
 ------

build.gradle.kts

@@ -107,6 +107,8 @@ allprojects {
         "**/Remove.kt",
         "**/Pets.kt",
         "**/SystemUiController.kt",
+        "**/RetainedStateHolderTest.kt",
+        "**/RetainedStateRestorationTester.kt",
       )
     }
   }

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

@@ -49,6 +49,7 @@ import com.slack.circuit.retained.LocalCanRetainChecker
 import com.slack.circuit.retained.LocalRetainedStateRegistry
 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
@@ -67,14 +68,6 @@ public fun <R : Record> NavigableCircuitContent(
   unavailableRoute: (@Composable (screen: Screen, modifier: Modifier) -> Unit) =
     circuit.onUnavailableContent,
 ) {
-  val activeContentProviders =
-    buildCircuitContentProviders(
-      backStack = backStack,
-      navigator = navigator,
-      circuit = circuit,
-      unavailableRoute = unavailableRoute,
-    )
-
   if (backStack.isEmpty) return
 
   /*
@@ -113,6 +106,14 @@ public fun <R : Record> NavigableCircuitContent(
   val outerRegistry = rememberRetained(key = outerKey) { RetainedStateRegistry() }
 
   CompositionLocalProvider(LocalRetainedStateRegistry provides outerRegistry) {
+    val activeContentProviders =
+      buildCircuitContentProviders(
+        backStack = backStack,
+        navigator = navigator,
+        circuit = circuit,
+        unavailableRoute = unavailableRoute,
+      )
+
     decoration.DecoratedContent(activeContentProviders, backStack.size, modifier) { provider ->
       val record = provider.record
 
@@ -174,6 +175,7 @@ private fun <R : Record> buildCircuitContentProviders(
   val lastUnavailableRoute by rememberUpdatedState(unavailableRoute)
 
   val saveableStateHolder = rememberSaveableStateHolder()
+  val retainedStateHolder = rememberRetainedStateHolder()
 
   fun createRecordContent() =
     movableContentOf<R> { record ->
@@ -186,32 +188,27 @@ private fun <R : Record> buildCircuitContentProviders(
         remember { MutableRecordLifecycle() }.apply { isActive = lastBackStack.topRecord == record }
 
       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() }
         saveableStateHolder.SaveableStateProvider(record.registryKey) {
-          CompositionLocalProvider(
-            LocalRetainedStateRegistry provides recordRetainedStateRegistry,
-            LocalCanRetainChecker provides CanRetainChecker.Always,
-            LocalRecordLifecycle provides lifecycle,
-          ) {
-            CircuitContent(
-              screen = record.screen,
-              navigator = lastNavigator,
-              circuit = lastCircuit,
-              unavailableContent = lastUnavailableRoute,
-              key = record.key,
-            )
+          // Provides a RetainedStateRegistry that is maintained independently for each record while
+          // the record exists in the back stack.
+          retainedStateHolder.RetainedStateProvider(record.registryKey) {
+            CompositionLocalProvider(LocalRecordLifecycle provides lifecycle) {
+              CircuitContent(
+                screen = record.screen,
+                navigator = lastNavigator,
+                circuit = lastCircuit,
+                unavailableContent = lastUnavailableRoute,
+                key = record.key,
+              )
+            }
           }
         }
-        // Remove saved states for records that are no longer in the back stack
-        DisposableEffect(record.registryKey) {
-          onDispose {
-            if (!lastBackStack.containsRecord(record, includeSaved = true)) {
-              saveableStateHolder.removeState(record.registryKey)
-            }
+      }
+      // Remove saved states for records that are no longer in the back stack
+      DisposableEffect(record.registryKey) {
+        onDispose {
+          if (!lastBackStack.containsRecord(record, includeSaved = true)) {
+            saveableStateHolder.removeState(record.registryKey)
           }
         }
       }

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,96 @@
+// 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 androidx.compose.runtime.key
+import androidx.compose.runtime.remember
+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
+
+  fun removeState(key: String)
+}
+
+/** Creates and remembers the instance of [RetainedStateHolder]. */
+@Composable
+internal fun rememberRetainedStateHolderWithReturn(): RetainedStateHolder {
+  return rememberRetained { RetainedStateHolderImpl() }
+}
+
+private class RetainedStateHolderImpl : RetainedStateHolder, RetainedStateRegistry {
+
+  private val registry = RetainedStateRegistry()
+
+  private val entries = mutableMapOf<String, Entry>()
+
+  @Composable
+  override fun <T> RetainedStateProvider(key: String, content: @Composable (() -> T)): T {
+    return withCompositionLocalProvider(LocalRetainedStateRegistry provides registry) {
+      key(key) {
+        val entry = remember { Entry() }
+        val childRegistry = rememberRetained(key = key) { RetainedStateRegistry() }
+        withCompositionLocalProvider(
+            LocalRetainedStateRegistry provides childRegistry,
+            LocalCanRetainChecker provides CanRetainChecker.Always,
+            content = content,
+          )
+          .also {
+            DisposableEffect(Unit) {
+              entries[key] = entry
+              onDispose {
+                if (entry.shouldSave) {
+                  registry.saveValue(key)
+                }
+                entries -= key
+              }
+            }
+          }
+      }
+    }
+  }
+
+  override fun removeState(key: String) {
+    val entry = entries[key]
+    if (entry != null) {
+      entry.shouldSave = false
+    } else {
+      registry.consumeValue(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(): Map<String, List<Any?>> {
+    return registry.saveAll()
+  }
+
+  override fun saveValue(key: String) {
+    registry.saveValue(key)
+  }
+
+  override fun forgetUnclaimedValues() {
+    registry.forgetUnclaimedValues()
+  }
+
+  private data class Entry(var shouldSave: Boolean = true)
+}