Use Main.immediate dispatcher with fallback to Main dispatcher

lifecycle-coroutines/src/commonMain/kotlin/com/arkivanov/essenty/lifecycle/coroutines/DispatchersExt.kt

@@ -0,0 +1,24 @@
+package com.arkivanov.essenty.lifecycle.coroutines
+
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.MainCoroutineDispatcher
+import kotlin.concurrent.Volatile
+
+@Volatile
+private var isImmediateSupported: Boolean = true
+
+@Suppress("UnusedReceiverParameter")
+internal val MainCoroutineDispatcher.immediateOrFallback: MainCoroutineDispatcher
+    get() {
+        if (isImmediateSupported) {
+            try {
+                return Dispatchers.Main.immediate
+            } catch (ignored: UnsupportedOperationException) {
+            } catch (ignored: NotImplementedError) {
+            }
+
+            isImmediateSupported = false
+        }
+
+        return Dispatchers.Main
+    }

lifecycle-coroutines/src/commonMain/kotlin/com/arkivanov/essenty/lifecycle/coroutines/FlowWithLifecycle.kt

@@ -11,13 +11,17 @@ import kotlin.coroutines.CoroutineContext
  * is at least at [minActiveState] state. The emissions will be stopped when the
  * [lifecycle] state falls below [minActiveState] state.
  *
+ * The [Flow] is collected on the specified [context], which defaults to
+ * [Dispatchers.Main.immediate][kotlinx.coroutines.MainCoroutineDispatcher.immediate]
+ * if available on the current platform, or to [Dispatchers.Main] otherwise.
+ *
  * See the [AndroidX documentation](https://developer.android.com/reference/kotlin/androidx/lifecycle/package-summary#(kotlinx.coroutines.flow.Flow).flowWithLifecycle(androidx.lifecycle.Lifecycle,androidx.lifecycle.Lifecycle.State))
  * for more information.
  */
 fun <T> Flow<T>.withLifecycle(
     lifecycle: Lifecycle,
     minActiveState: Lifecycle.State = Lifecycle.State.STARTED,
-    context: CoroutineContext = Dispatchers.Main
+    context: CoroutineContext = Dispatchers.Main.immediateOrFallback,
 ): Flow<T> = callbackFlow {
     lifecycle.repeatOnLifecycle(minActiveState, context) {
         this@withLifecycle.collect {
@@ -35,7 +39,7 @@ fun <T> Flow<T>.withLifecycle(
 fun <T> Flow<T>.flowWithLifecycle(
     lifecycle: Lifecycle,
     minActiveState: Lifecycle.State = Lifecycle.State.STARTED,
-    context: CoroutineContext = Dispatchers.Main
+    context: CoroutineContext = Dispatchers.Main.immediateOrFallback,
 ): Flow<T> {
     return withLifecycle(lifecycle, minActiveState, context)
 }

lifecycle-coroutines/src/commonMain/kotlin/com/arkivanov/essenty/lifecycle/coroutines/RepeatOnLifecycle.kt

@@ -6,6 +6,7 @@ import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.suspendCancellableCoroutine
 import kotlinx.coroutines.sync.Mutex
@@ -19,7 +20,7 @@ import kotlin.coroutines.resume
  */
 suspend fun LifecycleOwner.repeatOnLifecycle(
     minActiveState: Lifecycle.State = Lifecycle.State.STARTED,
-    context: CoroutineContext = Dispatchers.Main,
+    context: CoroutineContext = Dispatchers.Main.immediateOrFallback,
     block: suspend CoroutineScope.() -> Unit,
 ) {
     lifecycle.repeatOnLifecycle(minActiveState = minActiveState, context = context, block = block)
@@ -31,12 +32,16 @@ suspend fun LifecycleOwner.repeatOnLifecycle(
  *
  * The [block] will cancel and re-launch as the [Lifecycle] moves in and out of the [minActiveState].
  *
+ * The [block] is called on the specified [context], which defaults to
+ * [Dispatchers.Main.immediate][kotlinx.coroutines.MainCoroutineDispatcher.immediate]
+ * if available on the current platform, or to [Dispatchers.Main] otherwise.
+ *
  * See the [AndroidX documentation](https://developer.android.com/reference/kotlin/androidx/lifecycle/package-summary#(androidx.lifecycle.Lifecycle).repeatOnLifecycle(androidx.lifecycle.Lifecycle.State,kotlin.coroutines.SuspendFunction1))
  * for more information.
  */
 suspend fun Lifecycle.repeatOnLifecycle(
     minActiveState: Lifecycle.State = Lifecycle.State.STARTED,
-    context: CoroutineContext = Dispatchers.Main,
+    context: CoroutineContext = Dispatchers.Main.immediateOrFallback,
     block: suspend CoroutineScope.() -> Unit
 ) {
     require(minActiveState != Lifecycle.State.INITIALIZED) {

lifecycle-coroutines/src/commonTest/kotlin/com/arkivanov/essenty/lifecycle/coroutines/DispatchersExtTest.kt

@@ -0,0 +1,38 @@
+package com.arkivanov.essenty.lifecycle.coroutines
+
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestDispatcher
+import kotlinx.coroutines.test.resetMain
+import kotlinx.coroutines.test.setMain
+import kotlin.test.*
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@Suppress("TestFunctionName")
+class DispatchersExtTest {
+
+    @AfterTest
+    fun after() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun WHEN_immediateOrDefault_called_multiple_times_THEN_returns_same_dispatcher() {
+        val dispatcher1 = Dispatchers.Main.immediateOrFallback
+        val dispatcher2 = Dispatchers.Main.immediateOrFallback
+
+        assertSame(dispatcher1, dispatcher2)
+    }
+
+    @Test
+    fun GIVEN_Main_dispatcher_changed_WHEN_immediateOrDefault_called_THEN_returns_updated_dispatcher() {
+        val oldDispatcher = Dispatchers.Main.immediateOrFallback
+        val testDispatcher = StandardTestDispatcher()
+        Dispatchers.setMain(testDispatcher)
+
+        val newDispatcher = Dispatchers.Main.immediateOrFallback
+
+        assertNotSame(oldDispatcher, newDispatcher)
+    }
+}

lifecycle-coroutines/src/commonTest/kotlin/com/arkivanov/essenty/lifecycle/coroutines/LifecycleCoroutinesExtTest.kt

@@ -168,7 +168,7 @@ class LifecycleCoroutinesExtTest {
             flow {
                 repeat(2) { emit(lifecycleState) }
             }
-                .withLifecycle(registry, lifecycleState)
+                .withLifecycle(registry, lifecycleState, testDispatcher)
                 .collect { actual.add(it) }
         }