Simple derivedStateOf navigation

Router.kt

import androidx.activity.OnBackPressedDispatcher
import androidx.activity.OnBackPressedDispatcherOwner
import androidx.activity.compose.BackHandler
import androidx.activity.compose.LocalOnBackPressedDispatcherOwner
import androidx.compose.animation.Crossfade
import androidx.compose.foundation.background
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.Spacer
import androidx.compose.foundation.layout.height
import androidx.compose.foundation.layout.padding
import androidx.compose.foundation.text.BasicText
import androidx.compose.material.Button
import androidx.compose.material.Text
import androidx.compose.runtime.*
import androidx.compose.runtime.saveable.LocalSaveableStateRegistry
import androidx.compose.runtime.saveable.SaveableStateRegistry
import androidx.compose.runtime.saveable.Saver
import androidx.compose.runtime.saveable.rememberSaveable
import androidx.compose.runtime.snapshots.SnapshotStateList
import androidx.compose.runtime.snapshots.SnapshotStateMap
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.platform.LocalLifecycleOwner
import androidx.compose.ui.text.TextStyle
import androidx.compose.ui.tooling.preview.Preview
import androidx.compose.ui.unit.dp
import androidx.lifecycle.LifecycleOwner

/**
 * Build and [rememberSaveable] a [RouterState] with a starting location of [startLocation]
 * and optional additional initial state constructed by [builder]. If the underlying [RouterState]
 * already exists it will be returned and [startLocation] and [builder] will be ignored.
 * Generally [startLocation] should be a constant matching the constant used to define the
 * initial "home" route in the [Router] that this [RouterState] is used with.
 */
@Composable
inline fun rememberSavedRouterState(
    startLocation: String,
    crossinline builder: RouterState.() -> Unit = {}
): RouterState =
    rememberSaveable(saver = RouterState.Saver) { RouterState(startLocation).apply(builder) }

/**
 * [Saveable][rememberSaveable] state for use with the [Router] composable.
 *
 * [RouterState] holds the current user-visible state of navigation: the current navigation
 * location and back stack, as well as the [saveable state][rememberSaveable] of each screen
 * that the user has visited.
 *
 * Use [Saver] to persist a [RouterState] using [rememberSaveable] or as a component of another
 * state object, or see [rememberSavedRouterState] for a convenient shortcut.
 *
 * [RouterState] operations operate on [androidx.compose.runtime.snapshots.Snapshot] state.
 * Use [androidx.compose.runtime.snapshots.Snapshot.withMutableSnapshot] when you need to perform
 * several operations atomically.
 */
