Consumable.kt

import androidx.compose.runtime.Stable
import kotlinx.atomicfu.AtomicBoolean
import kotlinx.atomicfu.atomic
import kotlinx.coroutines.flow.FlowCollector
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.filter

@Stable
interface Consumable<T> {
    val value: T
    val isConsumed: Boolean
    fun consume()
}

fun consumableUnit(): Consumable<Unit> = ConsumableImpl(rawValue = Unit)
fun consumedUnit(): Consumable<Unit> = ConsumableImpl(rawValue = Unit, initialConsumed = true)

fun <T> consumableOf(value: T): Consumable<T> = ConsumableImpl(rawValue = value)

fun <T> derivedConsumableOf(value: T, origin: Consumable<*>): Consumable<T> = DerivedConsumable(
    rawValue = value,
    origin = origin,
)
fun derivedConsumableUnit(origin: Consumable<*>): Consumable<Unit> = DerivedConsumable(
    rawValue = null,
    origin = origin,
)

fun <T> consumed(): Consumable<T> = ConsumableImpl(rawValue = null, initialConsumed = true)

@Suppress("FunctionName")
fun <T> MutableStateFlowConsumable(): MutableStateFlow<Consumable<T>> = MutableStateFlow(consumed())

suspend inline fun <T> MutableStateFlow<Consumable<T>>.emit(value: T) = emit(consumableOf(value))
suspend inline fun <T> MutableStateFlow<Consumable<T>>.emitAsConsumed() = emit(consumed())

suspend fun <T> StateFlow<Consumable<T>>.collectAsConsumed(collector: FlowCollector<T>) {
    this.filter { !it.isConsumed }
        .collect {
            it.consume()
            collector.emit(it.value)
        }
}

private class ConsumableImpl<T>(
    private val rawValue: T?,
    initialConsumed: Boolean = false,
) : Consumable<T> {

    override val value: T get() = checkNotNull(rawValue) { "Could not resolve value. The rawValue was null." }
    override val isConsumed: Boolean get() = atomicConsumed.value || rawValue == null

    private val atomicConsumed: AtomicBoolean = atomic(initialConsumed)

    override fun consume() {
        atomicConsumed.compareAndSet(false, true)
    }

    override fun equals(other: Any?): Boolean {
        if (this === other) {
            return true
        }

        if (other !is ConsumableImpl<*>) {
            return false
        }

        return when {
            !isConsumed && !other.isConsumed -> {
                rawValue == other.rawValue
            }

            isConsumed && other.isConsumed -> {
                true
            }

            else -> {
                false
            }
        }
    }

    override fun toString(): String = "ConsumableImpl(rawValue=$rawValue, isConsumed=$isConsumed)"
}

@Suppress("EqualsOrHashCode")
private class DerivedConsumable<T>(
    private val rawValue: T?,
    private val origin: Consumable<*>,
) : Consumable<T> {
    override val value: T get() = checkNotNull(rawValue) { "Could not resolve value. The rawValue was null." }

    override val isConsumed: Boolean get() = origin.isConsumed || rawValue == null

    override fun consume() {
        origin.consume()
    }

    override fun equals(other: Any?): Boolean {
        if (this === other) {
            return true
        }

        if (other !is DerivedConsumable<*>) {
            return false
        }

        return when {
            !isConsumed && !other.isConsumed -> {
                rawValue == other.rawValue
            }

            isConsumed && other.isConsumed -> {
                true
            }

            else -> {
                false
            }
        }
    }

    override fun toString(): String = "DerivedConsumable(rawValue=$rawValue, origin=$origin)"
}

ConsumableState.kt

import androidx.compose.runtime.Composable
import androidx.compose.runtime.Stable
import androidx.compose.runtime.derivedStateOf
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableIntStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.setValue

@Stable
interface ConsumableState {
    val shouldConsume: Boolean
    fun consume()
}

@Composable
fun rememberConsumableState(
    consumable: Consumable<*>,
): ConsumableState {

    var recomposeKey by remember { mutableIntStateOf(0) }
    val shouldConsume by remember(consumable, recomposeKey) { derivedStateOf { !consumable.isConsumed } }

    return ConsumableStateImpl(
        shouldConsume,
        onConsume = {
            consumable.consume()
            recomposeKey++
        },
    )
}

private class ConsumableStateImpl(
    override val shouldConsume: Boolean,
    private val onConsume: () -> Unit,
) : ConsumableState {
    override fun consume() {
        onConsume()
    }
}