Jetpack Compose: Provide initial value for TextField

I want to achieve the following use case: A payment flow where you start with a screen to enter the amount (AmountScreen) to pay and some other screens to enter other values for the payment. At the end of the flow, a summary screen (SummaryScreen) is shown where you can modify the values inline. For the sake of simplicity we will assume there is only AmountScreen followed by SummaryScreen.

Now the following requirements should be realized:

on AmountScreen you don't loose your input on configuration change
when changing a value in SummaryScreen and go back to AmountScreen (using system back), the input is set to the changed value
AmountScreen and SummaryScreen must not know about the viewModel of the payment flow (PaymentFlowViewModel, see below)

So the general problem is: we have a screen with an initial value for an input field. The initial value can be changed on another (later) screen and when navigating back to the first screen, the initial value should be set to the changed value.

I tried various approaches to achieve this without reverting to Kotlin flows (or LiveData). Is there an approach without flows to achieve this (I am quite new to compose so I might be overlooking something obvious). If flows is the correct approach, would I keep a MutableStateFlow inside the PaymentFlowViewModel for amount instead of a simple string?

Here is the approach I tried (stripped and simplified from the real world example).

General setup:

internal class PaymentFlowViewModel : ViewModel() {
    var amount: String = ""
}

@Composable
internal fun NavigationGraph(viewModel: PaymentFlowViewModel = viewModel()) {
    val navController = rememberNavController()
    NavHost(
        navController = navController,
        startDestination = "AMOUNT_INPUT_SCREEN"
    ) {
        composable("AMOUNT_INPUT_SCREEN") {
            AmountInputRoute(
                // called when the Continue button is clicked
                onAmountConfirmed = {
                    viewModel.amount = it
                    navController.navigate("SUMMARY_SCREEN")
                },
                // apply the entered amount as the initial value for the input text
                initialAmount = viewModel.amount
            )
        }
        composable("SUMMARY_SCREEN") {
            SummaryRoute(
                // called when the amount is changed inline
                onAmountChanged = {
                    viewModel.amount = it
                },
                // apply the entered amount as the initial value for the input text
                amount = viewModel.amount
            )
        }
    }
}

The classes of the AmountScreen look like this:

@Composable
internal fun AmountInputRoute(
    initialAmount: String,
    onAmountConfirmed: (String) -> Unit
) {
    // without the "LaunchedEffect" statement below this fulfils all requirements
    // except that the changed value from the SummaryScreen is not applied
    val amountInputState: MutableState<String> = rememberSaveable { mutableStateOf(initialAmount) }
    // inserting this fulfils the req. that the changed value from SummaryScreen is
    // applied, but breaks keeping the entered value on configuration change
    LaunchedEffect(Unit) {
        amountInputState.value = initialAmount
    }
    Column {
         AmountInputView(
            amountInput = amountInputState.value,
            onAmountChange = { amountInput ->
                amountInputState.value = amountInput
            }
        )

        Button(onClick = { onAmountConfirmed(amountInputState.value) }) {
            Text(text = "Continue")
        }
    }
}
```

Solution

The natural approach for this scenario would be to use rememberSaveable(input = inititalAmount) {mutableStateOf(initialAmount)}. Unfortunately, when the state gets restored from a saveable (e.g. when navigating back), the change of the initial value is not considered because restoring takes precedence over re-initializing. Even though this behaviour is not the behaviour one would expect (at least from my point of view) it is not clear, if Google will fix this (https://issuetracker.google.com/issues/152014032).

Fortunately Google did provide me with a workaround for this particular scenario in a comment to the issue above. You can wrap rememberSaveable inside a key(initialAmount) composable. As Google explained: "The key composable will change the currentCompositeKeyHash we use inside of rememberSaveable, meaning that when initialValue changes the key changes and the value is not restored."

Using this approach our scenario can be implemented like this:

/**
 * Allows to create a saveable with an initial value where updating the initial value will lead to updating the
 * state *even if* the composable gets restored from saveable later on.
 *
 * rememberSaveableWithInitialValue should be used in a composable, when you want to initialize a mutable state
 * (e.g. for holding the value of a textinput field) with an initial value AND need the user input to survive
 * configuration changes AND want to allow changes to the initial value while being on a later screen
 * (i.e. while this composable is not active).
 */
@Composable
public fun <T : Any> rememberSaveableWithVolatileInitialValue(
    initialValue: T
): MutableState<T> {
    return key(initialValue) {
        rememberSaveable {
            mutableStateOf(initialValue)
        }
    }
}

This will be used inside a Screen-Composable like that:

@Composable
internal fun AmountInputRoute(
    initialAmount: String,
    onAmountConfirmed: (String) -> Unit
) {
    val amountInput by rememberSaveableWithVolatileInitialValue(initialAmount)
    Column {
         AmountInputView(
            amountInput = amountInput,
            onAmountChange = { changedInput ->
                amountInput = changedInput
            }
        )

        Button(onClick = {              
                onAmountConfirmed(amountInput)
            }) {
            Text(text = "Continue")
        }
    }
}