How to save UI state within a chip while using navigation in compose?

I have a chip that changes configurations when being clicked on. However, as the user navigates back (popbackstack) or navigates to another screen, the clicked chip re-sets its previous position. As I understood, I will have to use a viewmodel to pass the chips state and in that way save it? I cant find a way though to store the "remembersavable" in a viewmodel.

How can I achieve this? Appreciate any feedback!

My example chip:

@Composable
fun CatsChip() {

    val textChipRememberOneState = rememberSaveable { mutableStateOf(false) }

    TextChip(
        isSelected = textChipRememberOneState.value,
        shape = Shapes(medium = RoundedCornerShape(15.dp)),
        text = "Cats",
        selectedColor = LightGreen,
        onChecked = {
            textChipRememberOneState.value = it
        },
    )
}

You can keep the state in a MutableStateFlow in a ViewModel.

To use ViewModels in Compose you need to add the following dependency to your app/build.gradle file

implementation("androidx.lifecycle:lifecycle-viewmodel-compose:2.6.1")

Now you can use the viewModel() function to get the ViewModel instance in your composables.

With ViewModels you do not need to use rememberSaveable anymore, since the state will be kept in
the ViewModel, however, if you want the state to persist even across process death (not just configuration changes), then you have to save the state in the SavedStateHandle.

Here is an example of a ViewModel that only keeps the state in memory, but does not save it in the SavedStateHandle.

class MemoryOnlyViewModel : ViewModel () {
    val checkedState = MutableStateFlow(false)

    fun onCheckedChange(isChecked: Boolean) = checkedState.update { isChecked }
}

Here is an example of a ViewModel that saves the state in the SavedStateHandle.

class SavedStateViewModel(private val state: SavedStateHandle) : ViewModel() {
    val checkedState = state.getStateFlow(key = CHECKED_STATE_KEY, initialValue = false)

    fun onCheckedChange(isSelected: Boolean) = state.set(key = CHECKED_STATE_KEY, value = isSelected)

    companion object {
        private const val CHECKED_STATE_KEY = "checkedState"
    }
}

Then the usage would look like this

import androidx.compose.runtime.getValue
import androidx.compose.runtime.setValue
        
@Composable
fun CatsChip() {
    val vm: SavedStateViewModel = viewModel() // or: val vm = viewModel<SavedStateViewModel>()
    val catsChipState by vm.checkedState.collectAsState()

    TextChip(
        isSelected = catsChipState,
        shape = Shapes(medium = RoundedCornerShape(15.dp)),
        text = "Cats",
        selectedColor = LightGreen,
        onChecked = vm::onCheckedChange, // or: onChecked = { vm.onCheckedChange(it) }
    )
}

For more use cases see also the Business Logic section of the Compose State Hoisting documentation

<hr>

Here is a demo Composable using Compose navigation and showcasing the two view models from above,
comparing it with rememberSaveable, scoping them in two different ways, to the parent context and to
the NavBackStackEntry. This shows how different scopes affect the lifecycle of ViewModels.

Requires the Compose navigation dependency in your app/build.gradle file

implementation("androidx.navigation:navigation-compose:2.5.3")

You can check the demo by calling Demo() in some composable content of your app.
Click the buttons to navigate and see how the backstack changes. The ViewModels and also the rememberSaveable that are scoped to the parent context will preserve the state all the time, whereas those that are scoped to each NavBackStackEntry will preserve state only for their own navigation destinations, which can be seen when navigating back.
Also the state saved in the MemoryOnlyViewModels will not survive process death, which you can check in the following way:

1. Send the application to the background by pressing the Home button (but do not close it in the app switcher)
2. Kill the process by running the following command in the IDE Terminal tab with the package name of your application

adb shell am kill <package_name>

3. Open the application again from the app switcher

If you followed the steps correctly and managed to kill and restore the process in this way, then you should notice that only MemoryOnlyViewModels have lost/reset their state.

Here is the whole demo code. Just copy and paste to a new Kotlin file and call the Demo() composable from a composable content.

import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.ColumnScope
import androidx.compose.foundation.layout.Row
import androidx.compose.foundation.layout.padding
import androidx.compose.foundation.shape.RoundedCornerShape
import androidx.compose.material.icons.Icons
import androidx.compose.material.icons.filled.AddCircle
import androidx.compose.material.icons.filled.Clear
import androidx.compose.material3.Button
import androidx.compose.material3.Icon
import androidx.compose.material3.Surface
import androidx.compose.material3.Text
import androidx.compose.runtime.Composable
import androidx.compose.runtime.collectAsState
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.saveable.rememberSaveable
import androidx.compose.runtime.setValue
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.unit.dp
import androidx.lifecycle.SavedStateHandle
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewmodel.compose.viewModel
import androidx.navigation.compose.NavHost
import androidx.navigation.compose.composable
import androidx.navigation.compose.currentBackStackEntryAsState
import androidx.navigation.compose.rememberNavController
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.update