class RouterState private constructor(
    private val routeStack: SnapshotStateList<BackStackRecord>
) {

    /**
     * Construct a [RouterState] with [startLocation] as the [currentRoute].
     */
    constructor(startLocation: String) : this(mutableStateListOf(BackStackRecord(startLocation)))

    /**
     * The [BackStackRecord] at the top of this [RouterState]'s back stack.
     */
    val currentRoute: BackStackRecord
        get() = routeStack.last()

    /**
     * `true` if [currentRoute] is the root destination of this [RouterState].
     * Calls to [pop] if [isAtRoot] are no-ops.
     */
    val isAtRoot: Boolean
        get() = routeStack.size == 1

    /**
     * Push a new [BackStackRecord] onto this [RouterState]'s back stack.
     * [currentRoute] will reflect the newly created record.
     */
    fun push(name: String, args: Map<String, Any?>? = null, tag: String? = null) {
        currentRoute.freeze()
        routeStack += BackStackRecord(name, args, tag)
    }

    /**
     * Pop [currentRoute] off of the back stack and discard it if [isAtRoot] is `false`.
     */
    fun pop() {
        if (!isAtRoot) routeStack.removeLast()
    }

    /**
     * Pop [currentRoute] until a [BackStackRecord] with a [BackStackRecord.tag] matching [tag]
     * is reached if at least one matching record exists. If there are no records matching [tag]
     * in the stack, [popToTag] does nothing.
     */
    fun popToTag(tag: String, inclusive: Boolean = false) {
        val inclusiveIndex = if (inclusive) 0 else 1
        val lastIndex = routeStack.indexOfLast { it.tag == tag } + inclusiveIndex
        if (lastIndex > 0) {
            routeStack.removeRange(lastIndex, routeStack.size)
        }
    }

    /**
     * Pop [currentRoute] until a [BackStackRecord] with a [BackStackRecord.location] matching
     * [location] is reached if at least one matching record exists. If there are no records
     * matching [location] in the stack, [popToLocation] does nothing.
     */
    fun popToLocation(location: String, inclusive: Boolean = false) {
        val inclusiveIndex = if (inclusive) 0 else 1
        val lastIndex = routeStack.indexOfLast { it.location == location } + inclusiveIndex
        if (lastIndex > 0) {
            routeStack.removeRange(lastIndex, routeStack.size)
        }
    }

    /**
     * Pop [currentRoute] until [isAtRoot].
     */
    fun popToRoot() {
        if (!isAtRoot) routeStack.removeRange(1, routeStack.size)
    }

    /**
     * Record of a [location] that the user has navigated to.
     *
     * * [location] is the name of a navigation location matching the name of a
     *   [RouteBuilder.route] defined by the [Router]. If no matching route is defined
     *   the [Router] will display a fallback or error state.
     * * [args] is a map of named values safe for persistence via [rememberSaveable].
     *   These may be interpreted by the [RouteBuilder.route] defined for [location].
     * * [tag] may be used to mark a position in the back stack. [RouterState.popToTag]
     *   pops all [BackStackRecord]s from the top of the back stack to the first [BackStackRecord]
     *   with a matching [tag].
     */
    @Stable
    class BackStackRecord private constructor(
        val location: String,
        val args: Map<String, Any?>,
        val tag: String?,
        private val restored: SnapshotStateMap<String, List<Any?>>
    ) {
        constructor(
            location: String,
            args: Map<String, Any?>? = null,
            tag: String? = null
        ) : this(
            location = location,
            args = args ?: emptyMap(),
            tag = tag,
            restored = mutableStateMapOf()
        )

        override fun hashCode(): Int = location.hashCode()
        override fun equals(other: Any?): Boolean {
            if (other !is BackStackRecord) return false
            return location == other.location && restored == other.restored
        }

        /**
         * The [location] of this [BackStackRecord].
         */
        operator fun component1(): String = location

        /**
         * The [args] of this [BackStackRecord].
         */
        operator fun component2(): Map<String, Any?> = args

        /**
         * Save the current state now. Performed automatically when navigating away
         * or when [currentRoute] state is being saved.
         */
        fun freeze() {
            restored.putAll(saveableStateRegistry.performSave())
        }

        /**
         * The [SaveableStateRegistry] provided as [LocalSaveableStateRegistry] by [Router]
         * when this [BackStackRecord] is the [RouterState.currentRoute].
         */
        val saveableStateRegistry = object : SaveableStateRegistry {

            /**
             * Implementation note:
             * valueProviders is NOT backed by snapshot state, as registering and unregistering
             * providers is always performed outside of composition.
             * Access should be synchronized on the collection itself.
             */
            private val valueProviders = mutableMapOf<String, MutableList<() -> Any?>>()

            /**
             * This method doesn't do us much good since [RouterState] and the [BackStackRecord]s
             * contained within are independent of any one place in composition. There is no
             * parent [SaveableStateRegistry] available for us to appeal to, so this method
             * always optimistically returns `true`. This makes for some nasty surprises at
             * actual state save time, but this would happen anyway when we're dealing with
             * arbitrary nested objects that may mutate over time.
             */
            override fun canBeSaved(value: Any): Boolean = true // lol

            /**
             * Implementation note:
             * [consumeRestored] takes advantage of [restored] as snapshot state, as it
             * is called during composition to make [rememberSaveable] state immediately available.
             */
            override fun consumeRestored(key: String): Any? = restored.remove(key)?.let { list ->
                list.first().also {
                    if (list.size > 1) {
                        restored[key] = list.drop(1)
                    }
                }
            }

            override fun performSave(): Map<String, List<Any?>> {
                val map = restored.toMutableMap()
                synchronized(valueProviders) {
                    valueProviders.forEach { (key, list) ->
                        if (list.size == 1) {
                            val value = list[0].invoke()
                            if (value != null) {
                                map[key] = arrayListOf<Any?>(value)
                            }
                        } else {
                            // nulls hold empty spaces
                            map[key] = list.map { it() }
                        }
                    }
                }
                return map
            }

            override fun registerProvider(
                key: String,
                valueProvider: () -> Any?
            ): SaveableStateRegistry.Entry {
                require(key.isNotBlank()) { "Registered key is empty or blank" }
                synchronized(valueProviders) {
                    @Suppress("UNCHECKED_CAST")
                    valueProviders.getOrPut(key) { mutableListOf() }.add(valueProvider)
                }
                return object : SaveableStateRegistry.Entry {
                    override fun unregister() {
                        synchronized(valueProviders) {
                            val list = valueProviders.remove(key)
                            list?.remove(valueProvider)
                            if (list != null && list.isNotEmpty()) {
                                // if there are other providers for this key return list
                                // back to the map
                                valueProviders[key] = list
                            }
                        }
                    }
                }
            }
        }

        companion object {
            fun saveFrozen(record: BackStackRecord): List<Any?> = with(record) {
                listOf(
                    location,
                    args.takeIf { it.isNotEmpty() },
                    tag,
                    restored
                )
            }

            @Suppress("UNCHECKED_CAST")
            fun restore(list: List<Any?>): BackStackRecord = BackStackRecord(
                location = list[0] as String,
                args = list[1] as Map<String, Any?>? ?: emptyMap(),
                tag = list[2] as String?,
                restored = mutableStateMapOf<String, List<Any?>>().apply {
                    putAll(list[3] as Map<String, List<Any?>>)
                }
            )
        }
    }

    companion object {
        val Saver = Saver<RouterState, List<List<Any?>>>(
            save = { value ->
                value.currentRoute.freeze()
                value.routeStack.map { BackStackRecord.saveFrozen(it) }
            },
            restore = { value ->
                RouterState(
                    mutableStateListOf<BackStackRecord>().apply {
                        addAll(value.map { BackStackRecord.restore(it) })
                    }
                )
            }
        )
    }
}

