Add MutableStateFlow integration to SavedStateHandle

SavedStateHandle.kt

package dev.kaitei.utils

import androidx.lifecycle.SavedStateHandle
import kotlinx.coroutines.flow.MutableStateFlow

private fun <T> SavedStateHandle.getMutableStateFlow(
    key: String,
): MutableStateFlow<T?> = getMutableStateFlowInternal(
    key = key,
    hasInitialValue = true,
    initialValue = null,
)

@Suppress("UNCHECKED_CAST")
private fun <T> SavedStateHandle.getMutableStateFlow(
    key: String,
    initialValue: T,
): MutableStateFlow<T> = getMutableStateFlowInternal(
    key = key,
    hasInitialValue = true,
    initialValue = initialValue,
) as MutableStateFlow<T>

private fun <T> SavedStateHandle.getMutableStateFlowInternal(
    key: String,
    hasInitialValue: Boolean,
    initialValue: T?,
): MutableStateFlow<T?> = when {
    key in this -> createMutableSavedStateFlow(key, get(key))
    hasInitialValue -> createMutableSavedStateFlow(key, initialValue)
    else -> createMutableSavedStateFlow(key)
}

private fun <T> SavedStateHandle.createMutableSavedStateFlow(
    key: String,
): MutableStateFlow<T?> =
    MutableSavedStateFlowImpl(
        key = key,
        savedStateHandle = this,
        delegate = MutableStateFlow(null)
    )

private fun <T> SavedStateHandle.createMutableSavedStateFlow(
    key: String,
    initialValue: T
): MutableStateFlow<T> =
    MutableSavedStateFlowImpl(
        key = key,
        savedStateHandle = this,
        delegate = MutableStateFlow(initialValue)
    )

private class MutableSavedStateFlowImpl<T>(
    private val savedStateHandle: SavedStateHandle,
    private val key: String,
    private val delegate: MutableStateFlow<T>,
) : MutableStateFlow<T> by delegate {

    override var value: T
        get() = delegate.value
        set(value) {
            delegate.value = value
            savedStateHandle.set(key, value)
        }

    override suspend fun emit(value: T) {
        delegate.emit(value)
        savedStateHandle.set(key, value)
    }

    override fun tryEmit(value: T): Boolean {
        savedStateHandle.set(key, value)
        return delegate.tryEmit(value)
    }

    override fun compareAndSet(expect: T, update: T): Boolean {
        val newValueSet = delegate.compareAndSet(expect, update)
        if (newValueSet) {
            savedStateHandle.set(key, value)
        }
        return newValueSet
    }
}