Improved Navigation with MVVM with jetpack compose

MainActivity.kt

package com.debduttapanda.powernavigation

import android.os.Bundle
import androidx.activity.ComponentActivity
import androidx.activity.compose.setContent
import androidx.activity.viewModels
import androidx.compose.foundation.layout.fillMaxSize
import androidx.compose.material.MaterialTheme
import androidx.compose.material.Surface
import androidx.compose.material.Text
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.text.font.FontWeight
import androidx.compose.ui.unit.sp
import androidx.navigation.NavType
import androidx.navigation.compose.NavHost
import androidx.navigation.compose.composable
import androidx.navigation.compose.rememberNavController
import androidx.navigation.navArgument
import com.debduttapanda.powernavigation.ui.theme.PowerNavigationTheme

class MainActivity : ComponentActivity() {
    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        val pageAViewModel: PageAViewModel by viewModels()
        val pageBViewModel: PageBViewModel by viewModels()
        setContent {
            PowerNavigationTheme {
                Surface(
                    modifier = Modifier.fillMaxSize(),
                    color = MaterialTheme.colors.background
                ) {
                    Text(
                        "With MVVM",
                        fontSize = 24.sp,
                        color = Color(0xfff44336),
                        fontWeight = FontWeight.Bold
                    )
                    val navController = rememberNavController()
                    NavHost(navController = navController, startDestination = "page_a"){
                        composable(
                            "page_a"
                        ){
                            PageA(navController,pageAViewModel)
                        }
                        composable(
                            "page_b/{money}?bonus={bonus}",
                            arguments = listOf(
                                navArgument("money"){
                                    type = NavType.IntType
                                },
                                navArgument("bonus"){
                                    type = NavType.IntType
                                    defaultValue = 0
                                }
                            )
                        ){backStackEntry->
                            PageB(
                                navController,
                                pageBViewModel,
                                backStackEntry.arguments?.getInt("money"),
                                backStackEntry.arguments?.getInt("bonus"),
                            )
                        }
                    }
                }
            }
        }
    }
}

navigation_extension.kt

package com.debduttapanda.powernavigation

import androidx.compose.runtime.MutableState
import androidx.lifecycle.LifecycleOwner
import androidx.navigation.NavHostController

typealias NavigationCallback = (NavHostController,LifecycleOwner)->Unit

fun MutableState<NavigationCallback?>.navigate(block: NavigationCallback?){
    this.value = {navHostController, lifecycleOwner ->
        block?.invoke(navHostController,lifecycleOwner)
        this.value = null
    }
}

fun MutableState<NavigationCallback?>.forward(navHostController: NavHostController,lifecycleOwner: LifecycleOwner){
    this.value?.invoke(navHostController,lifecycleOwner)
}

PageA.kt

package com.debduttapanda.powernavigation

import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.Row
import androidx.compose.foundation.layout.fillMaxSize
import androidx.compose.foundation.text.KeyboardOptions
import androidx.compose.material.*
import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.platform.LocalLifecycleOwner
import androidx.compose.ui.text.font.FontWeight
import androidx.compose.ui.text.input.KeyboardType
import androidx.compose.ui.unit.sp
import androidx.navigation.NavHostController

@Composable
fun PageA(navController: NavHostController, pageAViewModel: PageAViewModel) {
    val owner = LocalLifecycleOwner.current
    LaunchedEffect(key1 = pageAViewModel.navigation.value){
        pageAViewModel.navigation.forward(navController,owner)
    }
    Column(
        modifier = Modifier.fillMaxSize(),
        horizontalAlignment = Alignment.CenterHorizontally,
        verticalArrangement = Arrangement.Center
    ){
        TextField(
            value = pageAViewModel.money.value.toString(),
            onValueChange = {
                try {
                    pageAViewModel.money.value = it.toInt()
                } catch (e: Exception) {
                    pageAViewModel.money.value = 0
                }
            },
            keyboardOptions = KeyboardOptions(
                keyboardType = KeyboardType.Number
            )
        )
        Row(
            verticalAlignment = Alignment.CenterVertically
        ){
            Text("Bonus")
            Checkbox(
                checked = pageAViewModel.sendBonus.value,
                onCheckedChange = {
                    pageAViewModel.sendBonus.value = it
                }
            )
        }

        Text(
            "Page A: Send Money",
            color = Color(0xfff44336),
            fontSize = 24.sp,
            fontWeight = FontWeight.Bold
        )


        Button(
            onClick = {
                pageAViewModel.onSendClick()
            },
            colors = ButtonDefaults.buttonColors(
                backgroundColor = Color(0xfff44336),
                contentColor = Color.White
            )
        ) {
            Text("Send Money")
        }
        if((pageAViewModel.totalReceived.value)>0){
            Text("Acknowledgement of ${pageAViewModel.totalReceived.value} received")
        }
    }
}

