multiChildStack for Decompose

MultiChildStack.kt

package ru.mobileup.template.core.utils

import com.arkivanov.decompose.Child
import com.arkivanov.decompose.ComponentContext
import com.arkivanov.decompose.router.stack.ChildStack
import com.arkivanov.decompose.router.stack.StackNavigationSource
import com.arkivanov.decompose.router.stack.childStack
import com.arkivanov.decompose.value.MutableValue
import com.arkivanov.decompose.value.Value
import com.arkivanov.essenty.lifecycle.Lifecycle
import com.arkivanov.essenty.lifecycle.doOnDestroy
import kotlinx.serialization.KSerializer
import kotlinx.serialization.builtins.PairSerializer
import kotlinx.serialization.builtins.serializer
import okio.withLock
import java.util.concurrent.locks.ReentrantLock

// Creates a child stack that supports duplicated configurations.
// It uses a normal child stack under the hood with configurations transformed to Pair<C, Int>:
// For example [A, B, B, B, C, C] is transformed to [(A, 0), (B, 0), (B, 1), (B, 2), (C, 0), (C, 1)]
fun <C : Any, T : Any> ComponentContext.multiChildStack(
    source: StackNavigationSource<C>,
    serializer: KSerializer<C>?,
    initialStack: () -> List<C>,
    key: String = "DefaultChildStack",
    handleBackButton: Boolean = false,
    childFactory: (configuration: C, ComponentContext) -> T,
): Value<ChildStack<C, T>> {
    return childStack(
        source = MappedStackNavigationSource(source),
        serializer = serializer?.let { PairSerializer(it, Int.serializer()) },
        initialStack = { initialStack().mapToPaired() },
        key = key,
        handleBackButton = handleBackButton,
        childFactory = { configuration: Pair<C, Int>, componentContext: ComponentContext ->
            childFactory(configuration.first, componentContext)
        }
    ).mapToNormal(lifecycle)
}

fun <C : Any, T : Any> ComponentContext.multiChildStack(
    source: StackNavigationSource<C>,
    serializer: KSerializer<C>?,
    initialConfiguration: C,
    key: String = "DefaultChildStack",
    handleBackButton: Boolean = false,
    childFactory: (configuration: C, ComponentContext) -> T
): Value<ChildStack<C, T>> =
    multiChildStack(
        source = source,
        serializer = serializer,
        initialStack = { listOf(initialConfiguration) },
        key = key,
        handleBackButton = handleBackButton,
        childFactory = childFactory,
    )

private fun <C : Any, T : Any> Value<ChildStack<Pair<C, Int>, T>>.mapToNormal(lifecycle: Lifecycle): Value<ChildStack<C, T>> {
    val result = MutableValue(this.value.mapToNormal())

    if (lifecycle.state != Lifecycle.State.DESTROYED) {
        val cancellation = observe { result.value = it.mapToNormal() }
        lifecycle.doOnDestroy {
            cancellation.cancel()
        }
    }

    return result
}

private fun <C : Any, T : Any> ChildStack<Pair<C, Int>, T>.mapToNormal(): ChildStack<C, T> {
    return ChildStack(
        active = Child.Created(this.active.configuration.first, this.active.instance),
        backStack = this.backStack.map { Child.Created(it.configuration.first, it.instance) }
    )
}

private fun <C : Any> List<Pair<C, Int>>.mapToNormal(): List<C> {
    return map { it.first }
}

private fun <C : Any> List<C>.mapToPaired(): List<Pair<C, Int>> {
    return mapIndexed { index, configuration ->
        Pair(configuration, countBeforePosition(configuration, index))
    }
}

private fun <C : Any> List<C>.countBeforePosition(configuration: C, position: Int): Int {
    return withIndex().count { (index, config) -> index < position && config == configuration }
}

private typealias EventObserver<T> = (StackNavigationSource.Event<T>) -> Unit

private class MappedStackNavigationSource<C : Any>(
    private val original: StackNavigationSource<C>
) : StackNavigationSource<Pair<C, Int>> {

    private val lock = ReentrantLock() // TODO: use cross-platform Lock
    private var observersMapping = mutableMapOf<EventObserver<Pair<C, Int>>, EventObserver<C>>()

    override fun subscribe(observer: EventObserver<Pair<C, Int>>) = lock.withLock {
        val mappedObserver = observer.mapObserver()
        observersMapping[observer] = mappedObserver
        original.subscribe(mappedObserver)
    }

    override fun unsubscribe(observer: EventObserver<Pair<C, Int>>) = lock.withLock {
        val mappedObserver = observersMapping[observer] ?: return@withLock
        observersMapping.remove(observer)
        original.unsubscribe(mappedObserver)
    }
}

private fun <C : Any> EventObserver<Pair<C, Int>>.mapObserver(): EventObserver<C> {
    return { event: StackNavigationSource.Event<C> ->
        invoke(
            StackNavigationSource.Event(
                transformer = { stack ->
                    event.transformer(stack.mapToNormal()).mapToPaired()
                },
                onComplete = { newStack, oldStack ->
                    event.onComplete(newStack.mapToNormal(), oldStack.mapToNormal())
                }
            )
        )
    }
}