Weekend navigation experiments extracted from home automation project

BackStack.kt

import androidx.compose.runtime.Stable

/**
 * A caller-supplied stack of [Record]s for presentation with the [Navigator] composable.
 * Iteration order is [top]-first.
 */
@Stable
interface BackStack<R : BackStack.Record> : Iterable<R> {
    /**
     * The number of records contained in this [BackStack] that will be seen by an iterator.
     */
    val size: Int

    /**
     * Attempt to pop the top item off of the back stack, returning the popped [Record]
     * if popping was successful or `null` if no entry was popped.
     */
    fun pop(): R?

    interface Record {
        /**
         * A value that identifies this record uniquely, even if it shares the same
         * [route] with another record. This key may be used by [BackStackRecordLocalProvider]s
         * to associate presentation data with a record across composition recreation.
         *
         * [key] MUST NOT change for the life of the record.
         */
        val key: String

        /**
         * The name of the route that should present this record
         */
        val route: String
    }
}

/**
 * `true` if the [BackStack] contains no records. [BackStack.top] will return `null`.
 */
val BackStack<*>.isEmpty: Boolean get() = size == 0

/**
 * `true` if the [BackStack] contains exactly one record.
 */
val BackStack<*>.isAtRoot: Boolean get() = size == 1

BackStackRecordLocalProvider.kt

import androidx.compose.runtime.Composable
import androidx.compose.runtime.ProvidedValue
import androidx.compose.runtime.compositionLocalOf
import androidx.compose.runtime.key

fun interface BackStackRecordLocalProvider<in R : BackStack.Record> {
    @Composable
    fun providedValuesFor(record: R): ProvidedValues
}

fun interface ProvidedValues {
    @Composable
    fun provideValues(): List<ProvidedValue<*>>
}

class CompositeProvidedValues(
    private val list: List<ProvidedValues>
) : ProvidedValues {
    @Composable
    override fun provideValues(): List<ProvidedValue<*>> = buildList {
        list.forEach {
            addAll(key(it) { it.provideValues() })
        }
    }
}

@Composable
fun <R : BackStack.Record> providedValuesForBackStack(
    backStack: BackStack<R>,
    stackLocalProviders: List<BackStackRecordLocalProvider<R>> = emptyList(),
    includeDefaults: Boolean = true,
): Map<R, ProvidedValues> = buildMap {
    backStack.forEach { record ->
        key(record) {
            put(
                record,
                CompositeProvidedValues(
                    buildList {
                        if (includeDefaults) {
                            LocalBackStackRecordLocalProviders.current.forEach {
                                add(key(it) { it.providedValuesFor(record) })
                            }
                        }
                        stackLocalProviders.forEach {
                            add(key(it) { it.providedValuesFor(record) })
                        }
                    }
                )
            )
        }
    }
}

@Suppress("RemoveExplicitTypeArguments")
val LocalBackStackRecordLocalProviders =
    compositionLocalOf<List<BackStackRecordLocalProvider<BackStack.Record>>> {
        listOf(
            SaveableStateRegistryBackStackRecordLocalProvider,
            ViewModelBackStackRecordLocalProvider
        )
    }

Navigator.kt

import androidx.activity.compose.BackHandler
import androidx.compose.animation.Crossfade
import androidx.compose.foundation.background
import androidx.compose.foundation.text.BasicText
import androidx.compose.runtime.*
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.text.TextStyle

/**
 * [DslMarker] for the [NavigationBuilder] DSL.
 */
@DslMarker
annotation class NavigationBuilderDsl

/**
 * Declarative builder scope for navigation routes; used with [Navigator].
 */
@NavigationBuilderDsl
interface NavigationBuilder<R> {
    fun route(route: String, content: @Composable (R) -> Unit)
}

/**
 * Presentation logic for currently visible routes of a [Navigator].
 */
@Stable
interface NavigatorRouteDecoration {
    @Composable
    fun <T> DecoratedContent(
        arg: T,
        backStackDepth: Int,
        modifier: Modifier,
        content: @Composable (T) -> Unit
    )
}

