Navigate with Parcelable arguments in Jetpack Compose

NavGraph.kt

composable(route = "screen_1") {
    Button(
        onClick = {
            navController.navigate(
                route = "screen_2",
                args = parcelableArgs {
                    arg { MyParcelableArgument() }
                    arg("named") { MyParcelableArgument() }
                },
            )
            // or simpler version
            navController.navigate(
                route = "screen_2",
                arg = MyParcelableArgument(),
            )
        }
    ) {
        Text("goto screen 2")
    }
}
composable(route = "screen_2") { entry ->
    val arg1: MyParcelableArgument = entry.rememberRequiredArgument()
    val arg2: MyParcelableArgument = entry.rememberRequiredArgument("named")
    // TODO: do something with args
}

NavigationUtils.kt

@file:Suppress("UnnecessaryVariable", "PackageDirectoryMismatch")

package androidx.navigation

import android.os.Bundle
import android.os.Parcelable
import androidx.compose.runtime.Composable
import androidx.compose.runtime.remember

@JvmInline
value class ParcelableArgBuilder(
    val args: MutableMap<String, Parcelable> = mutableMapOf(),
) {
    inline fun <reified T : Parcelable> arg(
        argName: String = T::class.qualifiedName!!,
        buildArg: () -> T,
    ) {
        args[argName] = buildArg()
    }
}

inline fun parcelableArgs(
    build: ParcelableArgBuilder.() -> Unit,
): Map<String, Parcelable> = ParcelableArgBuilder().apply(build).args.toMap()

fun NavController.navigate(
    route: String,
    navOptions: NavOptions? = null,
    navigatorExtras: Navigator.Extras? = null,
    args: Map<String, Parcelable>? = null,
) {
    if (args == null || args.isEmpty()) {
        navigate(route, navOptions, navigatorExtras)
        return
    }
    navigate(route, navOptions, navigatorExtras)
    val addedEntry: NavBackStackEntry = backQueue.last()
    val argumentBundle: Bundle = addedEntry.arguments ?: Bundle().also {
        addedEntry.arguments = it
    }
    args.forEach { (key: String, arg: Parcelable) ->
        argumentBundle.putParcelable(key, arg)
    }
}

inline fun <reified T : Parcelable> NavController.navigate(
    route: String,
    navOptions: NavOptions? = null,
    navigatorExtras: Navigator.Extras? = null,
    arg: T? = null,
) {
    if (arg == null) {
        navigate(route, navOptions, navigatorExtras)
        return
    }
    navigate(
        route = route,
        navOptions = navOptions,
        navigatorExtras = navigatorExtras,
        args = mapOf(T::class.qualifiedName!! to arg),
    )
}

inline fun <reified T : Parcelable> NavController.navigate(
    route: String,
    navOptions: NavOptions? = null,
    navigatorExtras: Navigator.Extras? = null,
    arg: Pair<String, T>? = null,
) {
    if (arg == null) {
        navigate(route, navOptions, navigatorExtras)
        return
    }
    navigate(
        route = route,
        navOptions = navOptions,
        navigatorExtras = navigatorExtras,
        args = mapOf(arg),
    )
}

fun NavBackStackEntry.requiredArguments(): Bundle = arguments ?: throw IllegalStateException("Arguments were expected, but none were provided!")

@Composable
inline fun <reified T : Parcelable> NavBackStackEntry.rememberRequiredArgument(
    key: String = T::class.qualifiedName!!,
): T = remember {
    requiredArguments().getParcelable<T>(key) ?: throw IllegalStateException("Expected argument with key: $key of type: ${T::class.qualifiedName!!}")
}

@Composable
inline fun <reified T : Parcelable> NavBackStackEntry.rememberArgument(
    key: String = T::class.qualifiedName!!,
): T? = remember {
    arguments?.getParcelable(key)
}

val NavBackStackEntry?.routeName: String?
    get() {
        if (this == null) return null
        val routeName: String? = destination.route?.substringBefore('/')
        return routeName
    }