Change rememberRetain not to retain the value of removed node

[vulpeszerda]

Resolves #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)
  val lastCanRetainChecker by rememberUpdatedState(canRetainChecker)

  LifecycleStartEffect(vm) {
    onStopOrDispose {
      if (lastCanRetainChecker.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.

Fixes #1783

[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?

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

[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 👌

[vulpeszerda]

I’ve cherry-picked the RetainedStateHolderTest from #1168 and integrated it into this branch. The tests pass with the RetainedStateHolder implementation in this PR, and everything appears to be working fine.

To use the existing circuit’s RetainedStateRegistry without significantly modifying its implementation, I had to implement RetainedStateHolder differently than the one in previous PR.

From my understanding, the issue addressed in #1168 is resolved by using this RetainedStateHolder, so I went ahead and applied those changes in commit 9714ce3.

It looks like all the issues have been resolved. If there’s anything else you’d like me to work on, please let me know!

[ZacSweers]

Thanks a bunch! @alexvanyo is going to take one more look today as well then I think we can move ahead with this. Thanks a bunch for your patience in this. Josh is on vacation and we've had a bit of a chaotic month to start the new year

[ZacSweers]

nit: do you think we can/should reconcile these into one API?

[vulpeszerda]

I think it would be good to consolidate them, but there are a few considerations:

1. Just like SaveableStateHolder in Add withCompositionLocalProvider to avoid backward writes #1451 , we need to investigate further whether converting a RetainedStateHolder without a return value into one with a return value would cause any issues.

2. The RetainedStateHolder without a return value internally uses ReusableContent(key), while the one with a return value uses key(key). I’m not entirely sure if using ReusableContent(key) in the case of a return value would cause any issues.

[alexvanyo]

Thank you for the work here and investigating the different options! The resulting API and usage matches my expectations more, and is more consistent with the SaveableStateProvider setup.

[alexvanyo]

Nit: this comment is now a bit out of date, it could probably point to how RetainedStateProvider works

[vulpeszerda]

Um.. yes I changed a little bit.

[ZacSweers]

Sorry one last request - would you mind doing a writeup of this change into the CHANGELOG.md file? It can be a longer/more detailed section to explain it since it's more significant of a change :). Let's make sure we highlight both the behavior changes and the new APIs.

[alexvanyo]

Something very subtle I ran into experimenting with this change: I think this needs to have both entryCanRetainChecker and childRegistry as keys, or alterantively using rememberUpdatedState with them to avoid operating on stale references to old childRegistrys or entryCanRetainCheckers.

This is exacerbated by #1919 - if the CanRetainChecker instance changes, then this can get into a bad state where the child registry gets a new instance, and this disposable effect operates on the stale old instance.

[vulpeszerda]

Sorry for bringing up a difficult discussion again at the end of a long task, but to properly address the issue Alex raised, I think I need to understand how the Main Contributors think about the following concerns.

Should CanRetainChecker only affect RetainedStateRegistry?

CanRetainChecker originally seems to have existed for RetainedStateRegistry. However, since it is provided via CompositionLocal, the following code can occur:

val registry = remember { RetainedStateRegistry() }

CompositionLocalProvider(LocalRetainedStateRegistry provides registry) {

    CompositionLocalProvider(LocalCanRetainChecker provides remember { CanRetainChecker { false }}) {
          val value = rememberRetained { ... }
    }
}

In this case, when registry.saveAll() is called, the value should / should not be saved?

I believe that LocalCanRetainChecker should only affect to RetainedStateRegistry level, not for individual rememberRetained. (which means the value should be saved)

[ZacSweers]

I think right now we need to check in rememberRetained still. We could revisit being able to check it solely in RetainedStateRegistry but I think that's outside the scope of this PR maybe?

[vulpeszerda]

Yes right. I agree. For this update, I’ll make sure LocalCanRetainChecker continues to work as expected and still affects rememberRetained.

However, the changes I made earlier caused an issue where LocalCanRetainChecker no longer influenced rememberRetained, so I’ve fixed that.

[ZacSweers]

Sounds good. Don't forget the changelog writeup!

[vulpeszerda]

Of course! 🙂
I was actually writing a more detailed explanation of the changes I just uploaded.
But it’s taking a bit of time, so it’d be great if you could review it tomorrow or the day after!

[vulpeszerda]

Alex, since the code in RetainedStateHolder that referenced parent CanRetainChecker has been removed, we no longer need to handle the scenario of CanRetainChecker changes separately.

To elaborate, in your #1920 you addressed changes in CanRetainChecker within rememberRetained. In my revised approach, whether the child registry is saved in RetainedStateHolder is automatically determined by the CanRetainChecker applied to rememberRetained in below

rememberRetained(key = key) { RetainedStateRegistry() }

[alexvanyo]

I think I see - so the latest changes now fully rely upon CanRetainChecker being used in rememberRetained, instead of it also being checked inside RetainedStateHolder. Makes sense to me!

[vulpeszerda]

As I mentioned above,
I’ve added a test case for the scenario Alex originally commented on, where CanRetainChecker changes. During this process, I discovered and fixed unintended behavior change where LocalCanRetainChecker was not affecting rememberRetained.

Additionally, as part of this fix, I simplified the model by removing the unnecessary EntryCanRetainChecker. In the previous implementation of RetainedStateHolder, the parent scope’s LocalCanRetainChecker was already influencing the child registry’s canBeRetained, so now it only references the value for shouldSave.

As the final step of this long task, I’ll update the PR description to reflect the latest changes and add the changelog as requested by Zac
Done!

[alexvanyo]

I think I see - so the latest changes now fully rely upon CanRetainChecker being used in rememberRetained, instead of it also being checked inside RetainedStateHolder. Makes sense to me!

[ZacSweers]

Awesome work! Thanks again for your patience with us as well, really excited to get this finally in

[ZacSweers]

Great writeup!