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 function public and modified NavigableCircuitContent and PausableState to use it. (24d8547)

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, 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.

[vulpeszerda]

Timeout failure occurred when attempting to launch the emulator in the CI.
I’m not sure how to proceed from here. 😢

Not enough space to create userdata partition. Available: 6784.417969 MB at /home/runner/.android/avd/../avd/test.avd, need 7372.800000 MB.

It seems that have not enough space in CI machine.

[vulpeszerda]

@ZacSweers Can you help me about this?

[ZacSweers]

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