SharedViewModelHolder (Compose)

AppNavGraph.kt

// Inner navigation graph

val sharedViewModelHolder = SharedViewModelHolder<PostCreateViewModel>()
navGraphBuilder.navigation(
    route = PostCreateDestination.RootCreatePostRoute,
    startDestination = PostCreateDestination.SelectCreatablePost.createRoute()
) {
    composable(PostCreateDestination.SelectCreatablePost) {
        SelectCreatablePostScreen(
            modifier = modifier,
            navController = navController,
            sharedViewModel = sharedViewModelHolder.createNewSharedViewModel()
        )
    }
    composable(PostCreateDestination.CreatePost) { navBackStackEntry ->
        val arguments: PostCreateDestination.CreatePost.Arguments =
            PostCreateDestination.CreatePost.rememberArguments(navBackStackEntry)
        PostCreateScreen(
            modifier = modifier,
            navController = navController,
            sharedViewModel = sharedViewModelHolder.getSharedViewModel(),
            type = arguments.nonNullablePostType
        )
    }
    composable(PostCreateDestination.ChooseSpace) {
        ChooseSpaceScreen(
            modifier = modifier,
            navController = navController,
            sharedViewModel = sharedViewModelHolder.getSharedViewModel()
        )
    }
}

SharedViewModelHolder.kt

import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.lifecycle.ViewModel
import org.koin.androidx.compose.getViewModel
import java.lang.ref.WeakReference

class SharedViewModelHolder<T: ViewModel> {
    var sharedViewModelReference = WeakReference<T>(null)
}

@Composable
inline fun <reified T: ViewModel> SharedViewModelHolder<T>.createNewSharedViewModel(): T {
    val viewModel: T = getViewModel()
    LaunchedEffect(Unit) {
        if (sharedViewModelReference.get() !== viewModel) {
            sharedViewModelReference = WeakReference(viewModel)
        }
    }
    return viewModel
}

@Composable
inline fun <reified T : ViewModel> SharedViewModelHolder<T>.getSharedViewModel(): T {
    return sharedViewModelReference.get() ?: getViewModel()
}