Potential workaround for #3820 (#4054)

Replace the specific place where ARFU gets misexecuted by a specific Android toolchain

Fixes #3820

Author: qwwdfsad

kotlinx-coroutines-core/common/src/flow/StateFlow.kt

@@ -250,8 +250,14 @@ private class StateFlowSlot : AbstractSharedFlowSlot<StateFlowImpl<*>>() {
      *
      * It is important that default `null` value is used, because there can be a race between allocation
      * of a new slot and trying to do [makePending] on this slot.
+     *
+     * ===
+     * This should be `atomic<Any?>(null)` instead of the atomic reference, but because of #3820
+     * it is used as a **temporary** solution starting from 1.8.1 version.
+     * Depending on the fix rollout on Android, it will be removed in 1.9.0 or 2.0.0.
+     * See https://issuetracker.google.com/issues/325123736
      */
-    private val _state = atomic<Any?>(null)
+    private val _state = WorkaroundAtomicReference<Any?>(null)
 
     override fun allocateLocked(flow: StateFlowImpl<*>): Boolean {
         // No need for atomic check & update here, since allocated happens under StateFlow lock
@@ -290,7 +296,6 @@ private class StateFlowSlot : AbstractSharedFlowSlot<StateFlowImpl<*>>() {
         return state === PENDING
     }
 
-    @Suppress("UNCHECKED_CAST")
     suspend fun awaitPending(): Unit = suspendCancellableCoroutine sc@ { cont ->
         assert { _state.value !is CancellableContinuationImpl<*> } // can be NONE or PENDING
         if (_state.compareAndSet(NONE, cont)) return@sc // installed continuation, waiting for pending
@@ -306,7 +311,6 @@ private class StateFlowImpl<T>(
     private val _state = atomic(initialState) // T | NULL
     private var sequence = 0 // serializes updates, value update is in process when sequence is odd
 
-    @Suppress("UNCHECKED_CAST")
     public override var value: T
         get() = NULL.unbox(_state.value)
         set(value) { updateState(null, value ?: NULL) }

kotlinx-coroutines-core/common/src/internal/Concurrent.common.kt

@@ -19,3 +19,22 @@ internal expect fun <E> identitySet(expectedSize: Int): MutableSet<E>
 @OptionalExpectation
 @Target(AnnotationTarget.FIELD)
 internal expect annotation class BenignDataRace()
+
+// Used **only** as a workaround for #3820 in StateFlow. Do not use anywhere else
+internal expect class WorkaroundAtomicReference<T>(value: T) {
+    public fun get(): T
+    public fun set(value: T)
+    public fun getAndSet(value: T): T
+    public fun compareAndSet(expected: T, value: T): Boolean
+}
+
+@Suppress("UNUSED_PARAMETER", "EXTENSION_SHADOWED_BY_MEMBER")
+internal var <T> WorkaroundAtomicReference<T>.value: T
+    get() = this.get()
+    set(value) = this.set(value)
+
+internal inline fun <T> WorkaroundAtomicReference<T>.loop(action: WorkaroundAtomicReference<T>.(value: T) -> Unit) {
+    while (true) {
+        action(value)
+    }
+}

kotlinx-coroutines-core/jsAndWasmShared/src/internal/Concurrent.kt

@@ -11,3 +11,25 @@ internal class NoOpLock {
 
 internal actual fun <E> identitySet(expectedSize: Int): MutableSet<E> = HashSet(expectedSize)
 
+internal actual class WorkaroundAtomicReference<T> actual constructor(private var value: T) {
+
+    public actual fun get(): T = value
+
+    public actual fun set(value: T) {
+        this.value = value
+    }
+
+    public actual fun getAndSet(value: T): T {
+        val prev = this.value
+        this.value = value
+        return prev
+    }
+
+    public actual fun compareAndSet(expected: T, value: T): Boolean {
+        if (this.value === expected) {
+            this.value = value
+            return true
+        }
+        return false
+    }
+}

kotlinx-coroutines-core/jvm/src/internal/Concurrent.kt

@@ -10,6 +10,9 @@ internal actual typealias ReentrantLock = java.util.concurrent.locks.ReentrantLo
 
 internal actual inline fun <T> ReentrantLock.withLock(action: () -> T) = this.withLockJvm(action)
 
+@Suppress("ACTUAL_WITHOUT_EXPECT") // Visibility
+internal actual typealias WorkaroundAtomicReference<T> = java.util.concurrent.atomic.AtomicReference<T>
+
 @Suppress("NOTHING_TO_INLINE") // So that R8 can completely remove ConcurrentKt class
 internal actual inline fun <E> identitySet(expectedSize: Int): MutableSet<E> =
     Collections.newSetFromMap(IdentityHashMap(expectedSize))

kotlinx-coroutines-core/native/src/internal/Concurrent.kt

@@ -13,3 +13,18 @@ internal actual fun <E> identitySet(expectedSize: Int): MutableSet<E> = HashSet(
 
 @Suppress("ACTUAL_WITHOUT_EXPECT") // This suppress can be removed in 2.0: KT-59355
 internal actual typealias BenignDataRace = kotlin.concurrent.Volatile
+
+internal actual class WorkaroundAtomicReference<T> actual constructor(value: T) {
+
+    private val nativeAtomic = kotlin.concurrent.AtomicReference<T>(value)
+
+    public actual fun get(): T = nativeAtomic.value
+
+    public actual fun set(value: T) {
+        nativeAtomic.value = value
+    }
+
+    public actual fun getAndSet(value: T): T = nativeAtomic.getAndSet(value)
+
+    public actual fun compareAndSet(expected: T, value: T): Boolean = nativeAtomic.compareAndSet(expected, value)
+}