Change rememberRetain not to retain the value of removed node

[vulpeszerda]

As mentioned in #1783, I propose modifying the behavior of rememberRetained to improve consistency.

Changes to rememberRetained Behavior

The following explanation is based on the code sample below.

@Composable
private fun ConditionalRetainContent(registry: RetainedStateRegistry) {
  CompositionLocalProvider(LocalRetainedStateRegistry provides registry) {
    var showContent by remember { mutableStateOf(false) }
    Column {
      Button(modifier = Modifier.testTag(TAG_BUTTON_HIDE), onClick = { showContent = false }) {
        Text(text = "Hide content")
      }
      Button(modifier = Modifier.testTag(TAG_BUTTON_SHOW), onClick = { showContent = true }) {
        Text(text = "Show content")
      }
      if (showContent) {
        var count by rememberRetained { mutableIntStateOf(0) }
        Button(modifier = Modifier.testTag(TAG_BUTTON_INC), onClick = { count += 1 }) {
          Text(text = "Increment")
        }
        Text(modifier = Modifier.testTag(TAG_RETAINED_1), text = count.toString())
      }
    }
  }
}

Current Behavior

Before saveAll is called on the registry

• rememberRetained added when the condition is true will be removed when the condition changes to false.
• When the condition changes back to true and the same rememberRetained is called, the previous value is not retained.

Test code for this behavior

After saveAll is called on the registry

• rememberRetained added when the condition is true will be removed when the condition changes to false.
• Then, saveAll is called on the registry.
• When the condition changes back to true and the same rememberRetained is called, the previous value is retained.

Test code for this behavior

Challenges with the Current Behavior

• The timing of the saveAll call on the RetainedStateRegistry is unknown from the perspective of lower-level content, making retention depend on whether saveAll was called.
• In the case of remember / rememberSaveable, the node doesn’t retain the previous value when it’s removed and added back, making it easy to expect rememberRetained to behave similarly.

Due to these challenges, I propose modifying rememberRetained to behave like remember / rememberSaveable, so that rememberRetained does not retain values when it is hidden and reshown in the compose node based on a condition.

Internal Implementation Changes

Overall, I drew inspiration from the implementation of SaveableStateHolder and rememberSaveable for these modifications.

The following changes were made to achieve this goal:

1. Modify RetainableSaveableHolder to always unregister values from the RetainedStateRegistry when onForgotten is called.

In the previous implementation, if rememberRetained was removed from the node and onForgotten was called, values were not unregistered from the registry if canRetain was true. As a result, all values were retained in the registry regardless of whether rememberRetained was present in the node when saveAll was called.

Now, values are unregistered from the registry when onForgotten is called, so only values present in the composition node are retained when saveAll is called .

2. Change the timing of saveAll in RetainedStateRegistry

Previously, saveAll was called on the RetainedStateRegistry within RetainableSaveableHolder when onForgotten was invoked. However, due to the change in 1, onForgotten of all child rememberRetained nodes is called before onForgotten of rememberRetained { RetainedStateRegistry() }, which would result in no values being retained.

Therefore, I referenced the SaveableStateHolder implementation to create a separate DisposableEffect to run saveAll, as shown below:

val parentRegistry = LocalRetainedStateRegistry.current
val registry = rememberRetained(key) { RetainedStateRegistry() }

CompositionLocalProvider(LocalRetainedStateRegistry provides registry) {
    // child content
    ... 
}

DisposableEffect(key, registry) {
    onDispose { 
        registry.saveAll()
        parentRegistry.saveValue(key)
    }
}

By structuring it this way, saveAll can be called on the RetainedStateRegistry before the child content's dispose stage. Thus, even if onForgotten is called on the child content's rememberRetained and unregisters the value, it will still be retained.

Compared to the existing implementation, this change needs to be applied to all cases where the RetainedStateRegistry is redefined in a nested way. To facilitate this process, I made the RetainedStateProvider RetainedStateHolder function public and modified NavigableCircuitContent and PausableState to use it. (24d8547 ff654b4 )

3. Change the timing of saveAll in AndroidContinuity

For the same reason as in 2, continuityRetainedStateRegistry also needs to call saveAll using DisposableEffect after declaring the child content.

Example

val registry = continuityRetainedStateRegistry()

CompositionLocalProvider(LocalRetainedStateRegistry provides registry) {
    // child content
    ... 
}

DisposableEffect(registry) {
    onDispose { 
        registry.saveAll()
    }
}

However, in Android, continuityRegistry should always call saveAll when onStop is called. Since calling saveAll multiple times won’t cause issues, I modified it to call saveAll conveniently upon onStop or disposal.

@Composable
public fun continuityRetainedStateRegistry(
  key: String = Continuity.KEY,
  factory: ViewModelProvider.Factory = ContinuityViewModel.Factory,
  canRetainChecker: CanRetainChecker = LocalCanRetainChecker.current ?: rememberCanRetainChecker(),
): RetainedStateRegistry {
  @Suppress("ComposeViewModelInjection")
  val vm = viewModel<ContinuityViewModel>(key = key, factory = factory)

  LifecycleStartEffect(vm) {
    onStopOrDispose {
      if (canRetainChecker.canRetain(vm)) {
        vm.saveAll()
      }
    }
  }

  LaunchedEffect(vm) {
    withFrameNanos {}
    // This resumes after the just-composed frame completes drawing. Any unclaimed values at this
    // point can be assumed to be no longer used
    vm.forgetUnclaimedValues()
  }

  return vm
}

Changes to Test Code

Change 1

In RetainedTest.kt, parts where nestedRegistry is declared now use RetainedStateProvider RetainedStateHolder, as saveAll must be manually called.
affec25

