Fix value restoration after removeState call in RetainedStateHolder

[vulpeszerda]

I discovered an issue while running additional tests with RetainedStateHolder where the handling was incorrect.

@Test
fun removedStateShouldNotBeRestored() {
  var increment = 0
  val screen = Screens.Screen1
  var restorableStateHolder: RetainedStateHolder? = null
  var restorableNumberOnScreen1 = -1
  restorationTester.setContent {
    val holder = rememberRetainedStateHolder()
    restorableStateHolder = holder
    holder.RetainedStateProvider(screen.name) {
      restorableNumberOnScreen1 = rememberRetained { increment++ }
    }
  }

  composeTestRule.runOnIdle {
    assertThat(restorableNumberOnScreen1).isEqualTo(0)
    restorableNumberOnScreen1 = -1
    restorableStateHolder!!.removeState(screen.name)
  }

  restorationTester.emulateRetainedInstanceStateRestore()

  composeTestRule.runOnIdle { assertThat(restorableNumberOnScreen1).isEqualTo(1) }
}

In this test scenario, the RetainedStateHolder holds a child RetainedStateRegistry for Screen1. Since removeState is called, the corresponding Entry.shouldSave is set to false, preventing a call to saveValue in the subsequent DisposableEffect.

RetainedStateHolder.kt (before fix)

@Composable
override fun RetainedStateProvider(key: String, content: @Composable (() -> Unit)) {
  CompositionLocalProvider(LocalRetainedStateRegistry provides registry) {
    ReusableContent(key) {
      val entry = remember { Entry() }
      val childRegistry = rememberRetained(key = key) { RetainedStateRegistry() }
      CompositionLocalProvider(
        LocalRetainedStateRegistry provides childRegistry,
        LocalCanRetainChecker provides CanRetainChecker.Always,
        content = content,
      )
      DisposableEffect(Unit) {
        entries[key] = entry
        onDispose {
          if (entry.shouldSave) {
            registry.saveValue(key)
          }
          entries -= key
        }
      }
    }
  }
}

However, when the parent’s overall RetainedStateRegistry.saveAll() is invoked, the childRegistry—lacking a specifically set LocalCanRetainChecker—does not consider the Entry.shouldSave flag and ends up saving the value.

To fix this, I modified the code to provide a LocalCanRetainChecker for the childRegistry. Additionally, I have added few test logics for RetainedStateHolder.

I wish this had been handled better in #1794; sorry for the hassle, and please take a look when you have time!

[ZacSweers]

Thanks!