Simple forward and back navigation in Jetpack Compose

BackStack.kt

package com.example.yourapp.nav

import androidx.compose.runtime.Composable
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.setValue
import androidx.lifecycle.ViewModel

class BackStack : ViewModel() {
    val callback = BackStackCallback(this)
    private var stack: List<@Composable () -> Unit> by mutableStateOf(listOf())
    val activeScreen: @Composable () -> Unit
        get() = this.stack.last()
    val size: Int
        get() = this.stack.size
    fun pop() {
        this.stack = this.stack.toMutableList().also {
            it.removeLast()
        }
    }
    fun push(destination: @Composable () -> Unit) {
        this.stack = this.stack.toMutableList().also {
            it += destination
        }
    }
}

BackStackCallback.kt

package com.example.yourapp.nav

import androidx.activity.OnBackPressedCallback

class BackStackCallback(backStack: BackStack) : OnBackPressedCallback(true) {
    private val stack = backStack
    override fun handleOnBackPressed() {
        if (this.stack.size > 1) {
            this.stack.pop()
        }
    }
}

MainActivity.kt

class MainActivity : AppCompatActivity() {
    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)

        val backStack = BackStack()
        onBackPressedDispatcher.addCallback(this, backStack.callback)

        setContent {
            MaterialTheme {
                // A surface container using the 'background' color from the theme
                Surface(color = MaterialTheme.colors.background) {
                    Providers(AmbientBackStack provides backStack) {
                        NavHost(AmbientBackStack.current) { 
                            val backStack = AmbientBackStack.current
                            Button({ backStack.push({ Text("Use the back button to go back") }) }) { 
                                Text("Navigate forward")
                            }
                        }
                    }
                }
            }
        }
    }
}

NavHost.kt

package com.example.yourapp.nav

import androidx.compose.runtime.Composable
import androidx.compose.runtime.ambientOf
import androidx.compose.runtime.remember

val AmbientBackStack = ambientOf<BackStack> { error("BackStack must have its " +
        "BackStackCallback registered and be provided") }

@Composable
fun NavHost(backStack: BackStack, firstChild: @Composable () -> Unit) {
        remember { backStack.push(firstChild) }
        backStack.activeScreen()
}