private class NavigationMap<R> : NavigationBuilder<R> {
    val routes = mutableMapOf<String, @Composable (R) -> Unit>()

    override fun route(route: String, content: @Composable (R) -> Unit) {
        routes[route] = content
    }
}

/**
 * Default values and common alternatives used by the [Navigator] family of composables.
 */
object NavigatorDefaults {

    /**
     * The default [NavigatorRouteDecoration] used by [Navigator].
     */
    object DefaultDecoration : NavigatorRouteDecoration {
        @Composable
        override fun <T> DecoratedContent(
            arg: T,
            backStackDepth: Int,
            modifier: Modifier,
            content: @Composable (T) -> Unit
        ) {
            Crossfade(arg, modifier, content = content)
        }
    }

    /**
     * An empty [NavigatorRouteDecoration] that emits the content with no surrounding decoration
     * or logic.
     */
    object EmptyDecoration : NavigatorRouteDecoration {
        @Composable
        override fun <T> DecoratedContent(
            arg: T,
            backStackDepth: Int,
            modifier: Modifier,
            content: @Composable (T) -> Unit
        ) {
            content(arg)
        }
    }

    /**
     * Bright ugly error text telling a developer they didn't provide a route
     * that a [BackStack] asked for.
     */
    val UnavailableRoute: @Composable (String) -> Unit = { route ->
        BasicText(
            "Route not available: $route",
            Modifier.background(Color.Red),
            style = TextStyle(color = Color.Yellow)
        )
    }
}

@Composable
fun <R : BackStack.Record> Navigator(
    backStack: BackStack<R>,
    modifier: Modifier = Modifier,
    enableBackHandler: Boolean = true,
    providedValues: Map<R, ProvidedValues> = providedValuesForBackStack(backStack),
    decoration: NavigatorRouteDecoration = NavigatorDefaults.DefaultDecoration,
    unavailableRoute: @Composable (String) -> Unit = NavigatorDefaults.UnavailableRoute,
    navigationBuilder: NavigationBuilder<R>.() -> Unit
) {
    BackHandler(enabled = enableBackHandler && !backStack.isAtRoot) {
        backStack.pop()
    }

    BasicNavigator(
        backStack = backStack,
        providedValues = providedValues,
        modifier = modifier,
        decoration = decoration,
        unavailableRoute = unavailableRoute,
        navigationBuilder = navigationBuilder
    )
}

@Composable
fun <R : BackStack.Record> BasicNavigator(
    backStack: BackStack<R>,
    providedValues: Map<R, ProvidedValues>,
    modifier: Modifier = Modifier,
    decoration: NavigatorRouteDecoration = NavigatorDefaults.EmptyDecoration,
    unavailableRoute: @Composable (String) -> Unit = NavigatorDefaults.UnavailableRoute,
    navigationBuilder: NavigationBuilder<R>.() -> Unit
) {
    val currentBuilder by rememberUpdatedState(navigationBuilder)
    val routes by remember { derivedStateOf { NavigationMap<R>().apply(currentBuilder) } }

    val activeContentProviders = buildList {
        for (record in backStack) {
            val provider = key(record.key) {
                val routeName = record.route
                val currentRouteContent by rememberUpdatedState(
                    routes.routes[routeName] ?: { unavailableRoute(routeName) }
                )
                val currentRecord by rememberUpdatedState(record)
                remember {
                    movableContentOf {
                        currentRouteContent(currentRecord)
                    }
                }
            }
            add(record to provider)
        }
    }

    decoration.DecoratedContent(
        activeContentProviders.first(),
        backStack.size,
        modifier
    ) { (record, provider) ->
        val values = providedValues[record]?.provideValues()
        val providedLocals = remember(values) {
            values?.toTypedArray() ?: emptyArray()
        }
        CompositionLocalProvider(*providedLocals) {
            provider.invoke()
        }
    }
}

SaveableBackStack.kt

