DataStore wrapper for instant updates.

StateDataStore.kt

package be.digitalia.common.flow

import androidx.datastore.core.DataStore
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Deferred
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.async
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.collectLatest
import kotlinx.coroutines.flow.emitAll
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.update
import kotlinx.coroutines.launch

/**
 * This class updates data synchronously then writes it asynchronously to disk, ignoring errors.
 * Writes will be ignored if the initial value didn't finish loading from disk
 * (data didn't return any value yet).
 * This is ideal for storing UI states where updates need to be visible immediately
 * and always happen after building the UI from the read value.
 */
class StateDataStore<T>(private val dataStore: DataStore<T>, private val scope: CoroutineScope) {
    private val state: Deferred<MutableStateFlow<T>> = scope.async {
        MutableStateFlow(dataStore.data.first()).also {
            scope.launch {
                it.collectLatest { value ->
                    try {
                        dataStore.updateData { value }
                    } catch (e: Exception) {
                        if (e is CancellationException) {
                            throw e
                        }
                    }
                }
            }
        }
    }

    val data: Flow<T> = flow {
        emitAll(state.await())
    }

    @OptIn(ExperimentalCoroutinesApi::class)
    fun updateData(transform: (t: T) -> T): Boolean {
        return try {
            state.getCompleted().update(transform)
            true
        } catch (e: IllegalStateException) {
            false
        }
    }
}