class MemoryOnlyViewModel : ViewModel () {
    val checkedState = MutableStateFlow(false)

    fun onCheckedChange(isChecked: Boolean) = checkedState.update { isChecked }
}

class SavedStateViewModel(private val state: SavedStateHandle) : ViewModel() {
    val checkedState = state.getStateFlow(key = CHECKED_STATE_KEY, initialValue = false)

    fun onCheckedChange(isSelected: Boolean) = state.set(key = CHECKED_STATE_KEY, value = isSelected)

    companion object {
        private const val CHECKED_STATE_KEY = "checkedState"
    }
}
    
@Composable
fun Demo() {    
    @Composable
    fun SimpleChip(
        text: String,
        isSelected: Boolean,
        onChecked: (Boolean) -> Unit,
    ) {
        Surface(
            onClick = { onChecked(!isSelected) },
            modifier = Modifier.padding(4.dp),
            shape = RoundedCornerShape(16.dp),
            color = if (isSelected) Color(0xFF7986CB) else Color.LightGray,
        ) {
            Row(
                modifier = Modifier.padding(8.dp),
                horizontalArrangement = Arrangement.spacedBy(4.dp),
                verticalAlignment = Alignment.CenterVertically,
            ) {
                Text(text)
                Icon(
                    imageVector = if (isSelected) Icons.Default.Clear else Icons.Default.AddCircle,
                    contentDescription = null,
                )
            }
        }
    }

    // These VMs are scoped to the lifecycle of the parent context (likely a ComponentActivity)
    val parentMemoryOnlyVm: MemoryOnlyViewModel = viewModel()
    val parentSavedStateVm: SavedStateViewModel = viewModel()
    var parentSaveable by rememberSaveable { mutableStateOf(false) }

    @Composable
    @Suppress("UnusedReceiverParameter")
    fun ColumnScope.DemoScreen(text: String) {
        Text(text)

        val parentMemoryOnlyState by parentMemoryOnlyVm.checkedState.collectAsState()

        SimpleChip(text = "MemoryOnly VM (parent scoped)",
            isSelected = parentMemoryOnlyState, onChecked = parentMemoryOnlyVm::onCheckedChange)

        val navMemoryOnlyVm: MemoryOnlyViewModel = viewModel()
        val navMemoryOnlyState by navMemoryOnlyVm.checkedState.collectAsState()

        SimpleChip(text = "MemoryOnly VM (nav scoped)",
            isSelected = navMemoryOnlyState, onChecked = navMemoryOnlyVm::onCheckedChange)


        val parentSavedState by parentSavedStateVm.checkedState.collectAsState()

        SimpleChip(text = "SavedState VM (parent scoped)",
            isSelected = parentSavedState, onChecked = parentSavedStateVm::onCheckedChange)

        val navSavedStateVm: SavedStateViewModel = viewModel()
        val navSavedState by navSavedStateVm.checkedState.collectAsState()

        SimpleChip(text = "SavedState VM (nav scoped)",
            isSelected = navSavedState, onChecked = navSavedStateVm::onCheckedChange)


        SimpleChip(text = "rememberSaveable (parent scoped)",
            isSelected = parentSaveable, onChecked = { parentSaveable = it })

        var navSaveable by rememberSaveable { mutableStateOf(false) }

        SimpleChip(text = "rememberSaveable (nav scoped)",
            isSelected = navSaveable, onChecked = { navSaveable = it })
    }

    val navController = rememberNavController()

    @Composable
    fun BackButton() = Button(onClick = { navController.navigateUp() }) {
        Text("Go back")
    }

    @Composable
    fun NavButton(route: String) = Button(onClick = { navController.navigate(route) }) {
        Text("Navigate to $route")
    }

    Column {
        val currentEntry by navController.currentBackStackEntryAsState()
        val backStack = remember(currentEntry) {
            navController.backQueue
                .mapNotNull { it.destination.route }
                .joinToString(" > ")
        }
        Text(text = "Backstack: $backStack")

        NavHost(navController = navController, startDestination = "start") {
            composable("start") {
                Row { NavButton(route = "A"); NavButton(route = "B") }
            }
            composable("A") {
                Column {
                    Row { BackButton(); NavButton(route = "B") }
                    DemoScreen("Screen A")
                }
            }
            composable("B") {
                Column {
                    Row { BackButton(); NavButton(route = "A") }
                    DemoScreen("Screen B")
                }
            }
        }
    }
}