From 43e782ab6f31c3b0cf7b10374a3a72939164a9da Mon Sep 17 00:00:00 2001
From: InsanusMokrassar <ovsyannikov.alexey95@gmail.com>
Date: Mon, 4 Dec 2023 15:37:02 +0600
Subject: [PATCH] SpecialMutableStateFlow : MutableStateFlow

---
 CHANGELOG.md                                  |  1 +
 .../coroutines/compose/FlowState.kt           |  7 +---
 .../coroutines/SpecialMutableStateFlow.kt     | 41 +++++++++++++++----
 3 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6b2f1eb4d46..8c30ac202c3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,7 @@
 ## 0.20.18
 
 * `Coroutines`:
+    * `SpecialMutableStateFlow` now extends `MutableStateFlow`
     * `Compose`:
         * Deprecate `FlowState` due to its complexity in fixes
 
diff --git a/coroutines/compose/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/compose/FlowState.kt b/coroutines/compose/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/compose/FlowState.kt
index 99c8ae1b58f..c3729ac32cf 100644
--- a/coroutines/compose/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/compose/FlowState.kt
+++ b/coroutines/compose/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/compose/FlowState.kt
@@ -1,10 +1,7 @@
 package dev.inmo.micro_utils.coroutines.compose
 
 import androidx.compose.runtime.MutableState
-import androidx.compose.runtime.State
-import androidx.compose.runtime.derivedStateOf
 import dev.inmo.micro_utils.coroutines.SpecialMutableStateFlow
-import dev.inmo.micro_utils.coroutines.doInUI
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 
@@ -26,9 +23,9 @@ class FlowState<T>(
             tryEmit(value)
         }
 
-    override suspend fun onChange(value: T) {
+    override fun onChangeWithoutSync(value: T) {
         internalValue = value
-        super.onChange(value)
+        super.onChangeWithoutSync(value)
     }
 
     override fun component1(): T = value
diff --git a/coroutines/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/SpecialMutableStateFlow.kt b/coroutines/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/SpecialMutableStateFlow.kt
index e39d5cdaa48..c1ca2992ba9 100644
--- a/coroutines/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/SpecialMutableStateFlow.kt
+++ b/coroutines/src/commonMain/kotlin/dev/inmo/micro_utils/coroutines/SpecialMutableStateFlow.kt
@@ -3,10 +3,14 @@ package dev.inmo.micro_utils.coroutines
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.InternalCoroutinesApi
 import kotlinx.coroutines.channels.BufferOverflow
 import kotlinx.coroutines.flow.FlowCollector
 import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.internal.SynchronizedObject
+import kotlinx.coroutines.internal.synchronized
 
 /**
  * Works like [StateFlow], but guarantee that latest value update will always be delivered to
@@ -15,7 +19,9 @@ import kotlinx.coroutines.flow.StateFlow
 open class SpecialMutableStateFlow<T>(
     initialValue: T,
     internalScope: CoroutineScope = CoroutineScope(Dispatchers.Default)
-) : StateFlow<T>, FlowCollector<T>, MutableSharedFlow<T> {
+) : MutableStateFlow<T>, FlowCollector<T>, MutableSharedFlow<T> {
+    @OptIn(InternalCoroutinesApi::class)
+    private val syncObject = SynchronizedObject()
     protected val internalSharedFlow: MutableSharedFlow<T> = MutableSharedFlow(
         replay = 0,
         extraBufferCapacity = 2,
@@ -28,16 +34,37 @@ open class SpecialMutableStateFlow<T>(
     )
 
     protected var _value: T = initialValue
-    override val value: T
+    @OptIn(InternalCoroutinesApi::class)
+    override var value: T
         get() = _value
-    protected open suspend fun onChange(value: T) {
+        set(value) {
+            doOnChangeAction(value)
+        }
+
+    @OptIn(InternalCoroutinesApi::class)
+    override fun compareAndSet(expect: T, update: T): Boolean {
+        return synchronized(syncObject) {
+            if (expect == _value && update != _value) {
+                doOnChangeAction(update)
+            }
+            expect == _value
+        }
+    }
+
+    protected open fun onChangeWithoutSync(value: T) {
         _value = value
-        publicSharedFlow.emit(value)
+        publicSharedFlow.tryEmit(value)
+    }
+    @OptIn(InternalCoroutinesApi::class)
+    protected open fun doOnChangeAction(value: T) {
+        synchronized(syncObject) {
+            if (_value != value) {
+                onChangeWithoutSync(value)
+            }
+        }
     }
     protected val job = internalSharedFlow.subscribe(internalScope) {
-        if (_value != it) {
-            onChange(it)
-        }
+        doOnChangeAction(it)
     }
 
     override val replayCache: List<T>