import androidx.compose.runtime.Composable
import androidx.compose.runtime.mutableStateListOf
import androidx.compose.runtime.saveable.Saver
import androidx.compose.runtime.saveable.rememberSaveable
import java.util.*

@Composable
fun rememberSaveableBackStack(
    init: SaveableBackStack.() -> Unit
): SaveableBackStack = rememberSaveable(saver = SaveableBackStack.Saver) {
    SaveableBackStack().apply(init)
}

fun SaveableBackStack.push(
    route: String,
    args: Map<String, Any?> = emptyMap()
) {
    push(SaveableBackStack.Record(route, args))
}

inline fun SaveableBackStack.popUntil(predicate: (SaveableBackStack.Record) -> Boolean) {
    while (topRecord?.let(predicate) == false) pop()
}

/**
 * A [BackStack] that supports saving its state via [rememberSaveable].
 * See [rememberSaveableBackStack].
 */
class SaveableBackStack : BackStack<SaveableBackStack.Record> {

    private val entryList = mutableStateListOf<Record>()

    override val size: Int
        get() = entryList.size

    override fun iterator(): Iterator<Record> = entryList.iterator()

    val topRecord: Record?
        get() = entryList.firstOrNull()

    fun push(record: Record) {
        entryList.add(0, record)
    }

    override fun pop(): Record? = entryList.removeFirstOrNull()

    data class Record(
        override val route: String,
        val args: Map<String, Any?> = emptyMap(),
        override val key: String = UUID.randomUUID().toString()
    ) : BackStack.Record {
        companion object {
            val Saver: Saver<Record, List<Any>> = Saver(
                save = { value ->
                    buildList {
                        add(value.route)
                        add(value.args)
                        add(value.key)
                    }
                },
                restore = { list ->
                    @Suppress("UNCHECKED_CAST")
                    Record(
                        route = list[0] as String,
                        args = list[1] as Map<String, Any?>,
                        key = list[2] as String
                    )
                }
            )
        }
    }

    companion object {
        val Saver = Saver<SaveableBackStack, List<Any>>(
            save = { value ->
                value.entryList.map {
                    with(Record.Saver) {
                        save(it)!!
                    }
                }
            },
            restore = { list ->
                SaveableBackStack().also { backStack ->
                    list.mapTo(backStack.entryList) {
                        @Suppress("UNCHECKED_CAST")
                        Record.Saver.restore(it as List<Any>)!!
                    }
                }
            }
        )
    }
}

SaveableStateRegistryBackStackRecordLocalProvider.kt

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.SnapshotStateMap

/**
 * A [BackStackRecordLocalProvider] that provides a [SaveableStateRegistry] for each record.
 */
object SaveableStateRegistryBackStackRecordLocalProvider
    : BackStackRecordLocalProvider<BackStack.Record> {
    @Composable
    override fun providedValuesFor(record: BackStack.Record): ProvidedValues {
        val childRegistry = rememberSaveable(
            record,
            saver = BackStackRecordLocalSaveableStateRegistry.Saver,
            key = record.key
        ) {
            BackStackRecordLocalSaveableStateRegistry(mutableStateMapOf())
        }
        // This write depends on childRegistry.parentRegistry being snapshot state backed
        childRegistry.parentRegistry = LocalSaveableStateRegistry.current
        return remember(childRegistry) {
            @Suppress("ObjectLiteralToLambda")
            object : ProvidedValues {
                val list = listOf(LocalSaveableStateRegistry provides childRegistry)
                @Composable
                override fun provideValues(): List<ProvidedValue<*>> {
                    remember {
                        object : RememberObserver {
                            override fun onForgotten() {
                                childRegistry.saveForContentLeavingComposition()
                            }

                            override fun onRemembered() {}
                            override fun onAbandoned() {}
                        }
                    }
                    return list
                }
            }
        }
    }
}

