A `StateFlow` builder that can subscribe to other StateFlows that are read within the builder block. This is intended for deriving state from one or more other states.

ContactViewModel.kt

@OptIn(ExperimentalCoroutinesApi::class)
class ContactViewModel() : ViewModel() {
    val firstName = MutableStateFlow("")
    val lastName = MutableStateFlow("")
    val fullName by DerivedStateFlow {
        "${+::firstName} ${+::lastName}"
    }
}

DerivedStateFlow.kt

@file:OptIn(ExperimentalCoroutinesApi::class)

import android.os.Build
import androidx.annotation.RestrictTo
import androidx.annotation.RestrictTo.Scope.LIBRARY
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import kotlinx.coroutines.*
import kotlinx.coroutines.flow.*
import java.util.concurrent.atomic.AtomicReference
import kotlin.experimental.ExperimentalTypeInference
import kotlin.reflect.KProperty0

const val functionName = "FunctionName"
const val unchecked = "UNCHECKED_CAST"

@OptIn(
    ExperimentalCoroutinesApi::class,
    ExperimentalTypeInference::class,
)
@Suppress(functionName)
fun <T> ViewModel.DerivedStateFlow(
    lazyThreadSafetyMode: LazyThreadSafetyMode = LazyThreadSafetyMode.SYNCHRONIZED,
    @BuilderInference deriveValue: DerivedStateFlowBuilderScope<T>.() -> T,
): Lazy<StateFlow<T>> = viewModelScope.DerivedStateFlow(
    lazyThreadSafetyMode = lazyThreadSafetyMode,
    deriveValue = deriveValue,
)

@OptIn(
    ExperimentalCoroutinesApi::class,
    ExperimentalTypeInference::class,
)
@Suppress(functionName)
fun <T> ViewModel.DerivedStateFlow(
    initialValue: T,
    lazyThreadSafetyMode: LazyThreadSafetyMode = LazyThreadSafetyMode.SYNCHRONIZED,
    @BuilderInference deriveValue: suspend DerivedStateFlowBuilderScope<T>.() -> T,
): Lazy<StateFlow<T>> = viewModelScope.DerivedStateFlow(
    initialValue = initialValue,
    lazyThreadSafetyMode = lazyThreadSafetyMode,
    deriveValue = deriveValue,
)

@OptIn(
    ExperimentalCoroutinesApi::class,
    ExperimentalTypeInference::class,
)
@Suppress(functionName)
fun <T> CoroutineScope.DerivedStateFlow(
    lazyThreadSafetyMode: LazyThreadSafetyMode = LazyThreadSafetyMode.SYNCHRONIZED,
    @BuilderInference deriveValue: DerivedStateFlowBuilderScope<T>.() -> T,
): Lazy<StateFlow<T>> = lazy(lazyThreadSafetyMode) {
    DerivedStateFlowBuilderScope.WithoutInitialValue(
        coroutineScope = switchIfRequested(),
        deriveValue = deriveValue,
    ).stateFlow
}

@OptIn(
    ExperimentalCoroutinesApi::class,
    ExperimentalTypeInference::class,
)
@Suppress(functionName)
fun <T> CoroutineScope.DerivedStateFlow(
    initialValue: T,
    lazyThreadSafetyMode: LazyThreadSafetyMode = LazyThreadSafetyMode.SYNCHRONIZED,
    @BuilderInference deriveValue: suspend DerivedStateFlowBuilderScope<T>.() -> T,
): Lazy<StateFlow<T>> = lazy(lazyThreadSafetyMode) {
    DerivedStateFlowBuilderScope.WithInitialValue(
        initialValue = initialValue,
        coroutineScope = switchIfRequested(),
        deriveValue = deriveValue,
    ).stateFlow
}