PageAViewModel.kt

package com.debduttapanda.powernavigation

import androidx.compose.runtime.mutableStateOf
import androidx.lifecycle.ViewModel

class PageAViewModel: ViewModel() {
    val navigation = mutableStateOf<NavigationCallback?>(null)
    val money = mutableStateOf(0)
    val sendBonus = mutableStateOf(false)
    val totalReceived = mutableStateOf(0)
    fun onSendClick() {
        navigation.navigate{navHostController, lifecycleOwner ->
            navHostController
                .currentBackStackEntry
                ?.savedStateHandle
                ?.getLiveData<Int>("totalReceived")
                ?.observe(lifecycleOwner){
                    totalReceived.value = it
                    navHostController
                        .currentBackStackEntry
                        ?.savedStateHandle
                        ?.remove<Int>("totalReceived")
                }
            if(!sendBonus.value){
                navHostController.navigate("page_b/${money.value}")
            }
            else{
                navHostController.navigate("page_b/${money.value}?bonus=50")
            }
        }
    }
}

PageB.kt

package com.debduttapanda.powernavigation

import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.fillMaxSize
import androidx.compose.material.Button
import androidx.compose.material.ButtonDefaults
import androidx.compose.material.Text
import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.platform.LocalLifecycleOwner
import androidx.compose.ui.text.font.FontWeight
import androidx.compose.ui.unit.sp
import androidx.navigation.NavHostController

@Composable
fun PageB(
    navController: NavHostController,
    pageBViewModel: PageBViewModel,
    receivedMoney: Int?,
    bonus: Int?
) {
    val owner = LocalLifecycleOwner.current
    LaunchedEffect(key1 = Unit){
        pageBViewModel.setArguments(receivedMoney,bonus)
    }
    LaunchedEffect(key1 = pageBViewModel.navigation.value){
        pageBViewModel.navigation.forward(navController,owner)
    }
    Column(
        modifier = Modifier.fillMaxSize(),
        horizontalAlignment = Alignment.CenterHorizontally,
        verticalArrangement = Arrangement.Center
    ){
        Text(
            "Page B: received money ${pageBViewModel.receivedMoney.value} ${if(pageBViewModel.bonus.value>0) "got $bonus bonus" else "): no bonus this month"}",
            color = Color(0xfff44336),
            fontSize = 24.sp,
            fontWeight = FontWeight.Bold
        )
        Button(
            onClick = {
                pageBViewModel.onGoBack()
            },
            colors = ButtonDefaults.buttonColors(
                backgroundColor = Color(0xfff44336),
                contentColor = Color.White
            )
        ) {
            Text("Go back and pay me again")
        }
    }
}

PageBViewModel.kt

package com.debduttapanda.powernavigation

import androidx.compose.runtime.mutableStateOf
import androidx.lifecycle.ViewModel

class PageBViewModel: ViewModel() {
    val navigation = mutableStateOf<NavigationCallback?>(null)
    val receivedMoney = mutableStateOf(0)
    val bonus = mutableStateOf(0)
    fun onGoBack() {
        navigation.navigate { navHostController, lifecycleOwner ->
            navHostController
                .previousBackStackEntry
                ?.savedStateHandle
                ?.set("totalReceived",(receivedMoney.value)+(bonus.value))
            navHostController.navigateUp()
        }
    }

    fun setArguments(_receivedMoney: Int?, _bonus: Int?) {
        receivedMoney.value = _receivedMoney?:0
        bonus.value = _bonus?:0
    }
}