NavHost.kt

NavHost.kt

@Composable
fun <T: Any> rememberMutableStateListOf(vararg elements: T): SnapshotStateList<T> {
  return rememberSaveable(
    saver = listSaver(
      save = { stateList ->
        if (stateList.isNotEmpty()) {
          val first = stateList.first()
          if (!canBeSaved(first)) {
            throw IllegalStateException("${first::class} cannot be saved. By default only types which can be stored in the Bundle class can be saved.")
          }
        }
        stateList.toList()
      },
      restore = { it.toMutableStateList() }
    )
  ) {
    elements.toList().toMutableStateList()
  }
}

sealed interface AppScreen : Parcelable {
  @Parcelize
  object ScreenA : AppScreen

  @Parcelize
  data class ScreenB(val id: Int) : AppScreen
}


@Composable fun NavHost() {


          val stack: SnapshotStateList<AppScreen> = rememberMutableStateListOf(AppScreen.ScreenA)

          val currentScreenState = remember(stack) {
            derivedStateOf { stack.last() }
          }

          val saveableStateHolder = rememberSaveableStateHolder()

          val navigateToScreen: (AppScreen) -> Unit = remember {
            { screen: AppScreen ->
              stack.add(screen)
            }
          }

          val onNavigateBack: () -> Unit = remember {
            {
              val removed = stack.removeLast()
              saveableStateHolder.removeState(removed)
            }
          }

          saveableStateHolder.SaveableStateProvider(key = currentScreenState.value) {
            when (val screen = currentScreenState.value) {
              AppScreen.ScreenA -> {
                rememberSaveable {
                  mutableStateOf(0)
                }

                Button(onClick = {
                  navigateToScreen(AppScreen.ScreenB(1))
                }) {
                  Text(text = "A")
                }
              }

              is AppScreen.ScreenB -> {
                Button(onClick = onNavigateBack) {
                  Text(text = "B ${screen.id}")
                }
              }
            }
          }
 }