commonMain.net.folivo.trixnity.client.store.cache.ConcurrentObservableMap.kt Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of trixnity-client-jvm Show documentation
Multiplatform Kotlin SDK for matrix-protocol
There is a newer version: 4.7.1
package net.folivo.trixnity.client.store.cache

import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.flow.*
import net.folivo.trixnity.utils.concurrentMutableMap
import kotlin.contracts.ExperimentalContracts
import kotlin.contracts.InvocationKind
import kotlin.contracts.contract

internal class ConcurrentObservableMap {
    private val _values = concurrentMutableMap()

    val indexes = MutableStateFlow(listOf>())

    sealed interface CompareAndSetResult {
        data object TryAgain : CompareAndSetResult
        data object OnPut : CompareAndSetResult
        data object OnRemove : CompareAndSetResult
        data object NothingChanged : CompareAndSetResult
    }

    private suspend fun compareAndSet(key: K, expectedOldValue: V?, newValue: V?): CompareAndSetResult =
        _values.write {
            val oldValue = get(key)
            when {
                expectedOldValue != oldValue -> CompareAndSetResult.TryAgain
                newValue == null -> {
                    remove(key)
                    CompareAndSetResult.OnRemove
                }

                expectedOldValue != newValue -> {
                    put(key, newValue)
                    CompareAndSetResult.OnPut
                }

                else -> CompareAndSetResult.NothingChanged
            }
        }

    suspend fun getOrPut(key: K, defaultValue: () -> V): V =
        _values.read { get(key) }
            ?: checkNotNull(update(key) { it ?: defaultValue() })


    @OptIn(ExperimentalContracts::class)
    suspend fun update(
        key: K,
        updater: suspend (V?) -> V?,
    ): V? {
        contract {
            callsInPlace(updater, InvocationKind.AT_LEAST_ONCE)
        }
        return internalUpdate(key, updater = updater)
    }

    suspend fun remove(key: K, stale: Boolean = false) {
        internalUpdate(key, stale) { null }
    }

    @OptIn(ExperimentalContracts::class)
    private suspend fun internalUpdate(
        key: K,
        stale: Boolean = false,
        updater: suspend (V?) -> V?,
    ): V? {
        contract {
            callsInPlace(updater, InvocationKind.AT_LEAST_ONCE)
        }
        // inspired by [MutableStateFlow::update]
        while (true) {
            val oldValue = get(key)
            val newValue = updater(oldValue)
            val compareAndSetResult = compareAndSet(key, oldValue, newValue)
            when (compareAndSetResult) {
                is CompareAndSetResult.TryAgain,
                is CompareAndSetResult.NothingChanged -> {
                }

                is CompareAndSetResult.OnPut -> {
                    indexes.value.forEach { index -> index.onPut(key) }
                }

                is CompareAndSetResult.OnRemove -> {
                    indexes.value.forEach { index -> index.onRemove(key, stale) }
                }
            }
            if (compareAndSetResult !is CompareAndSetResult.TryAgain) {
                return newValue
            }
        }
    }

    suspend fun get(key: K): V? = _values.read { get(key) }

    suspend fun getAll(): Map = _values.read { toMap() }

    suspend fun removeAll() = _values.write {
        clear()
        indexes.value.forEach { index -> index.onRemoveAll() }
    }

    @OptIn(ExperimentalCoroutinesApi::class)
    suspend fun getIndexSubscriptionCount(key: K): Flow =
        indexes.flatMapLatest { indexesValue ->
            if (indexesValue.isEmpty()) flowOf(0)
            else combine(indexesValue.map { it.getSubscriptionCount(key) }) { subscriptionCounts ->
                subscriptionCounts.sum()
            }
        }
}