Simplify DeferredDispatch implementation and improve tests. (#2806)

* Simplify DeferredDispatch implementation and improve tests.

# Conflicts:
#	coil-compose-core/src/commonMain/kotlin/coil3/compose/AsyncImagePainter.kt
#	coil-compose-core/src/commonTest/kotlin/coil3/compose/internal/DeferredDispatchTest.kt

* Clean up comments.

* Attempt fix.

* Attempt fix.

* Use increment methods.

* Formatting.

* Tweak implementation.

* Change test.

* Fix flaky test.

coil-compose-core/api/coil-compose-core.klib.api

@@ -184,8 +184,6 @@ final val coil3.compose.internal/coil3_compose_internal_AbstractContentPainterNo
 final val coil3.compose.internal/coil3_compose_internal_AsyncImageState$stableprop // coil3.compose.internal/coil3_compose_internal_AsyncImageState$stableprop|#static{}coil3_compose_internal_AsyncImageState$stableprop[0]
 final val coil3.compose.internal/coil3_compose_internal_ContentPainterElement$stableprop // coil3.compose.internal/coil3_compose_internal_ContentPainterElement$stableprop|#static{}coil3_compose_internal_ContentPainterElement$stableprop[0]
 final val coil3.compose.internal/coil3_compose_internal_ContentPainterNode$stableprop // coil3.compose.internal/coil3_compose_internal_ContentPainterNode$stableprop|#static{}coil3_compose_internal_ContentPainterNode$stableprop[0]
-final val coil3.compose.internal/coil3_compose_internal_DeferredDispatchCoroutineContext$stableprop // coil3.compose.internal/coil3_compose_internal_DeferredDispatchCoroutineContext$stableprop|#static{}coil3_compose_internal_DeferredDispatchCoroutineContext$stableprop[0]
-final val coil3.compose.internal/coil3_compose_internal_DeferredDispatchCoroutineDispatcher$stableprop // coil3.compose.internal/coil3_compose_internal_DeferredDispatchCoroutineDispatcher$stableprop|#static{}coil3_compose_internal_DeferredDispatchCoroutineDispatcher$stableprop[0]
 final val coil3.compose.internal/coil3_compose_internal_ForwardingCoroutineContext$stableprop // coil3.compose.internal/coil3_compose_internal_ForwardingCoroutineContext$stableprop|#static{}coil3_compose_internal_ForwardingCoroutineContext$stableprop[0]
 final val coil3.compose.internal/coil3_compose_internal_SubcomposeContentPainterElement$stableprop // coil3.compose.internal/coil3_compose_internal_SubcomposeContentPainterElement$stableprop|#static{}coil3_compose_internal_SubcomposeContentPainterElement$stableprop[0]
 final val coil3.compose.internal/coil3_compose_internal_SubcomposeContentPainterNode$stableprop // coil3.compose.internal/coil3_compose_internal_SubcomposeContentPainterNode$stableprop|#static{}coil3_compose_internal_SubcomposeContentPainterNode$stableprop[0]
@@ -211,8 +209,6 @@ final fun coil3.compose.internal/coil3_compose_internal_AbstractContentPainterNo
 final fun coil3.compose.internal/coil3_compose_internal_AsyncImageState$stableprop_getter(): kotlin/Int // coil3.compose.internal/coil3_compose_internal_AsyncImageState$stableprop_getter|coil3_compose_internal_AsyncImageState$stableprop_getter(){}[0]
 final fun coil3.compose.internal/coil3_compose_internal_ContentPainterElement$stableprop_getter(): kotlin/Int // coil3.compose.internal/coil3_compose_internal_ContentPainterElement$stableprop_getter|coil3_compose_internal_ContentPainterElement$stableprop_getter(){}[0]
 final fun coil3.compose.internal/coil3_compose_internal_ContentPainterNode$stableprop_getter(): kotlin/Int // coil3.compose.internal/coil3_compose_internal_ContentPainterNode$stableprop_getter|coil3_compose_internal_ContentPainterNode$stableprop_getter(){}[0]
-final fun coil3.compose.internal/coil3_compose_internal_DeferredDispatchCoroutineContext$stableprop_getter(): kotlin/Int // coil3.compose.internal/coil3_compose_internal_DeferredDispatchCoroutineContext$stableprop_getter|coil3_compose_internal_DeferredDispatchCoroutineContext$stableprop_getter(){}[0]
-final fun coil3.compose.internal/coil3_compose_internal_DeferredDispatchCoroutineDispatcher$stableprop_getter(): kotlin/Int // coil3.compose.internal/coil3_compose_internal_DeferredDispatchCoroutineDispatcher$stableprop_getter|coil3_compose_internal_DeferredDispatchCoroutineDispatcher$stableprop_getter(){}[0]
 final fun coil3.compose.internal/coil3_compose_internal_ForwardingCoroutineContext$stableprop_getter(): kotlin/Int // coil3.compose.internal/coil3_compose_internal_ForwardingCoroutineContext$stableprop_getter|coil3_compose_internal_ForwardingCoroutineContext$stableprop_getter(){}[0]
 final fun coil3.compose.internal/coil3_compose_internal_SubcomposeContentPainterElement$stableprop_getter(): kotlin/Int // coil3.compose.internal/coil3_compose_internal_SubcomposeContentPainterElement$stableprop_getter|coil3_compose_internal_SubcomposeContentPainterElement$stableprop_getter(){}[0]
 final fun coil3.compose.internal/coil3_compose_internal_SubcomposeContentPainterNode$stableprop_getter(): kotlin/Int // coil3.compose.internal/coil3_compose_internal_SubcomposeContentPainterNode$stableprop_getter|coil3_compose_internal_SubcomposeContentPainterNode$stableprop_getter(){}[0]

coil-compose-core/src/androidInstrumentedTest/kotlin/coil3/compose/AsyncImagePainterTest.kt

@@ -824,12 +824,13 @@ class AsyncImagePainterTest {
         val key = Extras.Key(0)
         val value = AtomicInteger()
         val compositionCount = AtomicInteger()
+        val maxCompositionCount = AtomicInteger(3)
 
         composeTestRule.setContent {
             val painter = rememberAsyncImagePainter(
                 model = ImageRequest.Builder(LocalContext.current)
                     .data("https://example.com/image")
-                    .apply { extras[key] = value.getAndIncrement() }
+                    .apply { extras[key] = value.getAndSet(1) }
                     .build(),
                 imageLoader = imageLoader,
             )
@@ -841,17 +842,18 @@ class AsyncImagePainterTest {
 
             val state by painter.state.collectAsState()
 
-            when (compositionCount.incrementAndGet()) {
-                1 -> assertIs<State.Empty>(state)
-                2 -> assertIs<State.Loading>(state)
-                3 -> assertIs<State.Success>(state)
-                else -> error("too many compositions")
+            val compositions = compositionCount.incrementAndGet()
+            if (compositions > maxCompositionCount.get()) {
+                error("too many compositions")
+            }
+            if (state is State.Success) {
+                maxCompositionCount.set(compositions)
             }
         }
 
         waitForRequestComplete()
 
-        assertEquals(3, compositionCount.get())
+        assertEquals(maxCompositionCount.get(), compositionCount.get())
         assertEquals(1, requestTracker.startedRequests)
         assertEquals(1, requestTracker.finishedRequests)
     }

coil-compose-core/src/androidMain/kotlin/coil3/compose/AsyncImagePainter.android.kt

@@ -1,6 +1,7 @@
 package coil3.compose
 
 import android.graphics.drawable.Drawable
+import android.view.View
 import androidx.compose.ui.layout.ContentScale
 import coil3.request.SuccessResult
 import coil3.request.transitionFactory
@@ -38,6 +39,6 @@ internal actual fun maybeNewCrossfadePainter(
 }
 
 private val FakeTransitionTarget = object : TransitionTarget {
-    override val view get() = throw UnsupportedOperationException()
+    override val view: View get() = throw UnsupportedOperationException()
     override val drawable: Drawable? get() = null
 }

coil-compose-core/src/commonMain/kotlin/coil3/compose/AsyncImage.kt

@@ -151,7 +151,6 @@ private fun AsyncImage(
     filterQuality: FilterQuality,
     clipToBounds: Boolean,
 ) {
-    // Create and execute the image request.
     val request = requestOfWithSizeResolver(
         model = state.model,
         contentScale = contentScale,

coil-compose-core/src/commonMain/kotlin/coil3/compose/AsyncImagePainter.kt

@@ -26,8 +26,7 @@ import coil3.compose.AsyncImagePainter.Companion.DefaultTransform
 import coil3.compose.AsyncImagePainter.Input
 import coil3.compose.AsyncImagePainter.State
 import coil3.compose.internal.AsyncImageState
-import coil3.compose.internal.DeferredDispatchCoroutineScope
-import coil3.compose.internal.launchUndispatched
+import coil3.compose.internal.launchWithDeferredDispatch
 import coil3.compose.internal.onStateOf
 import coil3.compose.internal.previewHandler
 import coil3.compose.internal.requestOf
@@ -190,8 +189,8 @@ class AsyncImagePainter internal constructor(
     private val inputFlow: MutableStateFlow<Input> = MutableStateFlow(input)
     val input: StateFlow<Input> = inputFlow.asStateFlow()
 
-    private val _state: MutableStateFlow<State> = MutableStateFlow(State.Empty)
-    val state: StateFlow<State> = _state.asStateFlow()
+    private val stateFlow: MutableStateFlow<State> = MutableStateFlow(State.Empty)
+    val state: StateFlow<State> = stateFlow.asStateFlow()
 
     override val intrinsicSize: Size
         get() = painter?.intrinsicSize ?: Size.Unspecified
@@ -220,7 +219,7 @@ class AsyncImagePainter internal constructor(
     private fun launchJob() {
         val input = _input ?: return
 
-        rememberJob = DeferredDispatchCoroutineScope(scope.coroutineContext).launchUndispatched {
+        rememberJob = scope.launchWithDeferredDispatch {
             val previewHandler = previewHandler
             val state = if (previewHandler != null) {
                 // If we're in inspection mode use the preview renderer.
@@ -297,9 +296,9 @@ class AsyncImagePainter internal constructor(
     }
 
     private fun updateState(state: State) {
-        val previous = _state.value
+        val previous = stateFlow.value
         val current = transform(state)
-        _state.value = current
+        stateFlow.value = current
         painter = maybeNewCrossfadePainter(previous, current, contentScale) ?: current.painter
 
         // Manually forget and remember the old/new painters.

coil-compose-core/src/commonMain/kotlin/coil3/compose/SubcomposeAsyncImage.kt

@@ -169,7 +169,6 @@ private fun SubcomposeAsyncImage(
     clipToBounds: Boolean,
     content: @Composable SubcomposeAsyncImageScope.() -> Unit,
 ) {
-    // Create and execute the image request.
     val request = requestOfWithSizeResolver(
         model = state.model,
         contentScale = contentScale,

coil-compose-core/src/commonMain/kotlin/coil3/compose/internal/DeferredDispatch.kt

@@ -1,5 +1,3 @@
-@file:Suppress("NOTHING_TO_INLINE")
-
 package coil3.compose.internal
 
 import kotlin.coroutines.CoroutineContext
@@ -9,27 +7,38 @@ import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.CoroutineStart
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.InternalCoroutinesApi
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.Runnable
 import kotlinx.coroutines.launch
 
 /**
- * A [CoroutineScope] that does not dispatch until the [CoroutineDispatcher] in its
- * [CoroutineContext] changes.
+ * Launch [block] and defer dispatching until the context's [CoroutineDispatcher] changes.
  */
-internal fun DeferredDispatchCoroutineScope(
-    context: CoroutineContext,
-): CoroutineScope {
-    val originalDispatcher = context.dispatcher ?: Dispatchers.Unconfined
-    return CoroutineScope(DeferredDispatchCoroutineContext(context, originalDispatcher))
+internal fun CoroutineScope.launchWithDeferredDispatch(
+    block: suspend CoroutineScope.() -> Unit,
+): Job {
+    val originalDispatcher = coroutineContext.dispatcher
+    if (originalDispatcher == null || originalDispatcher == Dispatchers.Unconfined) {
+        return launch(
+            context = Dispatchers.Unconfined,
+            start = CoroutineStart.UNDISPATCHED,
+            block = block,
+        )
+    } else {
+        return CoroutineScope(DeferredDispatchCoroutineContext(coroutineContext)).launch(
+            context = DeferredDispatchCoroutineDispatcher(originalDispatcher),
+            start = CoroutineStart.UNDISPATCHED,
+            block = block,
+        )
+    }
 }
 
 /**
  * A special [CoroutineContext] implementation that automatically enables
  * [DeferredDispatchCoroutineDispatcher] dispatching if the context's [CoroutineDispatcher] changes.
  */
-internal class DeferredDispatchCoroutineContext(
+private class DeferredDispatchCoroutineContext(
     context: CoroutineContext,
-    val originalDispatcher: CoroutineDispatcher,
 ) : ForwardingCoroutineContext(context) {
 
     override fun newContext(
@@ -39,24 +48,18 @@ internal class DeferredDispatchCoroutineContext(
         val oldDispatcher = old.dispatcher
         val newDispatcher = new.dispatcher
         if (oldDispatcher is DeferredDispatchCoroutineDispatcher && oldDispatcher != newDispatcher) {
-            oldDispatcher.unconfined = oldDispatcher.unconfined &&
-                (newDispatcher == null || newDispatcher == Dispatchers.Unconfined)
+            oldDispatcher.unconfined = false
         }
 
-        return DeferredDispatchCoroutineContext(new, originalDispatcher)
+        return DeferredDispatchCoroutineContext(new)
     }
 }
 
-/** Launch [block] without dispatching. */
-internal inline fun CoroutineScope.launchUndispatched(
-    noinline block: suspend CoroutineScope.() -> Unit,
-) = launch(DeferredDispatchCoroutineDispatcher(Dispatchers.Unconfined), CoroutineStart.UNDISPATCHED, block)
-
 /**
  * A [CoroutineDispatcher] that delegates to [Dispatchers.Unconfined] while [unconfined] is true
  * and [delegate] when [unconfined] is false.
  */
-internal class DeferredDispatchCoroutineDispatcher(
+private class DeferredDispatchCoroutineDispatcher(
     private val delegate: CoroutineDispatcher,
 ) : CoroutineDispatcher() {
     private val _unconfined = atomic(true)

coil-compose-core/src/commonTest/kotlin/coil3/compose/internal/DeferredDispatchTest.kt

@@ -1,7 +1,6 @@
 package coil3.compose.internal
 
 import coil3.ImageLoader
-import coil3.compose.AsyncImagePainter
 import coil3.decode.DataSource
 import coil3.decode.DecodeResult
 import coil3.decode.Decoder
@@ -26,81 +25,91 @@ import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Runnable
 import kotlinx.coroutines.delay
-import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.withContext
 import okio.Buffer
 
 class DeferredDispatchTest : RobolectricTest() {
     private val testDispatcher = TestCoroutineDispatcher()
-    private val deferredDispatcher = DeferredDispatchCoroutineDispatcher(testDispatcher)
 
     @Test
-    fun doesNotDispatchWhen_unconfined_true() = runTest {
-        deferredDispatcher.unconfined = true
+    fun `does not dispatch if dispatcher does not change`() = runTest {
+        withContext(testDispatcher) {
+            launchWithDeferredDispatch {
+                assertEquals(1, testDispatcher.dispatchCount)
 
-        withContext(deferredDispatcher) {
-            delay(100.milliseconds)
-            assertEquals(0, testDispatcher.dispatchCount)
+                withContext(EmptyCoroutineContext) {
+                    delay(10.milliseconds)
+                }
+
+                assertEquals(1, testDispatcher.dispatchCount)
+            }.join()
         }
     }
 
     @Test
-    fun doesDispatchWhen_unconfined_false() = runTest {
-        deferredDispatcher.unconfined = false
+    fun `does dispatch if dispatcher changes`() = runTest {
+        withContext(testDispatcher) {
+            launchWithDeferredDispatch {
+                assertEquals(1, testDispatcher.dispatchCount)
+
+                withContext(Dispatchers.Default) {
+                    delay(10.milliseconds)
+                }
 
-        withContext(deferredDispatcher) {
-            delay(100.milliseconds)
-            assertEquals(2, testDispatcher.dispatchCount)
+                assertEquals(2, testDispatcher.dispatchCount)
+            }.join()
         }
     }
 
-    /** This test emulates the context that [AsyncImagePainter] launches its request into. */
     @Test
-    fun imageLoaderDoesNotDispatchIfContextDoesNotChange() = runTest {
-        deferredDispatcher.unconfined = true
-
-        DeferredDispatchCoroutineScope(coroutineContext + deferredDispatcher).launch {
-            val imageLoader = ImageLoader(context)
-            val request = ImageRequest.Builder(context)
-                .data(Unit)
-                .fetcherFactory(TestFetcher.Factory())
-                .decoderFactory(TestDecoder.Factory())
-                .coroutineContext(EmptyCoroutineContext)
-                .build()
-            val result = imageLoader.execute(request)
-
-            assertIs<SuccessResult>(result)
-            assertEquals(0, testDispatcher.dispatchCount)
-        }.join()
+    fun `image loader does not dispatch if dispatcher does not change`() = runTest {
+        withContext(testDispatcher) {
+            launchWithDeferredDispatch {
+                assertEquals(1, testDispatcher.dispatchCount)
+
+                val imageLoader = ImageLoader(context)
+                val request = ImageRequest.Builder(context)
+                    .data(Unit)
+                    .fetcherFactory(TestFetcher.Factory())
+                    .decoderFactory(TestDecoder.Factory())
+                    .coroutineContext(EmptyCoroutineContext)
+                    .build()
+                val result = imageLoader.execute(request)
+
+                assertIs<SuccessResult>(result)
+                assertEquals(1, testDispatcher.dispatchCount)
+            }.join()
+        }
     }
 
-    /** This test emulates the context that [AsyncImagePainter] launches its request into. */
     @Test
-    fun imageLoaderDoesDispatchIfContextChanges() = runTest {
-        deferredDispatcher.unconfined = true
-
-        DeferredDispatchCoroutineScope(coroutineContext + deferredDispatcher).launch {
-            val imageLoader = ImageLoader(context)
-            val request = ImageRequest.Builder(context)
-                .data(Unit)
-                .fetcherFactory(TestFetcher.Factory())
-                .decoderFactory(TestDecoder.Factory())
-                .decoderCoroutineContext(Dispatchers.Default)
-                .build()
-            val result = imageLoader.execute(request)
-
-            assertIs<SuccessResult>(result)
-            assertEquals(1, testDispatcher.dispatchCount)
-        }.join()
+    fun `image loader does dispatch if dispatcher changes`() = runTest {
+        withContext(testDispatcher) {
+            launchWithDeferredDispatch {
+                assertEquals(1, testDispatcher.dispatchCount)
+
+                val imageLoader = ImageLoader(context)
+                val request = ImageRequest.Builder(context)
+                    .data(Unit)
+                    .fetcherFactory(TestFetcher.Factory())
+                    .decoderFactory(TestDecoder.Factory())
+                    .decoderCoroutineContext(Dispatchers.Default)
+                    .build()
+                val result = imageLoader.execute(request)
+
+                assertIs<SuccessResult>(result)
+                assertEquals(2, testDispatcher.dispatchCount)
+            }.join()
+        }
     }
 
     private class TestCoroutineDispatcher : CoroutineDispatcher() {
-        private var _dispatchCount by atomic(0)
-        val dispatchCount get() = _dispatchCount
+        private val _dispatchCount = atomic(0)
+        val dispatchCount: Int by _dispatchCount
 
         override fun dispatch(context: CoroutineContext, block: Runnable) {
-            _dispatchCount++
+            _dispatchCount.incrementAndGet()
             block.run()
         }
     }

coil-core/src/commonTest/kotlin/coil3/disk/SimpleTestDispatcher.kt

@@ -24,10 +24,10 @@ class SimpleTestDispatcher : CoroutineDispatcher() {
     fun runNextTask() {
         val task = synchronized(lock) { tasks.removeFirst() }
         try {
-            inProgressTasks.getAndIncrement()
+            inProgressTasks.incrementAndGet()
             task.run()
         } finally {
-            inProgressTasks.getAndDecrement()
+            inProgressTasks.decrementAndGet()
         }
     }