Compose Navigation

Ambients.kt

import androidx.compose.Ambient
import androidx.navigation.NavController

object Ambients {
    val Navigator = Ambient.of<ComposableNavigator>()
    val NavController = Ambient.of<NavController>()
}

Arguments.kt

import androidx.compose.ambient
import androidx.compose.effectOf
import com.example.reddit.Ambients

fun <T> navArg(name: String) = effectOf<T> {
    val nav = +ambient(Ambients.NavController)
    val entry = nav.getBackStackEntry(nav.currentDestination!!.id)
    val args = entry.arguments
    val arg = args?.get(name) ?: error("No argument found with name $name")
    @Suppress("UNCHECKED_CAST")
    arg as T
}

fun <T> optionalNavArg(name: String) = effectOf<T?> {
    val nav = +ambient(Ambients.NavController)
    val entry = nav.getBackStackEntry(nav.currentDestination!!.id)
    val args = entry.arguments
    val arg = args?.get(name)
    @Suppress("UNCHECKED_CAST")
    arg as? T
}

ComposableNavigator.kt

import android.os.Bundle
import androidx.compose.Composable
import androidx.compose.FrameManager
import androidx.compose.Model
import androidx.navigation.*
import java.util.*

private typealias ComposableUnitLambda = @Composable () -> Unit
private val EmptyRoute: ComposableUnitLambda = {}

@Model
class ContentHolder(var value: ComposableUnitLambda = EmptyRoute)

@Navigator.Name("compose")
class ComposableNavigator : Navigator<Destination>() {

    private val stack = Stack<ComposableUnitLambda>()
    private var content = ContentHolder()

    val current: ComposableUnitLambda
        get() = content.value

    override fun createDestination(): Destination {
        return Destination(this)
    }

    override fun popBackStack(): Boolean {
        if (stack.empty()) {
            return false
        }

        content.value = stack.pop()
        return true
    }

    override fun navigate(
        destination: Destination,
        args: Bundle?,
        navOptions: NavOptions?,
        navigatorExtras: Extras?
    ): NavDestination? {
        FrameManager.ensureStarted()

        if (content.value !== EmptyRoute) {
            stack.push(content.value)
        }

        content.value = destination.content

        return destination
    }
}

@NavDestination.ClassType(Destination::class)
class Destination(navigator: ComposableNavigator) : NavDestination(navigator) {
    var content: ComposableUnitLambda = EmptyRoute
}

fun NavGraphBuilder.route(
    id: Int,
    content: @Composable () -> Unit) {
    addDestination(
        DestinationBuilder(provider[ComposableNavigator::class], id, content)
            .build())
}

fun NavGraphBuilder.route(
    id: Int,
    content: @Composable () -> Unit,
    builder: DestinationBuilder.() -> Unit) {
    addDestination(
        DestinationBuilder(provider[ComposableNavigator::class], id, content)
            .apply(builder)
            .build()
    )
}

class DestinationBuilder(
    navigator: ComposableNavigator,
    id: Int,
    val content: @Composable () -> Unit)
    : NavDestinationBuilder<Destination>(navigator, id) {

    override fun build(): Destination {
        return super.build().also { it.content = content }
    }
}

ComposeActivity.kt

import android.app.Activity
import android.os.Bundle
import android.util.Log
import androidx.compose.Composable
import androidx.compose.disposeComposition
import androidx.navigation.NavController
import androidx.navigation.NavGraphBuilder
import androidx.navigation.createGraph
import androidx.navigation.plusAssign
import androidx.ui.core.setContent
import com.example.reddit.Ambients

abstract class ComposeActivity : Activity() {
    @Composable
    abstract fun content(content: @Composable () -> Unit)
    abstract fun NavGraphBuilder.graph()
    abstract val initialRoute: Int

    private val graphId = 100
    private var navController: NavController? = null
    private val navigator = ComposableNavigator()

    override fun onBackPressed() {
        if (navController?.popBackStack() == true) {
            Log.v("Navigation", "Successfully navigated back")
        } else {
            super.onBackPressed()
        }
    }

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

        val controller = NavController(this).also {
            it.navigatorProvider += navigator
            it.graph = it.createGraph(graphId, initialRoute) {
                graph()
            }
        }

        navController = controller

        setContent {
            Ambients.NavController.Provider(controller) {
                Ambients.Navigator.Provider(navigator) {
                    content(navigator.current)
                }
            }
        }
    }

    override fun onDestroy() {
        super.onDestroy()
        navController = null
        disposeComposition()
    }
}

ExampleMainActivity.kt

// example usage
class MainActivity : ComposeActivity() {
    override val initialRoute = R.id.screen_one
    override fun NavGraphBuilder.graph() {
        route(R.id.screen_one) {
            val arg1 = +optionalNavArg<String>("someArg")
            ScreenOne(arg1)
        }
        route(R.id.screen_two) {
            val id = +navArg<String>("id")
            ScreenTwo(id)
        }
        route(R.id.screen_three) {
            ScreenThree()
        }
    }

    @Composable
    override fun content(content: @Composable () -> Unit) {
        AppTheme {
            Scaffold {
                content()
            }
        }
    }
}