@DslMarker
annotation class RouteBuilderDsl

/**
 * Receiver scope for the [Router] composable. Use [route] to associate composable functions
 * with route locations.
 */
@RouteBuilderDsl
interface RouteBuilder {
    /**
     * Associate [content] with [location]. When the [RouterState.currentRoute]'s
     * [RouterState.BackStackRecord.location] matches [location] the [Router]
     * will compose [content].
     */
    fun route(location: String, content: @Composable (RouterState.BackStackRecord) -> Unit)
}

private class Routes : RouteBuilder {
    val routeMap = mutableStateMapOf<String, @Composable (RouterState.BackStackRecord) -> Unit>()

    override fun route(
        location: String,
        content: @Composable (RouterState.BackStackRecord) -> Unit
    ) {
        routeMap[location] = content
    }
}

typealias RouterDecoration = @Composable (
    location: RouterState.BackStackRecord,
    modifier: Modifier,
    content: @Composable (RouterState.BackStackRecord) -> Unit
) -> Unit

/**
 * A router to [state]'s [RouterState.currentRoute] content as declared by [routeBuilder].
 */
@Suppress("NAME_SHADOWING")
@Composable
fun Router(
    state: RouterState,
    modifier: Modifier = Modifier,
    decoration: RouterDecoration = { location, modifier, content ->
        Crossfade(location, modifier, content = content)
    },
    unavailableRoute: @Composable (String, RouterState) -> Unit = { route, _ ->
        BasicText(
            "Route not available: $route",
            Modifier.background(Color.Red),
            style = TextStyle(color = Color.Yellow)
        )
    },
    routeBuilder: RouteBuilder.() -> Unit
) {
    val currentBuilder by rememberUpdatedState(routeBuilder)
    val routes by remember { derivedStateOf { Routes().apply(currentBuilder) } }
    val currentRoute = state.currentRoute

    decoration(currentRoute, modifier) {
        val currentRouteComposable = routes.routeMap[it.location]
        if (currentRouteComposable != null) key(it) {
            CompositionLocalProvider(LocalSaveableStateRegistry provides it.saveableStateRegistry) {
                currentRouteComposable(it)
            }
        } else {
            unavailableRoute(it.location, state)
        }
    }

    BackHandler(enabled = !state.isAtRoot) {
        state.pop()
    }
}

// Preview materials below

private val UniversalSnapshotStateListSaver = Saver<SnapshotStateList<*>, List<*>>(
    save = { it.toList() },
    restore = { mutableStateListOf<Any?>().apply { addAll(it) } }
)

@Suppress("UNCHECKED_CAST")
private fun <T> snapshotStateListSaver() =
    UniversalSnapshotStateListSaver as Saver<SnapshotStateList<T>, List<*>>

private class FakeOnBackPressedDispatcherOwner(
    lifecycleOwner: LifecycleOwner
) : OnBackPressedDispatcherOwner, LifecycleOwner by lifecycleOwner {

    private val dispatcher = OnBackPressedDispatcher()

    override fun getOnBackPressedDispatcher(): OnBackPressedDispatcher = dispatcher
}

@Preview
@Composable
fun RouterPreview() {
    var counter by rememberSaveable { mutableStateOf(0) }
    val items = rememberSaveable(saver = snapshotStateListSaver()) { mutableStateListOf<String>() }
    val routerState = rememberSavedRouterState("home")

    val lifecycleOwner = LocalLifecycleOwner.current
    val fakeObpo = remember { FakeOnBackPressedDispatcherOwner(lifecycleOwner) }
    CompositionLocalProvider(LocalOnBackPressedDispatcherOwner provides fakeObpo) {
        Router(routerState) {
            route("home") {
                Column {
                    Button(onClick = { routerState.push("details") }) {
                        Text("See details")
                    }
                    Spacer(Modifier.height(16.dp))
                    Button(onClick = { routerState.push("bogus") }) {
                        Text("Bogus route")
                    }
                    Spacer(Modifier.height(16.dp))
                    items.forEach { item ->
                        Button(onClick = { routerState.push(item) }) {
                            Text("See $item")
                        }
                    }
                    Spacer(Modifier.height(16.dp))
                    Button(onClick = { items += "item ${++counter}" }) {
                        Text("Add another item")
                    }
                }
            }
            route("details") {
                Button(onClick = routerState::pop) {
                    Text("Go back")
                }
            }
            items.forEach { item ->
                route(item) {
                    Column {
                        Column(Modifier.padding(vertical = 16.dp)) {
                            var count by rememberSaveable { mutableStateOf(0) }
                            Text("Count: $count")
                            Button(onClick = { count++ }) {
                                Text("++")
                            }
                        }
                        Button(onClick = routerState::pop) {
                            Text("Go back")
                        }
                    }
                }
            }
        }
    }
}