private class BackStackRecordLocalSaveableStateRegistry(
    // Note: restored is snapshot-backed because consumeRestored runs in composition
    // and must be rolled back if composition does not commit
    private val restored: SnapshotStateMap<String, List<Any?>>
) : SaveableStateRegistry {
    var parentRegistry: SaveableStateRegistry? by mutableStateOf(null)

    private val valueProviders = mutableMapOf<String, MutableList<() -> Any?>>()

    override fun canBeSaved(value: Any): Boolean = parentRegistry?.canBeSaved(value) != false

    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()
        saveInto(map)
        return map
    }

    override fun registerProvider(
        key: String,
        valueProvider: () -> Any?
    ): SaveableStateRegistry.Entry {
        require(key.isNotBlank()) { "Registered key is empty or blank" }
        synchronized(valueProviders) {
            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
                    }
                }
            }
        }
    }

    fun saveForContentLeavingComposition() {
        saveInto(restored)
    }

    private fun saveInto(map: MutableMap<String, List<Any?>>) {
        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() }
                }
            }
        }
    }

    companion object {
        val Saver = Saver<BackStackRecordLocalSaveableStateRegistry, Map<String, List<Any?>>>(
            save = { value ->
                value.performSave()
            },
            restore = { value ->
                BackStackRecordLocalSaveableStateRegistry(
                    mutableStateMapOf<String, List<Any?>>().apply { putAll(value) }
                )
            }
        )
    }
}

ViewModelBackStackRecordLocalProvider.kt

import android.app.Activity
import android.content.Context
import android.content.ContextWrapper
import androidx.compose.runtime.Composable
import androidx.compose.runtime.RememberObserver
import androidx.compose.runtime.remember
import androidx.compose.ui.platform.LocalContext
import androidx.lifecycle.ViewModel
import androidx.lifecycle.ViewModelStore
import androidx.lifecycle.ViewModelStoreOwner
import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
import androidx.lifecycle.viewmodel.compose.viewModel

private fun Context.findActivity(): Activity? {
    var context = this
    while (context is ContextWrapper) {
        if (context is Activity) return context
        context = context.baseContext
    }
    return null
}

/**
 * A [BackStackRecordLocalProvider] that provides [LocalViewModelStoreOwner]
 */
object ViewModelBackStackRecordLocalProvider : BackStackRecordLocalProvider<BackStack.Record> {
    @Composable
    override fun providedValuesFor(record: BackStack.Record): ProvidedValues {
        // Implementation note: providedValuesFor stays in the composition as long as the
        // back stack entry is present, which makes it safe for us to use composition
        // forget/abandon to clear the associated ViewModelStore if the host activity
        // isn't in the process of changing configurations.
        val containerViewModel = viewModel<BackStackRecordLocalProviderViewModel>()
        val viewModelStore = containerViewModel.viewModelStoreForKey(record.key)
        val activity = LocalContext.current.findActivity()
        remember(record, viewModelStore) {
            object : RememberObserver {
                override fun onAbandoned() {
                    disposeIfNotChangingConfiguration()
                }

                override fun onForgotten() {
                    disposeIfNotChangingConfiguration()
                }

                override fun onRemembered() {
                }

                fun disposeIfNotChangingConfiguration() {
                    if (activity?.isChangingConfigurations != true) {
                        containerViewModel.removeViewModelStoreOwnerForKey(record.key)?.clear()
                    }
                }
            }
        }
        return remember(viewModelStore) {
            val list = listOf(
                LocalViewModelStoreOwner provides ViewModelStoreOwner { viewModelStore }
            )
            ProvidedValues { list }
        }
    }
}

class BackStackRecordLocalProviderViewModel : ViewModel() {
    private val owners = mutableMapOf<String, ViewModelStore>()

    fun viewModelStoreForKey(key: String): ViewModelStore =
        owners.getOrPut(key) { ViewModelStore() }

    fun removeViewModelStoreOwnerForKey(key: String): ViewModelStore? = owners.remove(key)

    override fun onCleared() {
        owners.forEach { (_, store) ->
            store.clear()
        }
    }
}