@OptIn(ExperimentalCoroutinesApi::class)
sealed class DerivedStateFlowBuilderScope<T>(
    @RestrictTo(LIBRARY)
    val coroutineScope: CoroutineScope,
) {

    class WithInitialValue<T>(
        override val initialValue: T,
        coroutineScope: CoroutineScope,
        deriveValue: suspend DerivedStateFlowBuilderScope<T>.() -> T,
        ) : DerivedStateFlowBuilderScope<T>(
        coroutineScope = coroutineScope,
    ) {
        private val _deriveValue = deriveValue

        override suspend fun deriveValue(): T = _deriveValue()
    }

    class WithoutInitialValue<T>(
        coroutineScope: CoroutineScope,
        deriveValue: DerivedStateFlowBuilderScope<T>.() -> T,
    ) : DerivedStateFlowBuilderScope<T>(
        coroutineScope = coroutineScope,
    ) {
        private val _deriveValue = deriveValue

        override val initialValue: T
            get() = _deriveValue()

        override suspend fun deriveValue(): T = _deriveValue()
    }

    companion object {
        val empty = null to null
    }

    abstract val initialValue: T

    private val mutableStateFlow by lazy {
        MutableStateFlow(initialValue)
    }

    val stateFlow: StateFlow<T>
        get() = mutableStateFlow

    private fun updateValue(value: T) {
        mutableStateFlow.value = value
    }

    inline operator fun <reified T> KProperty0<StateFlow<T>>.unaryPlus(): T =
        getStateFlowWithKeyAndCollectIfNotCached(name).value

    inline operator fun <reified T> KProperty0<StateFlow<T>>.not(): T =
        collectStateFlowWithKeyAndCancelCached(name).value

    inline operator fun <reified T> NameAndScopeWithProperty<T>.unaryPlus(): T =
        property.getStateFlowWithKeyAndCollectIfNotCached(key).value

    inline operator fun <reified T> NameAndScopeWithProperty<T>.not(): T =
        property.collectStateFlowWithKeyAndCancelCached(key).value

    inline operator fun <Receiver : Any, reified T> Receiver.invoke(
        provideProperty: Receiver.() -> KProperty0<StateFlow<T>>
    ): NameAndScopeWithProperty<T> {
        val property = provideProperty()
        return NameAndScopeWithProperty(
            key = NameAndScope(
                name = property.name,
                scope = this,
            ),
            property = property
        )
    }

    class NameAndScopeWithProperty<T>(
        val key: NameAndScope,
        val property: KProperty0<StateFlow<T>>,
    )

    data class NameAndScope(
        val name: String,
        val scope: Any,
    )

    @RestrictTo(LIBRARY)
    val stateFlowsByPropertyName = mutableMapOf<Any, Pair<StateFlow<*>, Job>>()

    @RestrictTo(LIBRARY)
    fun existingFor(key: Any): Pair<StateFlow<*>?, Job?> =
        synchronized(stateFlowsByPropertyName) {
            stateFlowsByPropertyName[key] ?: empty
        }


    @RestrictTo(LIBRARY)
    operator fun set(key: Any, stateFlowAndJob: Pair<StateFlow<*>, Job>) {
        synchronized(stateFlowsByPropertyName) {
            stateFlowsByPropertyName[key] = stateFlowAndJob
        }
    }

    inline fun <reified R> KProperty0<StateFlow<R>>.getStateFlowWithKeyAndCollectIfNotCached(key: Any): StateFlow<R> =
        getCachedStateFlowWith(key) ?: getStateFlowAndCacheWith(key)

    inline fun <reified R> KProperty0<StateFlow<R>>.collectStateFlowWithKeyAndCancelCached(key: Any): StateFlow<R> {
        cancelStateCollectionJobWith(key)
        return getStateFlowAndCacheWith(key)
    }

    fun cancelStateCollectionJobWith(key: Any) {
        val (_, existingJob) = existingFor(key)
        existingJob?.cancel()
    }

    inline fun <reified R> getCachedStateFlowWith(key: Any): StateFlow<R>? {
        val (existingStateFlow, existingJob) = existingFor(key)
        if (existingStateFlow != null) {
            val existingValue = existingStateFlow.value
            if (existingValue is R) {
                @Suppress(unchecked)
                return existingStateFlow as StateFlow<R>
            } else existingJob?.cancel()
        }
        return null
    }

    inline fun <reified R> KProperty0<StateFlow<R>>.getStateFlowAndCacheWith(key: Any): StateFlow<R> {
        val builderScope = this@DerivedStateFlowBuilderScope
        val newStateFlow = get()
        builderScope[key] = newStateFlow to newStateFlow.deriveStateOnChange()
        return newStateFlow
    }

    fun <R> StateFlow<R>.deriveStateOnChange(): Job =
        coroutineScope.launch {
            drop(count = 1).collectLatest {
                updateValue(deriveValue())
            }
        }

    abstract suspend fun deriveValue(): T
}

private sealed class CoroutineScopeSwitcher {
    abstract operator fun CoroutineScope.invoke(): CoroutineScope
}

private object DontSwitch : CoroutineScopeSwitcher() {
    override fun CoroutineScope.invoke() = this
}

private object SwitchAllToGlobalCoroutineScope : CoroutineScopeSwitcher() {
    override fun CoroutineScope.invoke() = GlobalScope
}

private val coroutineScopeSwitcher by lazy {
    AtomicReference(
        if (Build.DEVICE.contains("layoutlib", ignoreCase = true)) {
            SwitchAllToGlobalCoroutineScope
        } else DontSwitch
    )
}

fun CoroutineScope.switchIfRequested(): CoroutineScope =
    coroutineScopeSwitcher.get().run { invoke() }