MutableSaveStateFlow - a workaround of using the SavedStateHandle with a StateFlow

MutableSaveStateFlow.kt

class MutableSaveStateFlow<T>(
    private val savedStateHandle: SavedStateHandle,
    private val key: String,
    defaultValue: T
) {
    private val _state: MutableStateFlow<T> =
        MutableStateFlow(savedStateHandle.get<T>(key) ?: defaultValue)
    
    var value: T
        get() = _state.value
        set(value) {
            _state.value = value
            savedStateHandle.set(key, value)
        }

    fun asStateFlow(): StateFlow<T> = _state
}

fun <T> SavedStateHandle.getStateFlow(key: String, initialValue: T): MutableSaveStateFlow<T> =
    MutableSaveStateFlow(this, key, initialValue)

class MyViewModel(
    private val savedStateHandle: SavedStateHandle,
) : ViewModel() {

    private val _state = savedStateHandle.getStateFlow("some-key", "default-value")
    val state: StateFlow<String> get() = _state.asStateFlow()

    init {
        _state.value = "some new value"
    }

}

// reactive query into a repository example 
// tweaked version of original example https://developer.android.com/topic/libraries/architecture/viewmodel-savedstate

class NewsRepository {
    fun fetchQuery(query: String): Flow<List<String>> = TODO()
}

class MyViewModel(
    savedStateHandle: SavedStateHandle,
    private val repository: NewsRepository
) : ViewModel() {

    private val query = savedStateHandle.getStateFlow("query-saved-state-key", "")
    private val _state = query.asStateFlow()
        .flatMapLatest { repository.fetchQuery(it) }
    val state: Flow<List<String>> get() = _state

    fun setQuery(newQuery: String) {
        query.value = newQuery
    }
}

Had the idea for this simple wrapper class when someone asked the question on reddit of how this could be achieved.