Replaced registry usage with RetainedStateProvider

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

@@ -19,7 +19,6 @@ import androidx.compose.animation.slideOutHorizontally
 import androidx.compose.animation.togetherWith
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
-import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.Immutable
 import androidx.compose.runtime.ProvidableCompositionLocal
 import androidx.compose.runtime.compositionLocalOf
@@ -40,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
@@ -175,35 +175,20 @@ private fun <R : Record> buildCircuitContentProviders(
 
       val lifecycle =
         remember { MutableRecordLifecycle() }.apply { isActive = lastBackStack.topRecord == record }
-      val parentRetainedStateRegistry = LocalRetainedStateRegistry.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 registryKey = record.registryKey
-        val recordRetainedStateRegistry =
-          rememberRetained(key = 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,
-          )
-        }
-
-        DisposableEffect(registryKey, recordRetainedStateRegistry) {
-          onDispose {
-            if (recordInBackStackRetainChecker.canRetain(recordRetainedStateRegistry)) {
-              parentRetainedStateRegistry.saveValue(registryKey)
-            }
+        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

@@ -5,13 +5,9 @@
 package com.slack.circuit.foundation
 
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.DisposableEffect
 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
 
@@ -61,19 +57,15 @@ public fun <T> pausableState(
   val state = remember(key) { MutableRef<T>(null) }
 
   val saveableStateHolder = rememberSaveableStateHolderWithReturn()
-  val retainedStateRegistry = rememberRetained(key = key) { RetainedStateRegistry() }
 
   return if (isActive || state.value == null) {
-    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
         state.value = it
-        DisposableEffect(retainedStateRegistry) { onDispose { retainedStateRegistry.saveAll() } }
       }
   } else {
     // Else, we just emit the last stored state instance

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)
+          }
+        }
+      }
+    }
+}

circuit-retained/src/commonMain/kotlin/com/slack/circuit/retained/RetainedStateProvider.kt

@@ -0,0 +1,38 @@
+// Copyright (C) 2024 Slack Technologies, LLC
+// SPDX-License-Identifier: Apache-2.0
+package com.slack.circuit.retained
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.runtime.DisposableEffect
+import kotlin.uuid.ExperimentalUuidApi
+import kotlin.uuid.Uuid
+
+/**
+ * Provides a [RetainedStateRegistry] for the child [content] based on the specified [key]. Before
+ * the provided registry is disposed, it calls [RetainedStateRegistry.saveValue] on the parent
+ * registry to save the current value, allowing it to be restored on the next visit with the same
+ * key.
+ */
+@Composable
+public fun <T> RetainedStateProvider(key: String? = null, content: @Composable () -> T) {
+  @OptIn(ExperimentalUuidApi::class)
+  val finalKey = key ?: rememberRetained { Uuid.random().toString() }
+  val canRetainChecker = LocalCanRetainChecker.current ?: CanRetainChecker.Always
+  val parentRegistry = LocalRetainedStateRegistry.current
+  val registry = rememberRetained(key = finalKey) { RetainedStateRegistry() }
+  CompositionLocalProvider(
+    LocalRetainedStateRegistry provides registry,
+    LocalCanRetainChecker provides CanRetainChecker.Always,
+  ) {
+    content()
+  }
+  DisposableEffect(finalKey, registry) {
+    onDispose {
+      registry.saveAll()
+      if (canRetainChecker.canRetain(registry)) {
+        parentRegistry.saveValue(finalKey)
+      }
+    }
+  }
+}

circuit-retained/src/commonMain/kotlin/com/slack/circuit/retained/RetainedStateRegistry.kt

@@ -110,6 +110,15 @@ internal class RetainedStateRegistryImpl(retained: MutableMap<String, List<Any?>
   }
 
   override fun saveAll() {
+    fun save(value: Any?) {
+      when (value) {
+        // If we get a RetainedHolder value, need to unwrap and call again
+        is RetainedValueHolder<*> -> save(value.value)
+        // Dispatch the call to nested registries
+        is RetainedStateRegistry -> value.saveAll()
+      }
+    }
+
     val values =
       valueProviders.mapValues { (_, list) ->
         // If we have multiple providers we should store null values as well to preserve
@@ -132,20 +141,11 @@ internal class RetainedStateRegistryImpl(retained: MutableMap<String, List<Any?>
   override fun saveValue(key: String) {
     val providers = valueProviders[key]
     if (providers != null) {
-      retained[key] = providers.map { it.invoke().also(::save) }
+      retained[key] = providers.map { it.invoke() }
       valueProviders.remove(key)
     }
   }
 
-  private fun save(value: Any?) {
-    when (value) {
-      // If we get a RetainedHolder value, need to unwrap and call again
-      is RetainedValueHolder<*> -> save(value.value)
-      // Dispatch the call to nested registries
-      is RetainedStateRegistry -> value.saveAll()
-    }
-  }
-
   override fun forgetUnclaimedValues() {
     fun clearValue(value: Any?) {
       when (value) {