Change 2

In NestedRetainWithPushAndPop and NestedRetainWithPushAndPopAndCannotRetain, the tests assume the same value is retained regardless of the showNestedContent value, so I set the same key in RetainedStateProvider to retain the values.

However, since these assumptions may change with this PR, it may need verification to ensure correctness.
affec25

Change 3

To test ImpressionEffect, I modified the function attempting to recreate it. Previously, the condition was first set to false to remove the child content and then saveAll was performed. In the modified behavior, saveAll is performed first and then the child content is removed.

465f98f

Additional Test Cases for the Reported Issue

I added test cases to cover the issue reported, which show which tests failed with the previous implementation and how they succeed with the modified implementation.

019c8cc, 5b279d4

---

When I initially reported this issue, I may have been somewhat aggressive due to the unexpected behavior. I apologize if it came across that way. Through working on this modification, I had the chance to explore the internal structure of rememberRetained and have come to appreciate how robust and well-designed many parts of the circuit library are. I have great respect for the efforts of the main contributors.

After considering various approaches, I believe this change to the behavior of rememberRetained is the right direction. However, since the implementation of rememberRetained is a core part of circuit, there may be differing perspectives on this.

While having this PR approved would be ideal, if there are differing views, I hope we can use this as a foundation for further discussion.

[salesforce-cla]

Thanks for the contribution! Unfortunately we can't verify the commit author(s): roy.tk <r***@k***.com>. One possible solution is to add that email to your GitHub account. Alternatively you can change your commits to another email and force push the change. After getting your commits associated with your GitHub account, sign the Salesforce Inc. Contributor License Agreement and this Pull Request will be revalidated.

[ZacSweers]

Please be patient and don't tag maintainers, we will get to it when we get to it!

[alexvanyo]

I agree with the goal of having rememberRetained match in behavior with rememberSaveable and remember: if one of those leaves composition in an if, the state is wiped, and I would expect the same for rememberRetained.

One thing not clear right now to me still after running some tests: if there is the if (...) { rememberRetained() } inside the RetainedStateProvider, I think by the same logic, state should also be lost.

However, in my tests, the state is still preserved in that case.

Maybe the LocalCanRetainChecker inside RetainedStateProvider shouldn't be set to the static CanRetainChecker.Always?

[alexvanyo]

What's the difference in behavior with this change?

[vulpeszerda]

I intended to notify onForgotten for all retained values that implment RememberObserver.

[vulpeszerda]

@alexvanyo
Thank you for the review.

I agree with this

One thing not clear right now to me still after running some tests: if there is the if (...) { rememberRetained() } inside the RetainedStateProvider, I think by the same logic, state should also be lost.

However, I couldn't reproduce the test case you mentioned. (I tested with this)
Could you please share the exact test case?

One thing not clear right now to me still after running some tests: if there is the if (...) { rememberRetained() } inside the RetainedStateProvider, I think by the same logic, state should also be lost.

However, in my tests, the state is still preserved in that case.

[alexvanyo]

The RetainedStateHolder API looks very similar to #1168 as its now mimicking the SaveableStateHolder API exactly, so curious to hear from Zac as to how that exploration went previously.

[alexvanyo]

Comparing this to SaveableStateHolder, SaveableStateHolder has a removeState function to clear out the state for a given key.

Would this RetainedStateHolder need a similar utility, for clearing out state for entries no longer on the backstack?

[vulpeszerda]

Sure do! fb672cf

[alexvanyo]

With the updated setup, I'm not sure what this CanRetainChecker.Always does, as opposed to not overriding LocalCanRetainChecker at all.

I tried removing it, and all of the tests still pass, so maybe this can now be removed?

[vulpeszerda]

When creating separate RetainedStateRegistry instances within RetainedStateProvider, the CanRetainChecker propagated to the child has a different meaning from the parent’s CanRetainChecker.

Meaning of the child’s CanRetainChecker:

The child’s CanRetainChecker determines whether rememberRetained declared under the RetainedStateProvider can be retained in cases where the key transitions from A -> B -> A. This behavior should always function correctly even if the parent’s CanRetainChecker always returns false, because the RetainedStateHolder itself has never been disposed. As long as the RetainedStateHolder is not disposed, rememberRetained declared under the RetainedStateProvider should always be restorable, so it seems appropriate to pass CanRetainChecker.Always.

Meaning of the parent’s CanRetainChecker:

The CanRetainChecker provided by the parent of the RetainedStateHolder determines whether the internal data of the RetainedStateHolder should be retained when the screen is recreated.

---

For these reasons, I think the existing code also passed CanRetainChecker.Always when creating a new registry. Following the intention of the original code, I believe redefining it as CanRetainChecker.Always in my code is necessary.

I tried removing it, and all of the tests still pass, so maybe this can now be removed?

I think the current test cases might be insufficient. I expect a difference in behavior for the following scenario:

1. The parent’s CanRetainChecker always returns false.
2. A RetainedStateHolder is defined.
3. The key in holder.RetainedStateProvider transitions from A -> B -> A.

When redefined with CanRetainChecker.Always ->
rememberRetained inside holder.RetainedStateProvider is preserved.

Without redefined ->
rememberRetained inside holder.RetainedStateProvider is not preserved.

[ZacSweers]

The RetainedStateHolder API looks very similar to #1168 as its now mimicking the SaveableStateHolder API exactly, so curious to hear from Zac as to how that exploration went previously.

Still want it! It essentially stalled because I hadn't picked it back up again. If that's something that naturally can fall out of this work then that'd be dope 👌