SavedStateHandleExt.kt

/*
 * Created by Attila Kruchió on 2022. 01. 09. 18:24
 * attila.kruchio.dev@gmail.com
 *
 * Copyright (c) 2022.
 * All rights reserved.
 */

package com.attila.kruchio.android.core.ext

import androidx.annotation.GuardedBy
import androidx.lifecycle.SavedStateHandle
import androidx.lifecycle.ViewModel
import kotlinx.coroutines.flow.MutableStateFlow
import kotlin.properties.ReadOnlyProperty
import kotlin.reflect.KProperty

/**
 * Usage: Use it just like a regular [MutableStateFlow].
 *
 * ```
 * private val count by savedStateHandle.getMutableStateFlow(0)
 *
 * fun updateValue() {
 *     count.value = //...
 *
 *     count.update {
 *         // ...
 *     }
 *
 *     val newValue = count.updateAndGet {
 *         // ...
 *     }
 * }
 * ```
 */
@Suppress("unused")
fun <T> SavedStateHandle.getMutableStateFlow(
    initialValue: T, key: String? = null
): MutableSavedStateFlowDelegate<T> {
    return MutableSavedStateFlowDelegateImpl(this, initialValue, key)
}

/**
 * Usage: Use it just like a regular [MutableStateFlow].
 *
 * ```
 * private val count = MutableSavedStateFlow(savedStateHandle, "key_count", 0)
 *
 * fun updateValue() {
 *     count.value = //...
 *
 *     count.update {
 *         // ...
 *     }
 *
 *     val newValue = count.updateAndGet {
 *         // ...
 *     }
 * }
 * ```
 */
class MutableSavedStateFlow<T> private constructor(
    private val savedStateHandle: SavedStateHandle,
    private val key: String,
    private val mutableStateFlow: MutableStateFlow<T>
) : MutableStateFlow<T> by mutableStateFlow {

    constructor(savedStateHandle: SavedStateHandle, key: String, initialValue: T) : this(
        savedStateHandle, key, MutableStateFlow(savedStateHandle[key] ?: initialValue)
    )

    override var value: T
        get() = mutableStateFlow.value
        set(value) {
            setSavedStateValue(value)
            mutableStateFlow.value = value
        }

    override fun compareAndSet(expect: T, update: T): Boolean {
        return mutableStateFlow.compareAndSet(expect, update).also { success ->
            if (success) {
                setSavedStateValue(update)
            }
        }
    }

    private fun setSavedStateValue(value: T) {
        savedStateHandle[key] = value
    }
}

interface MutableSavedStateFlowDelegate<T> : ReadOnlyProperty<ViewModel, MutableSavedStateFlow<T>>

private class MutableSavedStateFlowDelegateImpl<T>(
    private val savedStateHandle: SavedStateHandle,
    private val initialValue: T,
    private val key: String? = null
) : MutableSavedStateFlowDelegate<T> {

    private val lock = Any()

    @GuardedBy("lock")
    @Volatile
    private var INSTANCE: MutableSavedStateFlow<T>? = null

    override fun getValue(
        thisRef: ViewModel, property: KProperty<*>
    ): MutableSavedStateFlow<T> {
        return INSTANCE ?: synchronized(lock) {
            INSTANCE ?: run<MutableSavedStateFlow<T>> {
                val key = key ?: property.name

                MutableSavedStateFlow(savedStateHandle, key, initialValue).also {
                    INSTANCE = it
                }
            }
        }
    }
}