Enable RememberObserver to work with rememberRetained (#1210)

This PR manually calls `RememberObserver` calls as appropriate when used
in `rememberRetained`. The semantics of `rememberRetained` are obviously
different to `remember`, so we call them at different times:

- `onRemembered` will be called the first time that the retained item is
first `remember`ed. The difference here is that we will **not** call
`onRemembered` again once a the object is restored and remembered back
into composition.
- `onForgotten` will be called only once the retained state is cleared
from both composition and any retained registries.
- Thus, the logical lifecycle is maintained of created (`onRemembered`
called), through to being no longer retained/remembered (`onForgotten`
called).
- We do not call `onAbandoned` as retained state doesn't meet the
contract of the function.

I did think about creating a `RetainedObserver` interface, but figured
that using the existing `RememberObserver`, even with slightly different
semantics, is a better solution.

This has a variety of use-cases, but I have been playing around with a
`rememberRetainedCorotineScope` in Tivi:
chrisbanes/tivi#1763

circuit-retained/src/androidInstrumentedTest/kotlin/com/slack/circuit/retained/android/RetainedTest.kt

@@ -13,6 +13,7 @@ import androidx.compose.material.Text
 import androidx.compose.material.TextField
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.runtime.RememberObserver
 import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
@@ -281,6 +282,54 @@ class RetainedTest {
     composeTestRule.onNodeWithTag(TAG_RETAINED_1).assertTextContains("")
   }
 
+  @Test
+  fun rememberObserver() {
+    val subject =
+      object : RememberObserver {
+        var onRememberCalled: Int = 0
+          private set
+
+        var onForgottenCalled: Int = 0
+          private set
+
+        override fun onAbandoned() = Unit
+
+        override fun onForgotten() {
+          onForgottenCalled++
+        }
+
+        override fun onRemembered() {
+          onRememberCalled++
+        }
+      }
+
+    val content =
+      @Composable {
+        rememberRetained { subject }
+        Unit
+      }
+    setActivityContent(content)
+
+    assertThat(subject.onRememberCalled).isEqualTo(1)
+    assertThat(subject.onForgottenCalled).isEqualTo(0)
+
+    // Restart the activity
+    scenario.recreate()
+    // Compose our content again
+    setActivityContent(content)
+
+    // Assert that onRemembered was not called again
+    assertThat(subject.onRememberCalled).isEqualTo(1)
+    assertThat(subject.onForgottenCalled).isEqualTo(0)
+
+    // Now finish the Activity
+    scenario.close()
+
+    // Assert that the observer was forgotten
+    assertThat(subject.onRememberCalled).isEqualTo(1)
+    assertThat(subject.onForgottenCalled).isEqualTo(1)
+  }
+
   private fun nestedRegistriesWithPopAndPush(useKeys: Boolean) {
     val content = @Composable { NestedRetainWithPushAndPop(useKeys = useKeys) }
     setActivityContent(content)

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

@@ -95,7 +95,14 @@ public fun <T : Any> rememberRetained(vararg inputs: Any?, key: String? = null,
       val restored = registry.consumeValue(finalKey) as? RetainableHolder.Value<*>
       val finalValue = restored?.value ?: init()
       val finalInputs = restored?.inputs ?: inputs
-      RetainableHolder(registry, canRetainChecker, finalKey, finalValue, finalInputs)
+      RetainableHolder(
+        registry = registry,
+        canRetainChecker = canRetainChecker,
+        key = finalKey,
+        value = finalValue,
+        inputs = finalInputs,
+        hasBeenRestored = restored != null,
+      )
     }
   val value = holder.getValueIfInputsAreEqual(inputs) ?: init()
   SideEffect { holder.update(registry, finalKey, value, inputs) }
@@ -111,6 +118,7 @@ private class RetainableHolder<T>(
   private var key: String,
   private var value: T,
   private var inputs: Array<out Any?>,
+  private var hasBeenRestored: Boolean = false,
 ) : RetainedValueProvider, RememberObserver {
   private var entry: RetainedStateRegistry.Entry? = null
 
@@ -124,6 +132,10 @@ private class RetainableHolder<T>(
       this.key = key
       entryIsOutdated = true
     }
+    if (this.value !== value) {
+      // If the value changes, clear the hasBeenRestored flag
+      hasBeenRestored = false
+    }
     this.value = value
     this.inputs = inputs
     if (entry != null && entryIsOutdated) {
@@ -149,18 +161,27 @@ private class RetainableHolder<T>(
     // If the value is a RetainedStateRegistry, we need to take care to retain it.
     // First we tell it to saveAll, to retain it's values. Then we need to tell the host
     // registry to retain the child registry.
-    if (value is RetainedStateRegistry) {
-      (value as RetainedStateRegistry).saveAll()
+    val v = value
+    if (v is RetainedStateRegistry) {
+      v.saveAll()
       registry?.saveValue(key)
     }
 
     if (registry != null && !canRetainChecker.canRetain(registry!!)) {
       entry?.unregister()
+      // If value is a RememberObserver, we notify that it has been forgotten
+      if (v is RememberObserver) v.onForgotten()
     }
   }
 
   override fun onRemembered() {
     register()
+
+    // If value is a RememberObserver, we notify that it has remembered
+    if (!hasBeenRestored) {
+      val v = value
+      if (v is RememberObserver) v.onRemembered()
+    }
   }
 
   override fun onForgotten() {

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

@@ -3,6 +3,7 @@
 package com.slack.circuit.retained
 
 import androidx.compose.runtime.ProvidableCompositionLocal
+import androidx.compose.runtime.RememberObserver
 import androidx.compose.runtime.staticCompositionLocalOf
 import com.slack.circuit.retained.RetainedStateRegistry.Entry
 
@@ -143,6 +144,8 @@ internal class RetainedStateRegistryImpl(retained: MutableMap<String, List<Any?>
   }
 
   override fun forgetUnclaimedValues() {
+    // Notify any RememberObservers that it has been forgotten
+    retained.asSequence().filterIsInstance<RememberObserver>().forEach { it.onForgotten() }
     retained.clear()
   }
 }