Nested navigation in 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.compose.NavHost
import androidx.navigation.compose.composable
import androidx.navigation.compose.navigation
import androidx.navigation.compose.rememberNavController
import com.debduttapanda.powernavigation.ui.theme.PowerNavigationTheme

class MainActivity : ComponentActivity() {
    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        val pageAViewModel: PageAViewModel by viewModels()
        val page_Navigation_1_1_ViewModel: Page_Navigation_1_1_ViewModel by viewModels()
        val page_Navigation_1_2_ViewModel: Page_Navigation_1_2_ViewModel 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)
                        }
                        navigation(
                            route = "nested_navigation",
                            startDestination = "nested_nav_1_1"
                        ){
                            composable(
                                "nested_nav_1_1"
                            ){backStackEntry->
                                Page_Navigation_1_1(
                                    navController,
                                    page_Navigation_1_1_ViewModel
                                )
                            }
                            composable(
                                "nested_nav_1_2"
                            ){backStackEntry->
                                Page_Navigation_1_2(
                                    navController,
                                    page_Navigation_1_2_ViewModel
                                )
                            }
                        }
                    }
                }
            }
        }
    }
}

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)
}

Page_Navigation_1_1.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 Page_Navigation_1_1(
    navController: NavHostController,
    pageBViewModel: Page_Navigation_1_1_ViewModel
) {
    val owner = LocalLifecycleOwner.current
    LaunchedEffect(key1 = pageBViewModel.navigation.value){
        pageBViewModel.navigation.forward(navController,owner)
    }
    Column(
        modifier = Modifier.fillMaxSize(),
        horizontalAlignment = Alignment.CenterHorizontally,
        verticalArrangement = Arrangement.Center
    ){
        Text(
            "Page navigation 1_1",
            color = Color(0xfff44336),
            fontSize = 24.sp,
            fontWeight = FontWeight.Bold
        )
        Button(
            onClick = {
                pageBViewModel.onNextClick()
            },
            colors = ButtonDefaults.buttonColors(
                backgroundColor = Color(0xfff44336),
                contentColor = Color.White
            )
        ) {
            Text("Next")
        }
        Button(
            onClick = {
                pageBViewModel.onGoBack()
            },
            colors = ButtonDefaults.buttonColors(
                backgroundColor = Color(0xfff44336),
                contentColor = Color.White
            )
        ) {
            Text("Back")
        }
    }
}

Page_Navigation_1_1_ViewModel.kt

package com.debduttapanda.powernavigation

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

class Page_Navigation_1_1_ViewModel: ViewModel() {
    val navigation = mutableStateOf<NavigationCallback?>(null)
    fun onGoBack() {
        navigation.navigate { navHostController, lifecycleOwner ->
            navHostController.navigateUp()
        }
    }

    fun onNextClick() {
        navigation.navigate { navHostController, lifecycleOwner ->
            navHostController.navigate("nested_nav_1_2")
        }
    }
}

Page_Navigation_1_2.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 Page_Navigation_1_2(
    navController: NavHostController,
    pageBViewModel: Page_Navigation_1_2_ViewModel
) {
    val owner = LocalLifecycleOwner.current
    LaunchedEffect(key1 = pageBViewModel.navigation.value){
        pageBViewModel.navigation.forward(navController,owner)
    }
    Column(
        modifier = Modifier.fillMaxSize(),
        horizontalAlignment = Alignment.CenterHorizontally,
        verticalArrangement = Arrangement.Center
    ){
        Text(
            "Page navigation 1_2",
            color = Color(0xfff44336),
            fontSize = 24.sp,
            fontWeight = FontWeight.Bold
        )
        Button(
            onClick = {
                pageBViewModel.onGoBack()
            },
            colors = ButtonDefaults.buttonColors(
                backgroundColor = Color(0xfff44336),
                contentColor = Color.White
            )
        ) {
            Text("Back")
        }
    }
}

Page_Navigation_1_2_ViewModel.kt

package com.debduttapanda.powernavigation

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

class Page_Navigation_1_2_ViewModel: ViewModel() {
    val navigation = mutableStateOf<NavigationCallback?>(null)
    fun onGoBack() {
        navigation.navigate { navHostController, lifecycleOwner ->
            navHostController.navigateUp()
        }
    }
}

PageA.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 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
    ){


        Text(
            "Page A",
            color = Color(0xfff44336),
            fontSize = 24.sp,
            fontWeight = FontWeight.Bold
        )
        Button(
            onClick = {
                pageAViewModel.onSendClick()
            },
            colors = ButtonDefaults.buttonColors(
                backgroundColor = Color(0xfff44336),
                contentColor = Color.White
            )
        ) {
            Text("Go to nested navigation")
        }
    }
}

PageAViewModel.kt

package com.debduttapanda.powernavigation

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

class PageAViewModel: ViewModel() {
    val navigation = mutableStateOf<NavigationCallback?>(null)
    fun onSendClick() {
        navigation.navigate{navHostController, lifecycleOwner ->
            navHostController.navigate("nested_navigation")
        }
    }
}